From 5a2fd1cd8870fa8ac2b53f4a432efbd7250b0f5f Mon Sep 17 00:00:00 2001 From: Niek van der Maas Date: Tue, 4 Aug 2026 14:01:05 +0200 Subject: [PATCH 1/3] fix: support SOCKS5 UDP over MASQUE --- connector/masque/association.go | 292 +++++++++++++++++++++++++++ connector/masque/association_test.go | 182 +++++++++++++++++ connector/masque/connector.go | 79 +++++++- dialer/http3/masque/client.go | 9 + 4 files changed, 554 insertions(+), 8 deletions(-) create mode 100644 connector/masque/association.go create mode 100644 connector/masque/association_test.go diff --git a/connector/masque/association.go b/connector/masque/association.go new file mode 100644 index 00000000..35205c4f --- /dev/null +++ b/connector/masque/association.go @@ -0,0 +1,292 @@ +package masque + +import ( + "context" + "errors" + "net" + "os" + "sync" + "time" +) + +const udpAssociationBufferSize = 64 * 1024 + +var errMissingDestination = errors.New("masque: destination address required") + +type udpAssociationDialFunc func(ctx context.Context, addr net.Addr) (net.PacketConn, error) + +type udpAssociationResult struct { + data []byte + addr net.Addr + err error +} + +type udpAssociationTunnel struct { + conn net.PacketConn + addr net.Addr +} + +type udpAssociationConn struct { + ctx context.Context + cancel context.CancelFunc + localAddr net.Addr + dial udpAssociationDialFunc + closeIdle func() error + + closed chan struct{} + closeOnce sync.Once + closeErr error + results chan udpAssociationResult + + mu sync.Mutex + tunnels map[string]*udpAssociationTunnel + readDeadline time.Time + readDeadlineChanged chan struct{} + writeDeadline time.Time +} + +func newUDPAssociationConn( + ctx context.Context, + localAddr net.Addr, + dial udpAssociationDialFunc, + closeIdle func() error, +) *udpAssociationConn { + ctx, cancel := context.WithCancel(context.WithoutCancel(ctx)) + c := &udpAssociationConn{ + ctx: ctx, + cancel: cancel, + localAddr: localAddr, + dial: dial, + closeIdle: closeIdle, + closed: make(chan struct{}), + results: make(chan udpAssociationResult, 32), + tunnels: make(map[string]*udpAssociationTunnel), + readDeadlineChanged: make(chan struct{}), + } + return c +} + +func (c *udpAssociationConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) { + for { + select { + case <-c.closed: + return 0, nil, net.ErrClosed + default: + } + + c.mu.Lock() + deadline := c.readDeadline + deadlineChanged := c.readDeadlineChanged + c.mu.Unlock() + + var ( + timer *time.Timer + timeout <-chan time.Time + ) + if !deadline.IsZero() { + remaining := time.Until(deadline) + if remaining <= 0 { + return 0, nil, os.ErrDeadlineExceeded + } + timer = time.NewTimer(remaining) + timeout = timer.C + } + + select { + case result := <-c.results: + if timer != nil { + timer.Stop() + } + if result.err != nil { + return 0, nil, result.err + } + return copy(b, result.data), result.addr, nil + case <-deadlineChanged: + if timer != nil { + timer.Stop() + } + continue + case <-timeout: + return 0, nil, os.ErrDeadlineExceeded + case <-c.closed: + if timer != nil { + timer.Stop() + } + return 0, nil, net.ErrClosed + } + } +} + +func (c *udpAssociationConn) WriteTo(b []byte, addr net.Addr) (int, error) { + if addr == nil { + return 0, errMissingDestination + } + + tunnel, err := c.tunnel(addr) + if err != nil { + return 0, err + } + return tunnel.conn.WriteTo(b, addr) +} + +func (c *udpAssociationConn) tunnel(addr net.Addr) (*udpAssociationTunnel, error) { + key := addr.String() + + c.mu.Lock() + defer c.mu.Unlock() + + select { + case <-c.closed: + return nil, net.ErrClosed + default: + } + + if tunnel := c.tunnels[key]; tunnel != nil { + return tunnel, nil + } + + conn, err := c.dial(c.ctx, addr) + if err != nil { + return nil, err + } + if conn == nil { + return nil, errors.New("masque: nil UDP tunnel") + } + if !c.writeDeadline.IsZero() { + if err := conn.SetWriteDeadline(c.writeDeadline); err != nil { + conn.Close() + return nil, err + } + } + + tunnel := &udpAssociationTunnel{conn: conn, addr: addr} + c.tunnels[key] = tunnel + go c.readTunnel(key, tunnel) + return tunnel, nil +} + +func (c *udpAssociationConn) readTunnel(key string, tunnel *udpAssociationTunnel) { + buf := make([]byte, udpAssociationBufferSize) + for { + n, _, err := tunnel.conn.ReadFrom(buf) + if err != nil { + select { + case <-c.closed: + return + default: + } + c.removeTunnel(key, tunnel) + select { + case c.results <- udpAssociationResult{err: err}: + case <-c.closed: + } + return + } + + data := append([]byte(nil), buf[:n]...) + select { + case c.results <- udpAssociationResult{data: data, addr: tunnel.addr}: + case <-c.closed: + return + } + } +} + +func (c *udpAssociationConn) removeTunnel(key string, tunnel *udpAssociationTunnel) { + c.mu.Lock() + if c.tunnels[key] == tunnel { + delete(c.tunnels, key) + } + c.mu.Unlock() + tunnel.conn.Close() +} + +func (c *udpAssociationConn) Read(b []byte) (int, error) { + n, _, err := c.ReadFrom(b) + return n, err +} + +func (c *udpAssociationConn) Write(b []byte) (int, error) { + return 0, errMissingDestination +} + +func (c *udpAssociationConn) Close() error { + c.closeOnce.Do(func() { + close(c.closed) + c.cancel() + + c.mu.Lock() + close(c.readDeadlineChanged) + tunnels := make([]*udpAssociationTunnel, 0, len(c.tunnels)) + for _, tunnel := range c.tunnels { + tunnels = append(tunnels, tunnel) + } + clear(c.tunnels) + c.mu.Unlock() + + var errs []error + for _, tunnel := range tunnels { + if err := tunnel.conn.Close(); err != nil { + errs = append(errs, err) + } + } + if c.closeIdle != nil { + if err := c.closeIdle(); err != nil { + errs = append(errs, err) + } + } + c.closeErr = errors.Join(errs...) + }) + return c.closeErr +} + +func (c *udpAssociationConn) LocalAddr() net.Addr { + return c.localAddr +} + +func (c *udpAssociationConn) RemoteAddr() net.Addr { + return &net.UDPAddr{} +} + +func (c *udpAssociationConn) SetDeadline(t time.Time) error { + if err := c.SetReadDeadline(t); err != nil { + return err + } + return c.SetWriteDeadline(t) +} + +func (c *udpAssociationConn) SetReadDeadline(t time.Time) error { + c.mu.Lock() + defer c.mu.Unlock() + select { + case <-c.closed: + return net.ErrClosed + default: + } + c.readDeadline = t + close(c.readDeadlineChanged) + c.readDeadlineChanged = make(chan struct{}) + return nil +} + +func (c *udpAssociationConn) SetWriteDeadline(t time.Time) error { + c.mu.Lock() + defer c.mu.Unlock() + select { + case <-c.closed: + return net.ErrClosed + default: + } + c.writeDeadline = t + for _, tunnel := range c.tunnels { + if err := tunnel.conn.SetWriteDeadline(t); err != nil { + return err + } + } + return nil +} + +var ( + _ net.Conn = (*udpAssociationConn)(nil) + _ net.PacketConn = (*udpAssociationConn)(nil) +) diff --git a/connector/masque/association_test.go b/connector/masque/association_test.go new file mode 100644 index 00000000..eaccc89c --- /dev/null +++ b/connector/masque/association_test.go @@ -0,0 +1,182 @@ +package masque + +import ( + "context" + "errors" + "net" + "sync" + "testing" + "time" +) + +func TestUDPAssociationConnRoutesByDestination(t *testing.T) { + type dialCall struct { + addr net.Addr + conn *fakePacketConn + } + + var ( + mu sync.Mutex + calls []dialCall + ) + dial := func(ctx context.Context, addr net.Addr) (net.PacketConn, error) { + conn := newFakePacketConn() + mu.Lock() + calls = append(calls, dialCall{addr: addr, conn: conn}) + mu.Unlock() + return conn, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + conn := newUDPAssociationConn(ctx, &net.UDPAddr{}, dial, nil) + cancel() // The route dial context ends before the UDP association is used. + t.Cleanup(func() { conn.Close() }) + + addr1 := &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 53} + addr2 := &net.UDPAddr{IP: net.ParseIP("8.8.8.8"), Port: 53} + if _, err := conn.WriteTo([]byte("first"), addr1); err != nil { + t.Fatalf("first WriteTo failed: %v", err) + } + if _, err := conn.WriteTo([]byte("second"), addr1); err != nil { + t.Fatalf("second WriteTo failed: %v", err) + } + if _, err := conn.WriteTo([]byte("third"), addr2); err != nil { + t.Fatalf("third WriteTo failed: %v", err) + } + + mu.Lock() + gotCalls := append([]dialCall(nil), calls...) + mu.Unlock() + if len(gotCalls) != 2 { + t.Fatalf("expected one tunnel per destination, got %d", len(gotCalls)) + } + if gotCalls[0].addr.String() != addr1.String() || gotCalls[1].addr.String() != addr2.String() { + t.Fatalf("unexpected tunnel destinations: %s, %s", gotCalls[0].addr, gotCalls[1].addr) + } + if got := gotCalls[0].conn.writtenData(); len(got) != 2 || string(got[0]) != "first" || string(got[1]) != "second" { + t.Fatalf("unexpected writes for first destination: %q", got) + } + if got := gotCalls[1].conn.writtenData(); len(got) != 1 || string(got[0]) != "third" { + t.Fatalf("unexpected writes for second destination: %q", got) + } + + gotCalls[1].conn.queueRead([]byte("reply"), &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 1}) + buf := make([]byte, 32) + n, addr, err := conn.ReadFrom(buf) + if err != nil { + t.Fatalf("ReadFrom failed: %v", err) + } + if string(buf[:n]) != "reply" { + t.Fatalf("unexpected reply: %q", buf[:n]) + } + if addr.String() != addr2.String() { + t.Fatalf("expected reply from %s, got %s", addr2, addr) + } +} + +func TestUDPAssociationConnClose(t *testing.T) { + upstream := newFakePacketConn() + idleClosed := 0 + conn := newUDPAssociationConn( + context.Background(), + &net.UDPAddr{}, + func(context.Context, net.Addr) (net.PacketConn, error) { return upstream, nil }, + func() error { + idleClosed++ + return nil + }, + ) + + if _, err := conn.WriteTo([]byte("request"), &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 53}); err != nil { + t.Fatalf("WriteTo failed: %v", err) + } + if err := conn.Close(); err != nil { + t.Fatalf("Close failed: %v", err) + } + if err := conn.Close(); err != nil { + t.Fatalf("second Close failed: %v", err) + } + if idleClosed != 1 { + t.Fatalf("expected idle stream to be closed once, got %d", idleClosed) + } + if !upstream.isClosed() { + t.Fatal("expected destination tunnel to be closed") + } + if _, err := conn.WriteTo([]byte("request"), &net.UDPAddr{}); !errors.Is(err, net.ErrClosed) { + t.Fatalf("expected net.ErrClosed after Close, got %v", err) + } +} + +type fakePacket struct { + data []byte + addr net.Addr +} + +type fakePacketConn struct { + reads chan fakePacket + closed chan struct{} + closeOnce sync.Once + mu sync.Mutex + writes [][]byte +} + +func newFakePacketConn() *fakePacketConn { + return &fakePacketConn{ + reads: make(chan fakePacket, 1), + closed: make(chan struct{}), + } +} + +func (c *fakePacketConn) queueRead(data []byte, addr net.Addr) { + c.reads <- fakePacket{data: append([]byte(nil), data...), addr: addr} +} + +func (c *fakePacketConn) writtenData() [][]byte { + c.mu.Lock() + defer c.mu.Unlock() + result := make([][]byte, len(c.writes)) + for i, data := range c.writes { + result[i] = append([]byte(nil), data...) + } + return result +} + +func (c *fakePacketConn) isClosed() bool { + select { + case <-c.closed: + return true + default: + return false + } +} + +func (c *fakePacketConn) ReadFrom(b []byte) (int, net.Addr, error) { + select { + case packet := <-c.reads: + return copy(b, packet.data), packet.addr, nil + case <-c.closed: + return 0, nil, net.ErrClosed + } +} + +func (c *fakePacketConn) WriteTo(b []byte, _ net.Addr) (int, error) { + select { + case <-c.closed: + return 0, net.ErrClosed + default: + } + c.mu.Lock() + c.writes = append(c.writes, append([]byte(nil), b...)) + c.mu.Unlock() + return len(b), nil +} + +func (c *fakePacketConn) Close() error { + c.closeOnce.Do(func() { close(c.closed) }) + return nil +} + +func (c *fakePacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (c *fakePacketConn) SetDeadline(time.Time) error { return nil } +func (c *fakePacketConn) SetReadDeadline(time.Time) error { return nil } +func (c *fakePacketConn) SetWriteDeadline(time.Time) error { return nil } diff --git a/connector/masque/connector.go b/connector/masque/connector.go index 8040bced..e4f31b2d 100644 --- a/connector/masque/connector.go +++ b/connector/masque/connector.go @@ -10,6 +10,7 @@ import ( "net/url" "strconv" "strings" + "sync" "time" "github.com/go-gost/core/connector" @@ -19,6 +20,7 @@ import ( masque_dialer "github.com/go-gost/x/dialer/http3/masque" masque_util "github.com/go-gost/x/internal/util/masque" "github.com/go-gost/x/registry" + "github.com/quic-go/quic-go/http3" ) func init() { @@ -172,9 +174,69 @@ func (c *masqueConnector) connectUDP(ctx context.Context, conn net.Conn, address return nil, ErrInvalidConnection } - // Get pre-opened stream from dialer (stream opening happens there for dead connection detection) - reqStream := masqueConn.GetRequestStream() - proxyHost := masqueConn.GetHost() + if address == "" { + return c.connectUDPAssociation(ctx, masqueConn, log), nil + } + + return c.connectUDPTarget( + ctx, + conn.LocalAddr(), + masqueConn.GetHost(), + masqueConn.GetRequestStream(), + address, + nil, + log, + ) +} + +func (c *masqueConnector) connectUDPAssociation(ctx context.Context, conn *masque_dialer.MasqueConn, log logger.Logger) net.Conn { + var ( + mu sync.Mutex + firstStream = conn.GetRequestStream() + ) + + nextStream := func(ctx context.Context) (*http3.RequestStream, error) { + mu.Lock() + defer mu.Unlock() + if firstStream != nil { + stream := firstStream + firstStream = nil + return stream, nil + } + return conn.OpenRequestStream(ctx) + } + + closeIdle := func() error { + mu.Lock() + defer mu.Unlock() + if firstStream == nil { + return nil + } + err := firstStream.Close() + firstStream = nil + return err + } + + dial := func(ctx context.Context, addr net.Addr) (net.PacketConn, error) { + stream, err := nextStream(ctx) + if err != nil { + return nil, err + } + return c.connectUDPTarget(ctx, conn.LocalAddr(), conn.GetHost(), stream, addr.String(), addr, log) + } + + return newUDPAssociationConn(ctx, conn.LocalAddr(), dial, closeIdle) +} + +func (c *masqueConnector) connectUDPTarget( + ctx context.Context, + localAddr net.Addr, + proxyHost string, + reqStream *http3.RequestStream, + address string, + remoteAddr net.Addr, + log logger.Logger, +) (*masque_util.DatagramConn, error) { // Apply connect timeout to the actual stream if c.md.connectTimeout > 0 { @@ -259,14 +321,15 @@ func (c *masqueConnector) connectUDP(ctx context.Context, conn net.Conn, address // Get the underlying HTTP/3 stream for datagrams stream := reqStream - // Resolve target address - raddr, err := net.ResolveUDPAddr("udp", address) - if err != nil { - return nil, err + if remoteAddr == nil { + remoteAddr, err = net.ResolveUDPAddr("udp", address) + if err != nil { + return nil, err + } } // Create datagram connection wrapping the request stream - datagramConn := masque_util.NewDatagramConnFromRequestStream(stream, conn.LocalAddr(), raddr) + datagramConn := masque_util.NewDatagramConnFromRequestStream(stream, localAddr, remoteAddr) success = true // Prevent defer from closing stream - datagramConn now owns it return datagramConn, nil diff --git a/dialer/http3/masque/client.go b/dialer/http3/masque/client.go index c482afd6..98f8d0a3 100644 --- a/dialer/http3/masque/client.go +++ b/dialer/http3/masque/client.go @@ -102,6 +102,15 @@ func (c *MasqueConn) GetRequestStream() *http3.RequestStream { return c.reqStream } +// OpenRequestStream opens another request stream on the shared HTTP/3 connection. +// CONNECT-UDP associations use one stream per destination. +func (c *MasqueConn) OpenRequestStream(ctx context.Context) (*http3.RequestStream, error) { + if c.clientConn == nil { + return nil, errors.New("masque: connection is closed") + } + return c.clientConn.OpenRequestStream(ctx) +} + // GetHost returns the proxy host. func (c *MasqueConn) GetHost() string { return c.host From 7e49dfacc91446f6b96b5f24c84ae0a0b660f37f Mon Sep 17 00:00:00 2001 From: Niek van der Maas Date: Tue, 4 Aug 2026 15:29:06 +0200 Subject: [PATCH 2/3] fix: complete MASQUE SOCKS5 proxy support --- connector/masque/association.go | 7 +++++-- connector/masque/association_test.go | 7 ++++--- connector/masque/connector.go | 20 ++++++++++++-------- connector/masque/connector_test.go | 23 +++++++++++++++++++++++ 4 files changed, 44 insertions(+), 13 deletions(-) create mode 100644 connector/masque/connector_test.go diff --git a/connector/masque/association.go b/connector/masque/association.go index 35205c4f..5e25bb3e 100644 --- a/connector/masque/association.go +++ b/connector/masque/association.go @@ -168,7 +168,7 @@ func (c *udpAssociationConn) tunnel(addr net.Addr) (*udpAssociationTunnel, error func (c *udpAssociationConn) readTunnel(key string, tunnel *udpAssociationTunnel) { buf := make([]byte, udpAssociationBufferSize) for { - n, _, err := tunnel.conn.ReadFrom(buf) + n, addr, err := tunnel.conn.ReadFrom(buf) if err != nil { select { case <-c.closed: @@ -183,9 +183,12 @@ func (c *udpAssociationConn) readTunnel(key string, tunnel *udpAssociationTunnel return } + if addr == nil { + addr = tunnel.addr + } data := append([]byte(nil), buf[:n]...) select { - case c.results <- udpAssociationResult{data: data, addr: tunnel.addr}: + case c.results <- udpAssociationResult{data: data, addr: addr}: case <-c.closed: return } diff --git a/connector/masque/association_test.go b/connector/masque/association_test.go index eaccc89c..c6b3dc0f 100644 --- a/connector/masque/association_test.go +++ b/connector/masque/association_test.go @@ -60,7 +60,8 @@ func TestUDPAssociationConnRoutesByDestination(t *testing.T) { t.Fatalf("unexpected writes for second destination: %q", got) } - gotCalls[1].conn.queueRead([]byte("reply"), &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 1}) + replyAddr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 1} + gotCalls[1].conn.queueRead([]byte("reply"), replyAddr) buf := make([]byte, 32) n, addr, err := conn.ReadFrom(buf) if err != nil { @@ -69,8 +70,8 @@ func TestUDPAssociationConnRoutesByDestination(t *testing.T) { if string(buf[:n]) != "reply" { t.Fatalf("unexpected reply: %q", buf[:n]) } - if addr.String() != addr2.String() { - t.Fatalf("expected reply from %s, got %s", addr2, addr) + if addr.String() != replyAddr.String() { + t.Fatalf("expected reply from %s, got %s", replyAddr, addr) } } diff --git a/connector/masque/connector.go b/connector/masque/connector.go index e4f31b2d..b1134ac8 100644 --- a/connector/masque/connector.go +++ b/connector/masque/connector.go @@ -75,6 +75,16 @@ func (c *masqueConnector) Connect(ctx context.Context, conn net.Conn, network, a return nil, fmt.Errorf("%w: %s", ErrUnsupportedNetwork, network) } +func newTCPConnectRequest(address string) *http.Request { + return &http.Request{ + Method: http.MethodConnect, + URL: &url.URL{}, + Host: address, + Header: http.Header{}, + ProtoMajor: 3, + } +} + func (c *masqueConnector) connectTCP(ctx context.Context, conn net.Conn, address string, log logger.Logger) (net.Conn, error) { log.Debugf("connect-tcp %s", address) @@ -105,13 +115,7 @@ func (c *masqueConnector) connectTCP(ctx context.Context, conn net.Conn, address // Create standard HTTP/3 CONNECT request (RFC 9114) // No :protocol pseudo-header for standard CONNECT - req := &http.Request{ - Method: http.MethodConnect, - URL: &url.URL{}, - Host: address, // Target address goes in :authority - Header: http.Header{}, - Proto: "HTTP/3.0", - } + req := newTCPConnectRequest(address) // Add proxy authentication if configured if c.options.Auth != nil { @@ -222,7 +226,7 @@ func (c *masqueConnector) connectUDPAssociation(ctx context.Context, conn *masqu if err != nil { return nil, err } - return c.connectUDPTarget(ctx, conn.LocalAddr(), conn.GetHost(), stream, addr.String(), addr, log) + return c.connectUDPTarget(ctx, conn.LocalAddr(), conn.GetHost(), stream, addr.String(), nil, log) } return newUDPAssociationConn(ctx, conn.LocalAddr(), dial, closeIdle) diff --git a/connector/masque/connector_test.go b/connector/masque/connector_test.go new file mode 100644 index 00000000..edc545f3 --- /dev/null +++ b/connector/masque/connector_test.go @@ -0,0 +1,23 @@ +package masque + +import ( + "net/http" + "testing" +) + +func TestNewTCPConnectRequestUsesStandardConnect(t *testing.T) { + req := newTCPConnectRequest("example.com:443") + + if req.Method != http.MethodConnect { + t.Fatalf("expected CONNECT method, got %s", req.Method) + } + if req.Host != "example.com:443" { + t.Fatalf("expected target authority, got %s", req.Host) + } + if req.Proto != "" { + t.Fatalf("standard CONNECT must not set :protocol, got %q", req.Proto) + } + if req.ProtoMajor != 3 { + t.Fatalf("expected HTTP/3 request, got HTTP/%d", req.ProtoMajor) + } +} From 173de7584f168ef8625d792757b5537be6f2ad69 Mon Sep 17 00:00:00 2001 From: Niek van der Maas Date: Tue, 4 Aug 2026 15:55:58 +0200 Subject: [PATCH 3/3] fix: isolate MASQUE UDP destinations --- connector/masque/association.go | 107 +++++++++++++------- connector/masque/association_test.go | 140 ++++++++++++++++++++++++++- connector/masque/connector.go | 4 +- connector/masque/connector_test.go | 1 + 4 files changed, 212 insertions(+), 40 deletions(-) diff --git a/connector/masque/association.go b/connector/masque/association.go index 5e25bb3e..0beb5e61 100644 --- a/connector/masque/association.go +++ b/connector/masque/association.go @@ -7,6 +7,8 @@ import ( "os" "sync" "time" + + "github.com/go-gost/core/logger" ) const udpAssociationBufferSize = 64 * 1024 @@ -18,12 +20,10 @@ type udpAssociationDialFunc func(ctx context.Context, addr net.Addr) (net.Packet type udpAssociationResult struct { data []byte addr net.Addr - err error } type udpAssociationTunnel struct { conn net.PacketConn - addr net.Addr } type udpAssociationConn struct { @@ -32,6 +32,8 @@ type udpAssociationConn struct { localAddr net.Addr dial udpAssociationDialFunc closeIdle func() error + log logger.Logger + timeout time.Duration closed chan struct{} closeOnce sync.Once @@ -42,14 +44,15 @@ type udpAssociationConn struct { tunnels map[string]*udpAssociationTunnel readDeadline time.Time readDeadlineChanged chan struct{} - writeDeadline time.Time } func newUDPAssociationConn( ctx context.Context, localAddr net.Addr, + timeout time.Duration, dial udpAssociationDialFunc, closeIdle func() error, + log logger.Logger, ) *udpAssociationConn { ctx, cancel := context.WithCancel(context.WithoutCancel(ctx)) c := &udpAssociationConn{ @@ -58,6 +61,8 @@ func newUDPAssociationConn( localAddr: localAddr, dial: dial, closeIdle: closeIdle, + log: log, + timeout: timeout, closed: make(chan struct{}), results: make(chan udpAssociationResult, 32), tunnels: make(map[string]*udpAssociationTunnel), @@ -97,9 +102,6 @@ func (c *udpAssociationConn) ReadFrom(b []byte) (n int, addr net.Addr, err error if timer != nil { timer.Stop() } - if result.err != nil { - return 0, nil, result.err - } return copy(b, result.data), result.addr, nil case <-deadlineChanged: if timer != nil { @@ -124,43 +126,72 @@ func (c *udpAssociationConn) WriteTo(b []byte, addr net.Addr) (int, error) { tunnel, err := c.tunnel(addr) if err != nil { - return 0, err + if c.isClosed() { + return 0, net.ErrClosed + } + c.logError(addr, err) + return len(b), nil + } + n, err := tunnel.conn.WriteTo(b, addr) + if err == nil { + return n, nil } - return tunnel.conn.WriteTo(b, addr) + if c.isClosed() { + return 0, net.ErrClosed + } + if c.removeTunnel(addr.String(), tunnel) { + c.logError(addr, err) + } + return len(b), nil } func (c *udpAssociationConn) tunnel(addr net.Addr) (*udpAssociationTunnel, error) { key := addr.String() c.mu.Lock() - defer c.mu.Unlock() - select { case <-c.closed: + c.mu.Unlock() return nil, net.ErrClosed default: } if tunnel := c.tunnels[key]; tunnel != nil { + c.mu.Unlock() return tunnel, nil } + c.mu.Unlock() - conn, err := c.dial(c.ctx, addr) + ctx := c.ctx + if c.timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, c.timeout) + defer cancel() + } + conn, err := c.dial(ctx, addr) if err != nil { return nil, err } if conn == nil { return nil, errors.New("masque: nil UDP tunnel") } - if !c.writeDeadline.IsZero() { - if err := conn.SetWriteDeadline(c.writeDeadline); err != nil { - conn.Close() - return nil, err - } - } - tunnel := &udpAssociationTunnel{conn: conn, addr: addr} + tunnel := &udpAssociationTunnel{conn: conn} + c.mu.Lock() + select { + case <-c.closed: + c.mu.Unlock() + conn.Close() + return nil, net.ErrClosed + default: + } + if current := c.tunnels[key]; current != nil { + c.mu.Unlock() + conn.Close() + return current, nil + } c.tunnels[key] = tunnel + c.mu.Unlock() go c.readTunnel(key, tunnel) return tunnel, nil } @@ -170,22 +201,15 @@ func (c *udpAssociationConn) readTunnel(key string, tunnel *udpAssociationTunnel for { n, addr, err := tunnel.conn.ReadFrom(buf) if err != nil { - select { - case <-c.closed: + if c.isClosed() { return - default: } - c.removeTunnel(key, tunnel) - select { - case c.results <- udpAssociationResult{err: err}: - case <-c.closed: + if c.removeTunnel(key, tunnel) { + c.logError(key, err) } return } - if addr == nil { - addr = tunnel.addr - } data := append([]byte(nil), buf[:n]...) select { case c.results <- udpAssociationResult{data: data, addr: addr}: @@ -195,13 +219,30 @@ func (c *udpAssociationConn) readTunnel(key string, tunnel *udpAssociationTunnel } } -func (c *udpAssociationConn) removeTunnel(key string, tunnel *udpAssociationTunnel) { +func (c *udpAssociationConn) isClosed() bool { + select { + case <-c.closed: + return true + default: + return false + } +} + +func (c *udpAssociationConn) logError(addr any, err error) { + if c.log != nil { + c.log.Warnf("masque: UDP tunnel %v: %v", addr, err) + } +} + +func (c *udpAssociationConn) removeTunnel(key string, tunnel *udpAssociationTunnel) bool { c.mu.Lock() + removed := c.tunnels[key] == tunnel if c.tunnels[key] == tunnel { delete(c.tunnels, key) } c.mu.Unlock() tunnel.conn.Close() + return removed } func (c *udpAssociationConn) Read(b []byte) (int, error) { @@ -273,19 +314,11 @@ func (c *udpAssociationConn) SetReadDeadline(t time.Time) error { } func (c *udpAssociationConn) SetWriteDeadline(t time.Time) error { - c.mu.Lock() - defer c.mu.Unlock() select { case <-c.closed: return net.ErrClosed default: } - c.writeDeadline = t - for _, tunnel := range c.tunnels { - if err := tunnel.conn.SetWriteDeadline(t); err != nil { - return err - } - } return nil } diff --git a/connector/masque/association_test.go b/connector/masque/association_test.go index c6b3dc0f..ebd983ba 100644 --- a/connector/masque/association_test.go +++ b/connector/masque/association_test.go @@ -28,7 +28,7 @@ func TestUDPAssociationConnRoutesByDestination(t *testing.T) { } ctx, cancel := context.WithCancel(context.Background()) - conn := newUDPAssociationConn(ctx, &net.UDPAddr{}, dial, nil) + conn := newUDPAssociationConn(ctx, &net.UDPAddr{}, 0, dial, nil, nil) cancel() // The route dial context ends before the UDP association is used. t.Cleanup(func() { conn.Close() }) @@ -81,11 +81,13 @@ func TestUDPAssociationConnClose(t *testing.T) { conn := newUDPAssociationConn( context.Background(), &net.UDPAddr{}, + 0, func(context.Context, net.Addr) (net.PacketConn, error) { return upstream, nil }, func() error { idleClosed++ return nil }, + nil, ) if _, err := conn.WriteTo([]byte("request"), &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 53}); err != nil { @@ -108,9 +110,138 @@ func TestUDPAssociationConnClose(t *testing.T) { } } +func TestUDPAssociationConnTunnelFailureIsIsolated(t *testing.T) { + var ( + mu sync.Mutex + calls []*fakePacketConn + ) + conn := newUDPAssociationConn( + context.Background(), + &net.UDPAddr{}, + 0, + func(context.Context, net.Addr) (net.PacketConn, error) { + pc := newFakePacketConn() + mu.Lock() + calls = append(calls, pc) + mu.Unlock() + return pc, nil + }, + nil, + nil, + ) + t.Cleanup(func() { conn.Close() }) + + addr1 := &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 443} + addr2 := &net.UDPAddr{IP: net.ParseIP("8.8.8.8"), Port: 443} + if _, err := conn.WriteTo([]byte("first"), addr1); err != nil { + t.Fatalf("first WriteTo failed: %v", err) + } + if _, err := conn.WriteTo([]byte("second"), addr2); err != nil { + t.Fatalf("second WriteTo failed: %v", err) + } + + mu.Lock() + first, second := calls[0], calls[1] + mu.Unlock() + first.queueReadError(errors.New("stream reset")) + select { + case <-first.closed: + case <-time.After(time.Second): + t.Fatal("failed tunnel was not removed") + } + + second.queueRead([]byte("reply"), addr2) + buf := make([]byte, 32) + n, addr, err := conn.ReadFrom(buf) + if err != nil { + t.Fatalf("healthy tunnel failed after peer reset: %v", err) + } + if string(buf[:n]) != "reply" || addr.String() != addr2.String() { + t.Fatalf("unexpected healthy tunnel reply %q from %v", buf[:n], addr) + } + + if _, err := conn.WriteTo([]byte("retry"), addr1); err != nil { + t.Fatalf("retry WriteTo failed: %v", err) + } + mu.Lock() + callCount := len(calls) + mu.Unlock() + if callCount != 3 { + t.Fatalf("expected failed destination to redial, got %d dials", callCount) + } +} + +func TestUDPAssociationConnDialTimeout(t *testing.T) { + dialErr := make(chan error, 1) + conn := newUDPAssociationConn( + context.Background(), + &net.UDPAddr{}, + 20*time.Millisecond, + func(ctx context.Context, _ net.Addr) (net.PacketConn, error) { + <-ctx.Done() + dialErr <- ctx.Err() + return nil, ctx.Err() + }, + nil, + nil, + ) + t.Cleanup(func() { conn.Close() }) + + payload := []byte("request") + n, err := conn.WriteTo(payload, &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 443}) + if err != nil || n != len(payload) { + t.Fatalf("timed out datagram was not dropped cleanly: n=%d err=%v", n, err) + } + if err := <-dialErr; !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected dial deadline, got %v", err) + } +} + +func TestUDPAssociationConnDialDoesNotHoldAssociationLock(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + conn := newUDPAssociationConn( + context.Background(), + &net.UDPAddr{}, + 0, + func(context.Context, net.Addr) (net.PacketConn, error) { + close(started) + <-release + return newFakePacketConn(), nil + }, + nil, + nil, + ) + t.Cleanup(func() { conn.Close() }) + + writeDone := make(chan struct{}) + go func() { + conn.WriteTo([]byte("request"), &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 443}) + close(writeDone) + }() + <-started + + deadlineDone := make(chan error, 1) + go func() { + deadlineDone <- conn.SetReadDeadline(time.Now().Add(time.Second)) + }() + select { + case err := <-deadlineDone: + if err != nil { + t.Fatalf("SetReadDeadline failed: %v", err) + } + case <-time.After(time.Second): + t.Fatal("dial held the association lock") + } + + close(release) + <-writeDone +} + type fakePacket struct { data []byte addr net.Addr + err error } type fakePacketConn struct { @@ -132,6 +263,10 @@ func (c *fakePacketConn) queueRead(data []byte, addr net.Addr) { c.reads <- fakePacket{data: append([]byte(nil), data...), addr: addr} } +func (c *fakePacketConn) queueReadError(err error) { + c.reads <- fakePacket{err: err} +} + func (c *fakePacketConn) writtenData() [][]byte { c.mu.Lock() defer c.mu.Unlock() @@ -154,6 +289,9 @@ func (c *fakePacketConn) isClosed() bool { func (c *fakePacketConn) ReadFrom(b []byte) (int, net.Addr, error) { select { case packet := <-c.reads: + if packet.err != nil { + return 0, nil, packet.err + } return copy(b, packet.data), packet.addr, nil case <-c.closed: return 0, nil, net.ErrClosed diff --git a/connector/masque/connector.go b/connector/masque/connector.go index b1134ac8..c586d60f 100644 --- a/connector/masque/connector.go +++ b/connector/masque/connector.go @@ -226,10 +226,10 @@ func (c *masqueConnector) connectUDPAssociation(ctx context.Context, conn *masqu if err != nil { return nil, err } - return c.connectUDPTarget(ctx, conn.LocalAddr(), conn.GetHost(), stream, addr.String(), nil, log) + return c.connectUDPTarget(ctx, conn.LocalAddr(), conn.GetHost(), stream, addr.String(), addr, log) } - return newUDPAssociationConn(ctx, conn.LocalAddr(), dial, closeIdle) + return newUDPAssociationConn(ctx, conn.LocalAddr(), c.md.connectTimeout, dial, closeIdle, log) } func (c *masqueConnector) connectUDPTarget( diff --git a/connector/masque/connector_test.go b/connector/masque/connector_test.go index edc545f3..96571fb5 100644 --- a/connector/masque/connector_test.go +++ b/connector/masque/connector_test.go @@ -15,6 +15,7 @@ func TestNewTCPConnectRequestUsesStandardConnect(t *testing.T) { t.Fatalf("expected target authority, got %s", req.Host) } if req.Proto != "" { + // quic-go treats any non-empty Proto as an extended CONNECT :protocol. t.Fatalf("standard CONNECT must not set :protocol, got %q", req.Proto) } if req.ProtoMajor != 3 {