diff --git a/dnstt-client/main.go b/dnstt-client/main.go index f7b042a..6ac971e 100644 --- a/dnstt-client/main.go +++ b/dnstt-client/main.go @@ -120,11 +120,400 @@ 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 smuxSession interface { + OpenStream() (*smux.Stream, error) + Close() error +} + +type sessionState struct { + count int + draining bool + conv uint32 +} + +type sessionManager struct { + pubkey []byte + domain dns.Name + remoteAddr net.Addr + pconn net.PacketConn + mtu int + + mu sync.RWMutex + createMu sync.Mutex + // createSessionFn overrides session creation in tests. + createSessionFn func(closeExisting bool) error + sessions map[smuxSession]*sessionState + conn *kcp.UDPSession + rw io.ReadWriteCloser + sess smuxSession + 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: &noClosePacketConn{PacketConn: 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() + if sm.sessions != nil { + delete(sm.sessions, sm.sess) + } + 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 { + return sm.createSessionWithClose(true, nil) +} + +func (sm *sessionManager) createSessionWithClose(closeExisting bool, expectedSess smuxSession) error { + if sm.createSessionFn != nil { + return sm.createSessionFn(closeExisting) + } + + 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 { + 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() + 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 + sm.conv = conv + sm.mu.Unlock() + + 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() + 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. +func (sm *sessionManager) getSession() (smuxSession, uint32, error) { + sm.mu.RLock() + sess := sm.sess + conv := sm.conv + sm.mu.RUnlock() + + if sess != nil { + return sess, conv, nil + } + + // Serialize session creation attempts. + sm.createMu.Lock() + defer sm.createMu.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 + } + + // Create new session + err := sm.createSession() + if err != nil { + return nil, 0, err + } + + 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. +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, nil, err + } + + // Try to open a stream stream, err := sess.OpenStream() + if err == 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, 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) + + sess, conv, err = sm.recreateSession(sess, conv, true) + if err != nil { + return nil, 0, nil, fmt.Errorf("recreating session: %v", err) + } + + // Try again with the (possibly new) session + stream, err = sess.OpenStream() + if err != nil { + return nil, 0, nil, fmt.Errorf("session %08x opening stream after recreate: %v", conv, err) + } + release := sm.trackStream(sess, conv) + return stream, conv, release, nil + } + + return nil, 0, nil, fmt.Errorf("session %08x opening stream: %v", conv, err) +} + +func handle(local *net.TCPConn, sm *sessionManager) error { + stream, conv, release, err := sm.openStream() if err != nil { - return fmt.Errorf("session %08x opening stream: %v", conv, err) + return fmt.Errorf("opening stream: %v", err) } + defer release() defer func() { log.Printf("end stream %08x:%d", conv, stream.ID()) stream.Close() @@ -178,48 +567,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 +587,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) } diff --git a/dnstt-client/sessionmanager_test.go b/dnstt-client/sessionmanager_test.go new file mode 100644 index 0000000..a82e84e --- /dev/null +++ b/dnstt-client/sessionmanager_test.go @@ -0,0 +1,166 @@ +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(closeExisting bool) error { + if !closeExisting { + t.Fatal("expected closeExisting true for getSession") + } + 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(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")} + 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(closeExisting bool) error { + if !closeExisting { + t.Fatal("expected closeExisting true for closed pipe") + } + 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") + } + } +} 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) } }