diff --git a/handler/tun/client.go b/handler/tun/client.go index 7b8250fa..657d6dad 100644 --- a/handler/tun/client.go +++ b/handler/tun/client.go @@ -45,10 +45,10 @@ func (h *tunHandler) handleClient(ctx context.Context, conn net.Conn, network st defer cc.Close() if network == "udp" { - ctx, cancel := context.WithCancel(ctx) - defer cancel() + iterCtx, iterCancel := context.WithCancel(ctx) + defer iterCancel() - go h.keepalive(ctx, cc, ips) + go h.keepalive(iterCtx, cc, ips) } return h.transportClient(ctx, conn, cc, log) @@ -100,11 +100,21 @@ func (h *tunHandler) keepalive(ctx context.Context, conn net.Conn, ips []net.IP) } func (h *tunHandler) transportClient(ctx context.Context, tun io.ReadWriter, conn net.Conn, log logger.Logger) error { + c, cancel := context.WithCancel(ctx) + defer cancel() + errc := make(chan error, 2) go func() { var b [MaxMessageSize]byte for { + select { + case <-c.Done(): + errc <- c.Err() + return + default: + } + err := func() error { n, err := tun.Read(b[:]) if err != nil { @@ -155,6 +165,13 @@ func (h *tunHandler) transportClient(ctx context.Context, tun io.ReadWriter, con go func() { var b [MaxMessageSize]byte for { + select { + case <-c.Done(): + errc <- c.Err() + return + default: + } + err := func() error { n, err := conn.Read(b[:]) if err != nil { @@ -214,14 +231,5 @@ func (h *tunHandler) transportClient(ctx context.Context, tun io.ReadWriter, con } }() - select { - case err := <-errc: - if err != nil && err == io.EOF { - err = nil - } - return err - - case <-ctx.Done(): - return ctx.Err() - } + return collectFirstError(errc, cancel) } diff --git a/handler/tun/handler.go b/handler/tun/handler.go index 9c3307e1..afb69e73 100644 --- a/handler/tun/handler.go +++ b/handler/tun/handler.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "io" "net" "sync" "time" @@ -111,6 +112,32 @@ func (h *tunHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler. return h.handleServer(ctx, conn, config, log) } +// collectFirstError drains errc (whose capacity must equal the number of +// goroutines writing to it) and returns the first error that is not io.EOF, +// context.Canceled, or context.DeadlineExceeded. It cancels the derived +// context after the first error to signal sibling goroutines to exit. +// A timeout on the second read prevents deadlock when a goroutine is stuck +// in a blocking read that is not context-aware. +func collectFirstError(errc <-chan error, cancel context.CancelFunc) error { + var firstErr error + // Wait for the first goroutine to exit. + err := <-errc + cancel() // signal the sibling goroutine to exit + if err != nil && err != io.EOF && err != context.Canceled && err != context.DeadlineExceeded { + firstErr = err + } + // Wait for the sibling goroutine, but don't block forever — the + // goroutine may be stuck in a blocking read that is not context-aware. + select { + case err := <-errc: + if err != nil && firstErr == nil && err != io.EOF && err != context.Canceled && err != context.DeadlineExceeded { + firstErr = err + } + case <-time.After(5 * time.Second): + } + return firstErr +} + type tunRouteKey [16]byte func ipToTunRouteKey(ip net.IP) (key tunRouteKey) { diff --git a/handler/tun/server.go b/handler/tun/server.go index 9ce0d2be..8970be8e 100644 --- a/handler/tun/server.go +++ b/handler/tun/server.go @@ -3,6 +3,7 @@ package tun import ( "bytes" "context" + "errors" "io" "net" "net/netip" @@ -29,7 +30,7 @@ func (h *tunHandler) handleServer(ctx context.Context, conn net.Conn, config *tu return h.transportServer(ctx, conn, pc, config, log) }() - if err == ErrTun { + if errors.Is(err, ErrTun) { return err } @@ -39,11 +40,21 @@ func (h *tunHandler) handleServer(ctx context.Context, conn net.Conn, config *tu } func (h *tunHandler) transportServer(ctx context.Context, tun io.ReadWriter, conn net.PacketConn, config *tun_util.Config, log logger.Logger) error { - errc := make(chan error, 1) + c, cancel := context.WithCancel(ctx) + defer cancel() + + errc := make(chan error, 2) go func() { var b [MaxMessageSize]byte for { + select { + case <-c.Done(): + errc <- c.Err() + return + default: + } + err := func() error { n, err := tun.Read(b[:]) if err != nil { @@ -110,6 +121,13 @@ func (h *tunHandler) transportServer(ctx context.Context, tun io.ReadWriter, con go func() { var b [MaxMessageSize]byte for { + select { + case <-c.Done(): + errc <- c.Err() + return + default: + } + err := func() error { n, addr, err := conn.ReadFrom(b[:]) if err != nil { @@ -231,11 +249,7 @@ func (h *tunHandler) transportServer(ctx context.Context, tun io.ReadWriter, con } }() - err := <-errc - if err != nil && err == io.EOF { - err = nil - } - return err + return collectFirstError(errc, cancel) } func (h *tunHandler) updateRoute(ip net.IP, addr net.Addr, log logger.Logger) {