Fix DNS exchange failure and recursion deadlock in connector

Co-authored-by: everyx <lunt.luo@gmail.com>

dns/transport/connector.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"net"
 	"sync"
+	"time"
+
+	E "github.com/sagernet/sing/common/exceptions"
 )
 
 type ConnectorCallbacks[T any] struct {
@@ -16,10 +19,11 @@ type Connector[T any] struct {
 	dial      func(ctx context.Context) (T, error)
 	callbacks ConnectorCallbacks[T]
 
-	access        sync.Mutex
-	connection    T
-	hasConnection bool
-	connecting    chan struct{}
+	access           sync.Mutex
+	connection       T
+	hasConnection    bool
+	connectionCancel context.CancelFunc
+	connecting       chan struct{}
 
 	closeCtx context.Context
 	closed   bool
@@ -47,6 +51,10 @@ func NewSingleflightConnector(closeCtx context.Context, dial func(context.Contex
 	})
 }
 
+type contextKeyConnecting struct{}
+
+var errRecursiveConnectorDial = E.New("recursive connector dial")
+
 func (c *Connector[T]) Get(ctx context.Context) (T, error) {
 	var zero T
 	for {
@@ -64,6 +72,14 @@ func (c *Connector[T]) Get(ctx context.Context) (T, error) {
 		}
 
 		c.hasConnection = false
+		if c.connectionCancel != nil {
+			c.connectionCancel()
+			c.connectionCancel = nil
+		}
+		if isRecursiveConnectorDial(ctx, c) {
+			c.access.Unlock()
+			return zero, errRecursiveConnectorDial
+		}
 
 		if c.connecting != nil {
 			connecting := c.connecting
@@ -79,10 +95,16 @@ func (c *Connector[T]) Get(ctx context.Context) (T, error) {
 			}
 		}
 
+		if err := ctx.Err(); err != nil {
+			c.access.Unlock()
+			return zero, err
+		}
+
 		c.connecting = make(chan struct{})
 		c.access.Unlock()
 
-		connection, err := c.dialWithCancellation(ctx)
+		dialContext := context.WithValue(ctx, contextKeyConnecting{}, c)
+		connection, cancel, err := c.dialWithCancellation(dialContext)
 
 		c.access.Lock()
 		close(c.connecting)
@@ -94,13 +116,21 @@ func (c *Connector[T]) Get(ctx context.Context) (T, error) {
 		}
 
 		if c.closed {
+			cancel()
 			c.callbacks.Close(connection)
 			c.access.Unlock()
 			return zero, ErrTransportClosed
 		}
+		if err = ctx.Err(); err != nil {
+			cancel()
+			c.callbacks.Close(connection)
+			c.access.Unlock()
+			return zero, err
+		}
 
 		c.connection = connection
 		c.hasConnection = true
+		c.connectionCancel = cancel
 		result := c.connection
 		c.access.Unlock()
 
@@ -108,19 +138,63 @@ func (c *Connector[T]) Get(ctx context.Context) (T, error) {
 	}
 }
 
-func (c *Connector[T]) dialWithCancellation(ctx context.Context) (T, error) {
-	dialCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
+func isRecursiveConnectorDial[T any](ctx context.Context, connector *Connector[T]) bool {
+	dialConnector, loaded := ctx.Value(contextKeyConnecting{}).(*Connector[T])
+	return loaded && dialConnector == connector
+}
 
-	go func() {
-		select {
-		case <-c.closeCtx.Done():
+func (c *Connector[T]) dialWithCancellation(ctx context.Context) (T, context.CancelFunc, error) {
+	var zero T
+	if err := ctx.Err(); err != nil {
+		return zero, nil, err
+	}
+	connCtx, cancel := context.WithCancel(c.closeCtx)
+
+	var (
+		stateAccess  sync.Mutex
+		dialComplete bool
+	)
+	stopCancel := context.AfterFunc(ctx, func() {
+		stateAccess.Lock()
+		if !dialComplete {
 			cancel()
-		case <-dialCtx.Done():
 		}
-	}()
+		stateAccess.Unlock()
+	})
+	select {
+	case <-ctx.Done():
+		stateAccess.Lock()
+		dialComplete = true
+		stateAccess.Unlock()
+		stopCancel()
+		cancel()
+		return zero, nil, ctx.Err()
+	default:
+	}
 
-	return c.dial(dialCtx)
+	connection, err := c.dial(valueContext{connCtx, ctx})
+	stateAccess.Lock()
+	dialComplete = true
+	stateAccess.Unlock()
+	stopCancel()
+	if err != nil {
+		cancel()
+		return zero, nil, err
+	}
+	return connection, cancel, nil
+}
+
+type valueContext struct {
+	context.Context
+	parent context.Context
+}
+
+func (v valueContext) Value(key any) any {
+	return v.parent.Value(key)
+}
+
+func (v valueContext) Deadline() (time.Time, bool) {
+	return v.parent.Deadline()
 }
 
 func (c *Connector[T]) Close() error {
@@ -132,6 +206,10 @@ func (c *Connector[T]) Close() error {
 	}
 	c.closed = true
 
+	if c.connectionCancel != nil {
+		c.connectionCancel()
+		c.connectionCancel = nil
+	}
 	if c.hasConnection {
 		c.callbacks.Close(c.connection)
 		c.hasConnection = false
@@ -144,6 +222,10 @@ func (c *Connector[T]) Reset() {
 	c.access.Lock()
 	defer c.access.Unlock()
 
+	if c.connectionCancel != nil {
+		c.connectionCancel()
+		c.connectionCancel = nil
+	}
 	if c.hasConnection {
 		c.callbacks.Reset(c.connection)
 		c.hasConnection = false