From ae95dda9cbf806bb0904b00c72e480c97b4af407 Mon Sep 17 00:00:00 2001 From: AliReza Beigy Date: Fri, 16 Jan 2026 10:14:01 +0330 Subject: [PATCH 1/8] fix: Implement automatic session reconnection in dnstt-client to handle closed pipe errors --- dnstt-client/main.go | 258 ++++++++++++++++++++++++++++++++++++------- 1 file changed, 218 insertions(+), 40 deletions(-) diff --git a/dnstt-client/main.go b/dnstt-client/main.go index f7b042a..c57978f 100644 --- a/dnstt-client/main.go +++ b/dnstt-client/main.go @@ -120,10 +120,220 @@ func sampleUTLSDistribution(spec string) (*utls.ClientHelloID, error) { return ids[sampleWeighted(weights)], nil } -func handle(local *net.TCPConn, sess *smux.Session, conv uint32) error { +// sessionManager manages the KCP connection, Noise channel, and smux session, +// and can recreate them if they become closed. +type sessionManager struct { + pubkey []byte + domain dns.Name + remoteAddr net.Addr + pconn net.PacketConn + mtu int + + mu sync.RWMutex + conn *kcp.UDPSession + rw io.ReadWriteCloser + sess *smux.Session + conv uint32 +} + +// newSessionManager creates a new session manager. +func newSessionManager(pubkey []byte, domain dns.Name, remoteAddr net.Addr, pconn net.PacketConn, mtu int) *sessionManager { + return &sessionManager{ + pubkey: pubkey, + domain: domain, + remoteAddr: remoteAddr, + pconn: pconn, + mtu: mtu, + } +} + +// closeSessionLocked closes the current session if it exists. +// Caller must hold sm.mu write lock. +func (sm *sessionManager) closeSessionLocked() { + if sm.sess != nil { + sm.sess.Close() + sm.sess = nil + } + if sm.rw != nil { + sm.rw.Close() + sm.rw = nil + } + if sm.conn != nil { + conv := sm.conv + log.Printf("end session %08x", conv) + sm.conn.Close() + sm.conn = nil + } + sm.conv = 0 +} + +// createSession creates a new KCP connection, Noise channel, and smux session. +// Caller must NOT hold sm.mu lock. +func (sm *sessionManager) createSession() error { + sm.mu.Lock() + // Close existing session if any + sm.closeSessionLocked() + + // Open a KCP conn on the PacketConn. + // We do this outside the lock to avoid holding it during I/O operations. + sm.mu.Unlock() + + conn, err := kcp.NewConn2(sm.remoteAddr, nil, 0, 0, sm.pconn) + if err != nil { + return fmt.Errorf("opening KCP conn: %v", err) + } + conv := conn.GetConv() + log.Printf("begin session %08x", conv) + + // Permit coalescing the payloads of consecutive sends. + conn.SetStreamMode(true) + // Disable the dynamic congestion window (limit only by the maximum of + // local and remote static windows). + conn.SetNoDelay( + 0, // default nodelay + 0, // default interval + 0, // default resend + 1, // nc=1 => congestion window off + ) + conn.SetWindowSize(turbotunnel.QueueSize/2, turbotunnel.QueueSize/2) + if rc := conn.SetMtu(sm.mtu); !rc { + conn.Close() + panic(rc) + } + + // Put a Noise channel on top of the KCP conn. + rw, err := noise.NewClient(conn, sm.pubkey) + if err != nil { + conn.Close() + return err + } + + // Start a smux session on the Noise channel. + smuxConfig := smux.DefaultConfig() + smuxConfig.Version = 2 + smuxConfig.KeepAliveTimeout = idleTimeout + smuxConfig.MaxStreamBuffer = 1 * 1024 * 1024 // default is 65536 + sess, err := smux.Client(rw, smuxConfig) + if err != nil { + rw.Close() + conn.Close() + return fmt.Errorf("opening smux session: %v", err) + } + + // Lock again to update the session + sm.mu.Lock() + sm.conn = conn + sm.rw = rw + sm.sess = sess + sm.conv = conv + sm.mu.Unlock() + + return nil +} + +// closeSession closes the current session if it exists. +func (sm *sessionManager) closeSession() { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.closeSessionLocked() +} + +// getSession returns the current session, creating one if needed. +func (sm *sessionManager) getSession() (*smux.Session, uint32, error) { + sm.mu.RLock() + sess := sm.sess + conv := sm.conv + sm.mu.RUnlock() + + if sess != nil { + return sess, conv, nil + } + + // Need to create a new session. Upgrade to write lock. + sm.mu.Lock() + defer sm.mu.Unlock() + + // Double-check after acquiring write lock + if sm.sess != nil { + return sm.sess, sm.conv, nil + } + + // Create new session + err := sm.createSession() + if err != nil { + return nil, 0, err + } + + return sm.sess, sm.conv, nil +} + +// openStream opens a new stream, recreating the session if necessary. +func (sm *sessionManager) openStream() (*smux.Stream, uint32, error) { + // Try to get existing session + sess, conv, err := sm.getSession() + if err != nil { + return nil, 0, err + } + + // Try to open a stream stream, err := sess.OpenStream() + if err == nil { + return stream, conv, nil + } + + // If opening stream failed, the session might be closed. + // Check if it's a closed pipe error or similar. + errStr := err.Error() + isClosedError := errors.Is(err, io.ErrClosedPipe) || + strings.Contains(errStr, "closed pipe") || + strings.Contains(errStr, "broken pipe") || + strings.Contains(errStr, "use of closed network connection") || + err == io.EOF + + if isClosedError { + log.Printf("session %08x appears closed, recreating: %v", conv, err) + + // Use write lock to serialize session recreation attempts + sm.mu.Lock() + // Double-check: another goroutine might have already recreated the session + if sm.sess != nil && sm.sess != sess { + // Session was recreated by another goroutine + sess = sm.sess + conv = sm.conv + sm.mu.Unlock() + } else { + // We need to recreate + sm.closeSessionLocked() + sm.mu.Unlock() + + // Create a new session (this acquires its own lock) + err = sm.createSession() + if err != nil { + return nil, 0, fmt.Errorf("recreating session: %v", err) + } + + // Get the new session + sm.mu.RLock() + sess = sm.sess + conv = sm.conv + sm.mu.RUnlock() + } + + // Try again with the (possibly new) session + stream, err = sess.OpenStream() + if err != nil { + return nil, 0, fmt.Errorf("session %08x opening stream after recreate: %v", conv, err) + } + return stream, conv, nil + } + + return nil, 0, fmt.Errorf("session %08x opening stream: %v", conv, err) +} + +func handle(local *net.TCPConn, sm *sessionManager) error { + stream, conv, err := sm.openStream() if err != nil { - return fmt.Errorf("session %08x opening stream: %v", conv, err) + return fmt.Errorf("opening stream: %v", err) } defer func() { log.Printf("end stream %08x:%d", conv, stream.ID()) @@ -178,48 +388,16 @@ func run(pubkey []byte, domain dns.Name, localAddr *net.TCPAddr, remoteAddr net. } log.Printf("effective MTU %d", mtu) - // Open a KCP conn on the PacketConn. - conn, err := kcp.NewConn2(remoteAddr, nil, 0, 0, pconn) - if err != nil { - return fmt.Errorf("opening KCP conn: %v", err) - } - defer func() { - log.Printf("end session %08x", conn.GetConv()) - conn.Close() - }() - log.Printf("begin session %08x", conn.GetConv()) - // Permit coalescing the payloads of consecutive sends. - conn.SetStreamMode(true) - // Disable the dynamic congestion window (limit only by the maximum of - // local and remote static windows). - conn.SetNoDelay( - 0, // default nodelay - 0, // default interval - 0, // default resend - 1, // nc=1 => congestion window off - ) - conn.SetWindowSize(turbotunnel.QueueSize/2, turbotunnel.QueueSize/2) - if rc := conn.SetMtu(mtu); !rc { - panic(rc) - } + // Create session manager + sm := newSessionManager(pubkey, domain, remoteAddr, pconn, mtu) + defer sm.closeSession() - // Put a Noise channel on top of the KCP conn. - rw, err := noise.NewClient(conn, pubkey) + // Create initial session + err = sm.createSession() if err != nil { return err } - // Start a smux session on the Noise channel. - smuxConfig := smux.DefaultConfig() - smuxConfig.Version = 2 - smuxConfig.KeepAliveTimeout = idleTimeout - smuxConfig.MaxStreamBuffer = 1 * 1024 * 1024 // default is 65536 - sess, err := smux.Client(rw, smuxConfig) - if err != nil { - return fmt.Errorf("opening smux session: %v", err) - } - defer sess.Close() - for { local, err := ln.Accept() if err != nil { @@ -230,7 +408,7 @@ func run(pubkey []byte, domain dns.Name, localAddr *net.TCPAddr, remoteAddr net. } go func() { defer local.Close() - err := handle(local.(*net.TCPConn), sess, conn.GetConv()) + err := handle(local.(*net.TCPConn), sm) if err != nil { log.Printf("handle: %v", err) } From 668f966c438f672aba921578cbe9d080f6ba07e6 Mon Sep 17 00:00:00 2001 From: Mygod Date: Fri, 16 Jan 2026 16:38:51 -0800 Subject: [PATCH 2/8] Fix getSession deadlock on session creation --- dnstt-client/main.go | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/dnstt-client/main.go b/dnstt-client/main.go index c57978f..8c8206a 100644 --- a/dnstt-client/main.go +++ b/dnstt-client/main.go @@ -251,12 +251,15 @@ func (sm *sessionManager) getSession() (*smux.Session, uint32, error) { // Need to create a new session. Upgrade to write lock. sm.mu.Lock() - defer sm.mu.Unlock() // Double-check after acquiring write lock if sm.sess != nil { - return sm.sess, sm.conv, nil + sess = sm.sess + conv = sm.conv + sm.mu.Unlock() + return sess, conv, nil } + sm.mu.Unlock() // Create new session err := sm.createSession() @@ -264,7 +267,11 @@ func (sm *sessionManager) getSession() (*smux.Session, uint32, error) { return nil, 0, err } - return sm.sess, sm.conv, nil + sm.mu.RLock() + sess = sm.sess + conv = sm.conv + sm.mu.RUnlock() + return sess, conv, nil } // openStream opens a new stream, recreating the session if necessary. From bbb1019b1f1f3f0f3313c80d5c0e63b6e9bb75b1 Mon Sep 17 00:00:00 2001 From: Mygod Date: Fri, 16 Jan 2026 16:42:28 -0800 Subject: [PATCH 3/8] Serialize session recreation --- dnstt-client/main.go | 45 +++++++++++++++++++++++++------------------- 1 file changed, 26 insertions(+), 19 deletions(-) diff --git a/dnstt-client/main.go b/dnstt-client/main.go index 8c8206a..2853874 100644 --- a/dnstt-client/main.go +++ b/dnstt-client/main.go @@ -129,11 +129,12 @@ type sessionManager struct { pconn net.PacketConn mtu int - mu sync.RWMutex - conn *kcp.UDPSession - rw io.ReadWriteCloser - sess *smux.Session - conv uint32 + mu sync.RWMutex + createMu sync.Mutex + conn *kcp.UDPSession + rw io.ReadWriteCloser + sess *smux.Session + conv uint32 } // newSessionManager creates a new session manager. @@ -249,17 +250,18 @@ func (sm *sessionManager) getSession() (*smux.Session, uint32, error) { return sess, conv, nil } - // Need to create a new session. Upgrade to write lock. - sm.mu.Lock() + // Serialize session creation attempts. + sm.createMu.Lock() + defer sm.createMu.Unlock() - // Double-check after acquiring write lock - if sm.sess != nil { - sess = sm.sess - conv = sm.conv - sm.mu.Unlock() + // Double-check after waiting for the creator. + sm.mu.RLock() + sess = sm.sess + conv = sm.conv + sm.mu.RUnlock() + if sess != nil { return sess, conv, nil } - sm.mu.Unlock() // Create new session err := sm.createSession() @@ -300,22 +302,26 @@ func (sm *sessionManager) openStream() (*smux.Stream, uint32, error) { if isClosedError { log.Printf("session %08x appears closed, recreating: %v", conv, err) - // Use write lock to serialize session recreation attempts - sm.mu.Lock() + // Serialize recreation attempts to avoid leaked sessions. + sm.createMu.Lock() // Double-check: another goroutine might have already recreated the session + sm.mu.RLock() if sm.sess != nil && sm.sess != sess { - // Session was recreated by another goroutine sess = sm.sess conv = sm.conv - sm.mu.Unlock() + sm.mu.RUnlock() } else { - // We need to recreate - sm.closeSessionLocked() + sm.mu.RUnlock() + sm.mu.Lock() + if sm.sess == sess { + sm.closeSessionLocked() + } sm.mu.Unlock() // Create a new session (this acquires its own lock) err = sm.createSession() if err != nil { + sm.createMu.Unlock() return nil, 0, fmt.Errorf("recreating session: %v", err) } @@ -325,6 +331,7 @@ func (sm *sessionManager) openStream() (*smux.Stream, uint32, error) { conv = sm.conv sm.mu.RUnlock() } + sm.createMu.Unlock() // Try again with the (possibly new) session stream, err = sess.OpenStream() From 70234ff3f5ca4cf8892f0aa7e06b67e9ade689b7 Mon Sep 17 00:00:00 2001 From: Mygod Date: Fri, 16 Jan 2026 16:54:37 -0800 Subject: [PATCH 4/8] Fix fallback log formatting --- dnstt-server/main.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dnstt-server/main.go b/dnstt-server/main.go index ef669e6..055986a 100644 --- a/dnstt-server/main.go +++ b/dnstt-server/main.go @@ -580,7 +580,7 @@ func (m *FallbackManager) HandlePacket(packet []byte, clientAddr net.Addr) { // Session doesn't exist, create a new one. newConn, err := net.ListenPacket("udp", ":0") if err != nil { - log.Printf("failed to create fallback socket for %s: %v", clientKey, err) + log.Printf("failed to create fallback socket for %v: %v", clientKey, err) return } proxyConn = newConn // Use the new connection @@ -601,7 +601,7 @@ func (m *FallbackManager) HandlePacket(packet []byte, clientAddr net.Addr) { // Forward the client's packet to the fallback address. _, err := proxyConn.WriteTo(packet, m.fallbackAddr) if err != nil { - log.Printf("fallback write to %s for client %s failed: %v", m.fallbackAddr, clientKey, err) + log.Printf("fallback write to %s for client %v failed: %v", m.fallbackAddr, clientKey, err) } } From aa335295ed0194741d9886d891049c22d3df7e54 Mon Sep 17 00:00:00 2001 From: Mygod Date: Fri, 16 Jan 2026 17:06:04 -0800 Subject: [PATCH 5/8] Keep PacketConn open across sessions --- dnstt-client/main.go | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/dnstt-client/main.go b/dnstt-client/main.go index 2853874..aecc906 100644 --- a/dnstt-client/main.go +++ b/dnstt-client/main.go @@ -137,13 +137,22 @@ type sessionManager struct { conv uint32 } +// noClosePacketConn prevents session teardown from closing a shared PacketConn. +type noClosePacketConn struct { + net.PacketConn +} + +func (c *noClosePacketConn) Close() error { + return nil +} + // newSessionManager creates a new session manager. func newSessionManager(pubkey []byte, domain dns.Name, remoteAddr net.Addr, pconn net.PacketConn, mtu int) *sessionManager { return &sessionManager{ pubkey: pubkey, domain: domain, remoteAddr: remoteAddr, - pconn: pconn, + pconn: &noClosePacketConn{PacketConn: pconn}, mtu: mtu, } } From 08b71f14db87585ae1fd64662caf1b70351bc522 Mon Sep 17 00:00:00 2001 From: Mygod Date: Fri, 16 Jan 2026 17:11:02 -0800 Subject: [PATCH 6/8] Handle smux goaway and closed errors --- dnstt-client/main.go | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/dnstt-client/main.go b/dnstt-client/main.go index aecc906..c63c06c 100644 --- a/dnstt-client/main.go +++ b/dnstt-client/main.go @@ -301,12 +301,18 @@ func (sm *sessionManager) openStream() (*smux.Stream, uint32, error) { // If opening stream failed, the session might be closed. // Check if it's a closed pipe error or similar. - errStr := err.Error() - isClosedError := errors.Is(err, io.ErrClosedPipe) || - strings.Contains(errStr, "closed pipe") || - strings.Contains(errStr, "broken pipe") || - strings.Contains(errStr, "use of closed network connection") || - err == io.EOF + isClosedError := errors.Is(err, smux.ErrGoAway) || + errors.Is(err, io.ErrClosedPipe) || + errors.Is(err, net.ErrClosed) || + errors.Is(err, os.ErrClosed) || + errors.Is(err, io.EOF) + if !isClosedError { + // Fallback to string matching for errors that don't wrap the sentinels. + errStr := err.Error() + isClosedError = strings.Contains(errStr, "closed pipe") || + strings.Contains(errStr, "broken pipe") || + strings.Contains(errStr, "use of closed network connection") + } if isClosedError { log.Printf("session %08x appears closed, recreating: %v", conv, err) From 64bab5e5bfbab8acc6d046f1764bad9e4d159fbb Mon Sep 17 00:00:00 2001 From: Mygod Date: Fri, 16 Jan 2026 17:17:46 -0800 Subject: [PATCH 7/8] Add session manager reconnection tests --- dnstt-client/main.go | 21 +++- dnstt-client/sessionmanager_test.go | 157 ++++++++++++++++++++++++++++ 2 files changed, 173 insertions(+), 5 deletions(-) create mode 100644 dnstt-client/sessionmanager_test.go diff --git a/dnstt-client/main.go b/dnstt-client/main.go index c63c06c..f7e3fe4 100644 --- a/dnstt-client/main.go +++ b/dnstt-client/main.go @@ -122,6 +122,11 @@ func sampleUTLSDistribution(spec string) (*utls.ClientHelloID, error) { // sessionManager manages the KCP connection, Noise channel, and smux session, // and can recreate them if they become closed. +type smuxSession interface { + OpenStream() (*smux.Stream, error) + Close() error +} + type sessionManager struct { pubkey []byte domain dns.Name @@ -131,10 +136,12 @@ type sessionManager struct { mu sync.RWMutex createMu sync.Mutex - conn *kcp.UDPSession - rw io.ReadWriteCloser - sess *smux.Session - conv uint32 + // createSessionFn overrides session creation in tests. + createSessionFn func() error + conn *kcp.UDPSession + rw io.ReadWriteCloser + sess smuxSession + conv uint32 } // noClosePacketConn prevents session teardown from closing a shared PacketConn. @@ -180,6 +187,10 @@ func (sm *sessionManager) closeSessionLocked() { // createSession creates a new KCP connection, Noise channel, and smux session. // Caller must NOT hold sm.mu lock. func (sm *sessionManager) createSession() error { + if sm.createSessionFn != nil { + return sm.createSessionFn() + } + sm.mu.Lock() // Close existing session if any sm.closeSessionLocked() @@ -249,7 +260,7 @@ func (sm *sessionManager) closeSession() { } // getSession returns the current session, creating one if needed. -func (sm *sessionManager) getSession() (*smux.Session, uint32, error) { +func (sm *sessionManager) getSession() (smuxSession, uint32, error) { sm.mu.RLock() sess := sm.sess conv := sm.conv diff --git a/dnstt-client/sessionmanager_test.go b/dnstt-client/sessionmanager_test.go new file mode 100644 index 0000000..c6809e4 --- /dev/null +++ b/dnstt-client/sessionmanager_test.go @@ -0,0 +1,157 @@ +package main + +import ( + "errors" + "io" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/xtaci/smux" +) + +type stubSession struct { + openErr error + openCalls int32 +} + +func (s *stubSession) OpenStream() (*smux.Stream, error) { + atomic.AddInt32(&s.openCalls, 1) + return nil, s.openErr +} + +func (s *stubSession) Close() error { + return nil +} + +func TestGetSessionSerializesCreate(t *testing.T) { + sm := &sessionManager{} + + var createCalls int32 + started := make(chan struct{}) + release := make(chan struct{}) + var startOnce sync.Once + + sm.createSessionFn = func() error { + atomic.AddInt32(&createCalls, 1) + startOnce.Do(func() { close(started) }) + <-release + + sm.mu.Lock() + sm.sess = &stubSession{openErr: errors.New("no stream")} + sm.conv = 1 + sm.mu.Unlock() + return nil + } + + const goroutines = 5 + var wg sync.WaitGroup + errs := make(chan error, goroutines) + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, _, err := sm.getSession() + errs <- err + }() + } + + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for session creation") + } + close(release) + + wg.Wait() + close(errs) + + if got := atomic.LoadInt32(&createCalls); got != 1 { + t.Fatalf("expected 1 createSession call, got %d", got) + } + + for err := range errs { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + } +} + +func TestOpenStreamRecreatesOnGoAway(t *testing.T) { + sm := &sessionManager{} + sm.sess = &stubSession{openErr: smux.ErrGoAway} + sm.conv = 1 + + var createCalls int32 + sm.createSessionFn = func() error { + atomic.AddInt32(&createCalls, 1) + sm.mu.Lock() + sm.sess = &stubSession{openErr: errors.New("still closed")} + sm.conv = 2 + sm.mu.Unlock() + return nil + } + + _, _, err := sm.openStream() + if err == nil { + t.Fatal("expected error") + } + if got := atomic.LoadInt32(&createCalls); got != 1 { + t.Fatalf("expected 1 createSession call, got %d", got) + } +} + +func TestOpenStreamRecreateSerializes(t *testing.T) { + sm := &sessionManager{} + sm.sess = &stubSession{openErr: io.ErrClosedPipe} + sm.conv = 1 + + var createCalls int32 + started := make(chan struct{}) + release := make(chan struct{}) + var startOnce sync.Once + + sm.createSessionFn = func() error { + atomic.AddInt32(&createCalls, 1) + startOnce.Do(func() { close(started) }) + <-release + + sm.mu.Lock() + sm.sess = &stubSession{openErr: errors.New("still closed")} + sm.conv = 2 + sm.mu.Unlock() + return nil + } + + const goroutines = 5 + var wg sync.WaitGroup + errs := make(chan error, goroutines) + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, _, err := sm.openStream() + errs <- err + }() + } + + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for session recreation") + } + close(release) + + wg.Wait() + close(errs) + + if got := atomic.LoadInt32(&createCalls); got != 1 { + t.Fatalf("expected 1 createSession call, got %d", got) + } + for err := range errs { + if err == nil { + t.Fatal("expected error") + } + } +} From 52f7a92dd639547871780e3f560ddd3fe19b9a37 Mon Sep 17 00:00:00 2001 From: Mygod Date: Fri, 16 Jan 2026 17:46:52 -0800 Subject: [PATCH 8/8] Drain sessions on smux goaway --- dnstt-client/main.go | 237 ++++++++++++++++++++++------ dnstt-client/sessionmanager_test.go | 19 ++- 2 files changed, 202 insertions(+), 54 deletions(-) diff --git a/dnstt-client/main.go b/dnstt-client/main.go index f7e3fe4..6ac971e 100644 --- a/dnstt-client/main.go +++ b/dnstt-client/main.go @@ -127,6 +127,12 @@ type smuxSession interface { Close() error } +type sessionState struct { + count int + draining bool + conv uint32 +} + type sessionManager struct { pubkey []byte domain dns.Name @@ -137,7 +143,8 @@ type sessionManager struct { mu sync.RWMutex createMu sync.Mutex // createSessionFn overrides session creation in tests. - createSessionFn func() error + createSessionFn func(closeExisting bool) error + sessions map[smuxSession]*sessionState conn *kcp.UDPSession rw io.ReadWriteCloser sess smuxSession @@ -169,6 +176,9 @@ func newSessionManager(pubkey []byte, domain dns.Name, remoteAddr net.Addr, pcon func (sm *sessionManager) closeSessionLocked() { if sm.sess != nil { sm.sess.Close() + if sm.sessions != nil { + delete(sm.sessions, sm.sess) + } sm.sess = nil } if sm.rw != nil { @@ -187,17 +197,22 @@ func (sm *sessionManager) closeSessionLocked() { // createSession creates a new KCP connection, Noise channel, and smux session. // Caller must NOT hold sm.mu lock. func (sm *sessionManager) createSession() error { + return sm.createSessionWithClose(true, nil) +} + +func (sm *sessionManager) createSessionWithClose(closeExisting bool, expectedSess smuxSession) error { if sm.createSessionFn != nil { - return sm.createSessionFn() + return sm.createSessionFn(closeExisting) } - sm.mu.Lock() - // Close existing session if any - sm.closeSessionLocked() - - // Open a KCP conn on the PacketConn. - // We do this outside the lock to avoid holding it during I/O operations. - sm.mu.Unlock() + if closeExisting { + sm.mu.Lock() + // Close existing session if any. + if expectedSess == nil || sm.sess == expectedSess { + sm.closeSessionLocked() + } + sm.mu.Unlock() + } conn, err := kcp.NewConn2(sm.remoteAddr, nil, 0, 0, sm.pconn) if err != nil { @@ -243,6 +258,10 @@ func (sm *sessionManager) createSession() error { // Lock again to update the session sm.mu.Lock() + if sm.sessions == nil { + sm.sessions = make(map[smuxSession]*sessionState) + } + sm.sessions[sess] = &sessionState{conv: conv} sm.conn = conn sm.rw = rw sm.sess = sess @@ -252,11 +271,139 @@ func (sm *sessionManager) createSession() error { return nil } +func (sm *sessionManager) ensureSessionStateLocked(sess smuxSession, conv uint32) *sessionState { + if sm.sessions == nil { + sm.sessions = make(map[smuxSession]*sessionState) + } + state, ok := sm.sessions[sess] + if !ok { + state = &sessionState{conv: conv} + sm.sessions[sess] = state + } else if state.conv == 0 && conv != 0 { + state.conv = conv + } + return state +} + +func (sm *sessionManager) markSessionDraining(sess smuxSession, conv uint32) (smuxSession, uint32) { + if sess == nil { + return nil, 0 + } + var closeSess smuxSession + var closeConv uint32 + sm.mu.Lock() + state := sm.ensureSessionStateLocked(sess, conv) + state.draining = true + if state.count == 0 { + delete(sm.sessions, sess) + closeSess = sess + closeConv = state.conv + } + sm.mu.Unlock() + return closeSess, closeConv +} + +func (sm *sessionManager) trackStream(sess smuxSession, conv uint32) func() { + sm.mu.Lock() + state := sm.ensureSessionStateLocked(sess, conv) + state.count++ + sm.mu.Unlock() + return func() { + sm.releaseStream(sess) + } +} + +func (sm *sessionManager) releaseStream(sess smuxSession) { + var closeSess smuxSession + var closeConv uint32 + sm.mu.Lock() + state, ok := sm.sessions[sess] + if !ok { + sm.mu.Unlock() + return + } + if state.count > 0 { + state.count-- + } + if state.draining && state.count == 0 { + delete(sm.sessions, sess) + closeSess = sess + closeConv = state.conv + } + sm.mu.Unlock() + if closeSess != nil { + log.Printf("end session %08x", closeConv) + closeSess.Close() + } +} + +func (sm *sessionManager) recreateSession(sess smuxSession, conv uint32, closeExisting bool) (smuxSession, uint32, error) { + sm.createMu.Lock() + defer sm.createMu.Unlock() + + var closeSess smuxSession + var closeConv uint32 + if !closeExisting { + closeSess, closeConv = sm.markSessionDraining(sess, conv) + } + + // Double-check: another goroutine might have already recreated the session. + sm.mu.RLock() + if sm.sess != nil && sm.sess != sess { + sess = sm.sess + conv = sm.conv + sm.mu.RUnlock() + if closeSess != nil { + log.Printf("end session %08x", closeConv) + closeSess.Close() + } + return sess, conv, nil + } + sm.mu.RUnlock() + + err := sm.createSessionWithClose(closeExisting, sess) + if err != nil { + if closeSess != nil { + log.Printf("end session %08x", closeConv) + closeSess.Close() + } + return nil, 0, err + } + + sm.mu.RLock() + sess = sm.sess + conv = sm.conv + sm.mu.RUnlock() + if closeSess != nil { + log.Printf("end session %08x", closeConv) + closeSess.Close() + } + return sess, conv, nil +} + // closeSession closes the current session if it exists. func (sm *sessionManager) closeSession() { sm.mu.Lock() - defer sm.mu.Unlock() - sm.closeSessionLocked() + sessions := sm.sessions + fallbackSess := sm.sess + fallbackConv := sm.conv + sm.sessions = nil + sm.conn = nil + sm.rw = nil + sm.sess = nil + sm.conv = 0 + sm.mu.Unlock() + + if sessions == nil && fallbackSess != nil { + log.Printf("end session %08x", fallbackConv) + fallbackSess.Close() + return + } + + for sess, state := range sessions { + log.Printf("end session %08x", state.conv) + sess.Close() + } } // getSession returns the current session, creating one if needed. @@ -297,23 +444,39 @@ func (sm *sessionManager) getSession() (smuxSession, uint32, error) { } // openStream opens a new stream, recreating the session if necessary. -func (sm *sessionManager) openStream() (*smux.Stream, uint32, error) { +func (sm *sessionManager) openStream() (*smux.Stream, uint32, func(), error) { // Try to get existing session sess, conv, err := sm.getSession() if err != nil { - return nil, 0, err + return nil, 0, nil, err } // Try to open a stream stream, err := sess.OpenStream() if err == nil { - return stream, conv, nil + release := sm.trackStream(sess, conv) + return stream, conv, release, nil + } + + if errors.Is(err, smux.ErrGoAway) { + log.Printf("session %08x goaway, starting new session for new streams: %v", conv, err) + + sess, conv, err = sm.recreateSession(sess, conv, false) + if err != nil { + return nil, 0, nil, fmt.Errorf("recreating session after goaway: %v", err) + } + + stream, err = sess.OpenStream() + if err != nil { + return nil, 0, nil, fmt.Errorf("session %08x opening stream after goaway: %v", conv, err) + } + release := sm.trackStream(sess, conv) + return stream, conv, release, nil } // If opening stream failed, the session might be closed. // Check if it's a closed pipe error or similar. - isClosedError := errors.Is(err, smux.ErrGoAway) || - errors.Is(err, io.ErrClosedPipe) || + isClosedError := errors.Is(err, io.ErrClosedPipe) || errors.Is(err, net.ErrClosed) || errors.Is(err, os.ErrClosed) || errors.Is(err, io.EOF) @@ -328,53 +491,29 @@ func (sm *sessionManager) openStream() (*smux.Stream, uint32, error) { if isClosedError { log.Printf("session %08x appears closed, recreating: %v", conv, err) - // Serialize recreation attempts to avoid leaked sessions. - sm.createMu.Lock() - // Double-check: another goroutine might have already recreated the session - sm.mu.RLock() - if sm.sess != nil && sm.sess != sess { - sess = sm.sess - conv = sm.conv - sm.mu.RUnlock() - } else { - sm.mu.RUnlock() - sm.mu.Lock() - if sm.sess == sess { - sm.closeSessionLocked() - } - sm.mu.Unlock() - - // Create a new session (this acquires its own lock) - err = sm.createSession() - if err != nil { - sm.createMu.Unlock() - return nil, 0, fmt.Errorf("recreating session: %v", err) - } - - // Get the new session - sm.mu.RLock() - sess = sm.sess - conv = sm.conv - sm.mu.RUnlock() + sess, conv, err = sm.recreateSession(sess, conv, true) + if err != nil { + return nil, 0, nil, fmt.Errorf("recreating session: %v", err) } - sm.createMu.Unlock() // Try again with the (possibly new) session stream, err = sess.OpenStream() if err != nil { - return nil, 0, fmt.Errorf("session %08x opening stream after recreate: %v", conv, err) + return nil, 0, nil, fmt.Errorf("session %08x opening stream after recreate: %v", conv, err) } - return stream, conv, nil + release := sm.trackStream(sess, conv) + return stream, conv, release, nil } - return nil, 0, fmt.Errorf("session %08x opening stream: %v", conv, err) + return nil, 0, nil, fmt.Errorf("session %08x opening stream: %v", conv, err) } func handle(local *net.TCPConn, sm *sessionManager) error { - stream, conv, err := sm.openStream() + stream, conv, release, err := sm.openStream() if err != nil { return fmt.Errorf("opening stream: %v", err) } + defer release() defer func() { log.Printf("end stream %08x:%d", conv, stream.ID()) stream.Close() diff --git a/dnstt-client/sessionmanager_test.go b/dnstt-client/sessionmanager_test.go index c6809e4..a82e84e 100644 --- a/dnstt-client/sessionmanager_test.go +++ b/dnstt-client/sessionmanager_test.go @@ -33,7 +33,10 @@ func TestGetSessionSerializesCreate(t *testing.T) { release := make(chan struct{}) var startOnce sync.Once - sm.createSessionFn = func() error { + sm.createSessionFn = func(closeExisting bool) error { + if !closeExisting { + t.Fatal("expected closeExisting true for getSession") + } atomic.AddInt32(&createCalls, 1) startOnce.Do(func() { close(started) }) <-release @@ -84,7 +87,10 @@ func TestOpenStreamRecreatesOnGoAway(t *testing.T) { sm.conv = 1 var createCalls int32 - sm.createSessionFn = func() error { + sm.createSessionFn = func(closeExisting bool) error { + if closeExisting { + t.Fatal("expected closeExisting false for goaway") + } atomic.AddInt32(&createCalls, 1) sm.mu.Lock() sm.sess = &stubSession{openErr: errors.New("still closed")} @@ -93,7 +99,7 @@ func TestOpenStreamRecreatesOnGoAway(t *testing.T) { return nil } - _, _, err := sm.openStream() + _, _, _, err := sm.openStream() if err == nil { t.Fatal("expected error") } @@ -112,7 +118,10 @@ func TestOpenStreamRecreateSerializes(t *testing.T) { release := make(chan struct{}) var startOnce sync.Once - sm.createSessionFn = func() error { + sm.createSessionFn = func(closeExisting bool) error { + if !closeExisting { + t.Fatal("expected closeExisting true for closed pipe") + } atomic.AddInt32(&createCalls, 1) startOnce.Do(func() { close(started) }) <-release @@ -131,7 +140,7 @@ func TestOpenStreamRecreateSerializes(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - _, _, err := sm.openStream() + _, _, _, err := sm.openStream() errs <- err }() }