fix(tun): close TUN gracefully on client shutdown

Author: ruib-7f79mj44pz

app/cmd/client.go

@@ -879,6 +879,10 @@ func runClient(v *viper.Viper) {
 
 	// Register modes
 	var runner clientModeRunner
+	var tunDone <-chan struct{}
+
+	tunCtx, cancelTun := context.WithCancel(context.Background())
+	defer cancelTun()
 	if config.SOCKS5 != nil {
 		runner.Add("SOCKS5 server", func() error {
 			return clientSOCKS5(*config.SOCKS5, c)
@@ -915,9 +919,14 @@ func runClient(v *viper.Viper) {
 		})
 	}
 	if config.TUN != nil {
+		done := make(chan struct{})
+		tunDone = done
+
 		runner.Add("TUN", func() error {
-			return clientTUN(*config.TUN, c)
+			defer close(done)
+			return clientTUN(tunCtx, *config.TUN, c)
 		})
+
 	}
 
 	signalChan := make(chan os.Signal, 1)
@@ -932,11 +941,20 @@ func runClient(v *viper.Viper) {
 	select {
 	case <-signalChan:
 		logger.Info("received signal, shutting down gracefully")
+		cancelTun()
+		if tunDone != nil {
+			<-tunDone
+		}
 	case r := <-runnerChan:
 		if r.OK {
 			logger.Info(r.Msg)
 		} else {
-			_ = c.Close() // Close the client here as Fatal will exit the program without running defer
+			cancelTun()
+
+			if tunDone != nil {
+				<-tunDone
+			}
+			_ = c.Close()
 			if r.Err != nil {
 				logger.Fatal(r.Msg, zap.Error(r.Err))
 			} else {
@@ -1149,7 +1167,7 @@ func clientTCPRedirect(config tcpRedirectConfig, c client.Client) error {
 	return p.ListenAndServe(laddr)
 }
 
-func clientTUN(config tunConfig, c client.Client) error {
+func clientTUN(ctx context.Context, config tunConfig, c client.Client) error {
 	supportedPlatforms := []string{"linux", "darwin", "windows", "android"}
 	if !slices.Contains(supportedPlatforms, runtime.GOOS) {
 		logger.Error("TUN is not supported on this platform", zap.String("platform", runtime.GOOS))
@@ -1232,7 +1250,7 @@ func clientTUN(config tunConfig, c client.Client) error {
 		}
 	}
 	logger.Info("TUN listening", zap.String("interface", config.Name))
-	return server.Serve()
+	return server.Serve(ctx)
 }
 
 // parseServerAddrString parses server address string.

app/internal/tun/server.go

@@ -2,10 +2,12 @@ package tun
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net"
 	"net/netip"
+	"sync"
 
 	tun "github.com/apernet/sing-tun"
 	"github.com/sagernet/sing/common/buf"
@@ -48,7 +50,7 @@ type EventLogger interface {
 	UDPError(addr string, err error)
 }
 
-func (s *Server) Serve() error {
+func (s *Server) Serve(ctx context.Context) error {
 	if !isIPv6Supported() {
 		s.Logger.Warn("tun-pre-check", zap.String("msg", "IPv6 is not supported or enabled on this system, TUN device is created without IPv6 support."))
 		s.Inet6Address = nil
@@ -74,14 +76,23 @@ func (s *Server) Serve() error {
 	if err != nil {
 		return fmt.Errorf("failed to create tun interface: %w", err)
 	}
-	defer tunIf.Close()
+	var closeTunOnce sync.Once
 
+	closeTun := func() {
+		closeTunOnce.Do(func() {
+			_ = tunIf.Close()
+		})
+	}
+	defer closeTun()
 	tunStack, err := tun.NewSystem(tun.StackOptions{
-		Context:    context.Background(),
+		Context:    ctx,
 		Tun:        tunIf,
 		TunOptions: tunOpts,
 		UDPTimeout: s.Timeout,
-		Handler:    &tunHandler{s},
+		Handler: &tunHandler{
+			Server:      s,
+			shutdownCtx: ctx,
+		},
 		Logger: &singLogger{
 			tag:       "tun-stack",
 			zapLogger: s.Logger,
@@ -93,11 +104,25 @@ func (s *Server) Serve() error {
 		return fmt.Errorf("failed to create tun stack: %w", err)
 	}
 	defer tunStack.Close()
-	return tunStack.(tun.StackRunner).Run()
+
+	stopClose := context.AfterFunc(ctx, closeTun)
+	defer stopClose()
+	err = tunStack.(tun.StackRunner).Run()
+	if ctx.Err() != nil {
+		return nil
+	}
+
+	return err
 }
 
 type tunHandler struct {
 	*Server
+	shutdownCtx context.Context
+}
+
+func (t *tunHandler) isShutdownCancellation(err error) bool {
+	return errors.Is(err, context.Canceled) &&
+		errors.Is(t.shutdownCtx.Err(), context.Canceled)
 }
 
 var _ tun.Handler = (*tunHandler)(nil)
@@ -110,10 +135,11 @@ func (t *tunHandler) NewConnection(ctx context.Context, conn net.Conn, m metadat
 	}
 	var closeErr error
 	defer func() {
-		if t.EventLogger != nil {
+		if t.EventLogger != nil && !t.isShutdownCancellation(closeErr) {
 			t.EventLogger.TCPError(addr, reqAddr, closeErr)
 		}
 	}()
+
 	rc, err := t.HyClient.TCP(reqAddr)
 	if err != nil {
 		closeErr = err
@@ -147,10 +173,11 @@ func (t *tunHandler) NewPacketConnection(ctx context.Context, conn network.Packe
 	}
 	var closeErr error
 	defer func() {
-		if t.EventLogger != nil {
+		if t.EventLogger != nil && !t.isShutdownCancellation(closeErr) {
 			t.EventLogger.UDPError(addr, closeErr)
 		}
 	}()
+
 	rc, err := t.HyClient.UDP()
 	if err != nil {
 		closeErr = err