From 6fd25387f7c6d8060d5ef4f8e9b093a7d66cf2b3 Mon Sep 17 00:00:00 2001 From: Bagas Date: Fri, 17 Jul 2026 22:22:25 +0700 Subject: [PATCH] fix(session): resolve race conditions (#151) ## Summary Comprehensive fixes for race conditions, concurrency bugs, and a critical session handling bug across the session package and its child packages. ## Changes ### Bug Fix: waitForTCPIPForward fails on non-tcpip-forward global requests - Changed `waitForTCPIPForward()` from a single select to a for loop that drains non-tcpip-forward requests gracefully - Non-tcpip-forward requests now receive a false reply and the function continues waiting for the actual forward request - The 500ms timeout resets after each request to handle slow clients with multiple early requests - Updated and added tests to cover the new behavior: Wrong_Request_Type_Then_Timeout, Multiple_Non-Forward_Requests_Then_Success, Timeout_After_Non-Forward_Requests ### Race Condition Fixes - Added sync.RWMutex to the forwarder struct to protect concurrent field access - Fixed data race in forwarder setters (`SetType()`, `SetForwardedPort()`, `SetListener()`) by adding `Lock()`/`Unlock()` protection - Fixed data race in forwarder getters (`TunnelType()`, `ForwardedPort()`, `Listener()`) by adding `RLock()`/`RUnlock()` protection - Added sync.RWMutex to the slug struct to protect concurrent field access - Fixed data race in `slug.Set()` by adding `Lock()`/`Unlock()` protection - Fixed data race in `slug.String()` by adding `RLock()`/`RUnlock()` protection ### Concurrency Safety Improvements - Changed `OpenForwardedChannel()` to use `f.ForwardedPort()` getter instead of direct field access - Fixed `Close()` in forwarder to use `f.Listener()` getter consistently instead of accessing `f.listener` directly - Fixed data race in `HandleTCPForward()` where the goroutine running `tcpServer.Serve(listener)` wrote to the outer err variable from the enclosing function scope - Changed err = `tcpServer.Serve(listener)` to `if err := tcpServer.Serve(listener)` so the goroutine uses a local variable Reviewed-on: https://git.fossy.my.id/bagas/tunnel-please/pulls/151 Co-authored-by: Bagas Co-committed-by: Bagas --- internal/session/forwarder/forwarder.go | 19 +++++++-- internal/session/session.go | 34 +++++++-------- internal/session/session_test.go | 57 ++++++++++++++++++++++++- internal/session/slug/slug.go | 7 +++ 4 files changed, 95 insertions(+), 22 deletions(-) diff --git a/internal/session/forwarder/forwarder.go b/internal/session/forwarder/forwarder.go index 6af3bae..f59d315 100644 --- a/internal/session/forwarder/forwarder.go +++ b/internal/session/forwarder/forwarder.go @@ -28,6 +28,7 @@ type Forwarder interface { Close() error } type forwarder struct { + mu sync.RWMutex listener net.Listener tunnelType types.TunnelType forwardedPort uint16 @@ -60,7 +61,7 @@ func (f *forwarder) copyWithBuffer(dst io.Writer, src io.Reader) (written int64, } func (f *forwarder) OpenForwardedChannel(ctx context.Context, origin net.Addr) (ssh.Channel, <-chan *ssh.Request, error) { - payload := createForwardedTCPIPPayload(origin, f.forwardedPort) + payload := createForwardedTCPIPPayload(origin, f.ForwardedPort()) type channelResult struct { channel ssh.Channel reqs <-chan *ssh.Request @@ -141,32 +142,44 @@ func (f *forwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel) { } func (f *forwarder) SetType(tunnelType types.TunnelType) { + f.mu.Lock() + defer f.mu.Unlock() f.tunnelType = tunnelType } func (f *forwarder) TunnelType() types.TunnelType { + f.mu.RLock() + defer f.mu.RUnlock() return f.tunnelType } func (f *forwarder) ForwardedPort() uint16 { + f.mu.RLock() + defer f.mu.RUnlock() return f.forwardedPort } func (f *forwarder) SetForwardedPort(port uint16) { + f.mu.Lock() + defer f.mu.Unlock() f.forwardedPort = port } func (f *forwarder) SetListener(listener net.Listener) { + f.mu.Lock() + defer f.mu.Unlock() f.listener = listener } func (f *forwarder) Listener() net.Listener { + f.mu.RLock() + defer f.mu.RUnlock() return f.listener } func (f *forwarder) Close() error { - if f.Listener() != nil { - return f.listener.Close() + if listener := f.Listener(); listener != nil { + return listener.Close() } return nil } diff --git a/internal/session/session.go b/internal/session/session.go index b31e829..80b3db6 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -158,7 +158,7 @@ func (s *session) setupInteractiveMode(channel ssh.NewChannel) error { } go func() { - err = s.HandleGlobalRequest(reqs) + err := s.HandleGlobalRequest(reqs) if err != nil { log.Printf("global request handler error: %v", err) } @@ -200,23 +200,22 @@ func (s *session) waitForSessionEnd() error { } func (s *session) waitForTCPIPForward() *ssh.Request { - select { - case req, ok := <-s.initialReq: - if !ok { - log.Println("Forwarding request channel closed") + for { + select { + case req, ok := <-s.initialReq: + if !ok { + log.Println("Forwarding request channel closed") + return nil + } + if req.Type == "tcpip-forward" { + return req + } + log.Printf("Ignoring unexpected global request: %s", req.Type) + _ = req.Reply(false, nil) + case <-time.After(500 * time.Millisecond): + log.Println("No tcpip-forward request received within timeout") return nil } - if req.Type == "tcpip-forward" { - return req - } - if err := req.Reply(false, nil); err != nil { - log.Printf("Failed to reply to request: %v", err) - } - log.Printf("Expected tcpip-forward request, got: %s", req.Type) - return nil - case <-time.After(500 * time.Millisecond): - log.Println("No forwarding request received") - return nil } } @@ -379,8 +378,7 @@ func (s *session) HandleTCPForward(req *ssh.Request, addr string, portToBind uin } go func() { - err = tcpServer.Serve(listener) - if err != nil { + if err := tcpServer.Serve(listener); err != nil { log.Printf("Failed serving tcp server: %s\n", err) } }() diff --git a/internal/session/session_test.go b/internal/session/session_test.go index 9cff8be..c217947 100644 --- a/internal/session/session_test.go +++ b/internal/session/session_test.go @@ -807,7 +807,7 @@ func (m *mockNewChanFail) Accept() (ssh.Channel, <-chan *ssh.Request, error) { } func TestWaitForTCPIPForward_EdgeCases(t *testing.T) { - t.Run("Wrong Request Type", func(t *testing.T) { + t.Run("Wrong Request Type Then Timeout", func(t *testing.T) { _, sReqs, _, cConn, cleanup := setupSSH(t) defer cleanup() @@ -817,10 +817,65 @@ func TestWaitForTCPIPForward_EdgeCases(t *testing.T) { _, _, _ = cConn.SendRequest("not-tcpip-forward", true, nil) }() + start := time.Now() req := s.waitForTCPIPForward() + elapsed := time.Since(start) + if req != nil { t.Error("expected nil request") } + if elapsed < 400*time.Millisecond { + t.Errorf("expected timeout ~500ms, got %v", elapsed) + } + }) + + t.Run("Multiple Non-Forward Requests Then Success", func(t *testing.T) { + _, sReqs, _, cConn, cleanup := setupSSH(t) + defer cleanup() + + s := &session{initialReq: sReqs} + + go func() { + time.Sleep(100 * time.Millisecond) + _, _, _ = cConn.SendRequest("keepalive@openssh.com", false, nil) + time.Sleep(100 * time.Millisecond) + _, _, _ = cConn.SendRequest("hostkeys-00@openssh.com", false, nil) + time.Sleep(100 * time.Millisecond) + _, _, _ = cConn.SendRequest("tcpip-forward", true, nil) + }() + + req := s.waitForTCPIPForward() + if req == nil { + t.Error("expected tcpip-forward request, got nil") + } + if req != nil && req.Type != "tcpip-forward" { + t.Errorf("expected tcpip-forward, got %s", req.Type) + } + }) + + t.Run("Timeout After Non-Forward Requests", func(t *testing.T) { + _, sReqs, _, cConn, cleanup := setupSSH(t) + defer cleanup() + + s := &session{initialReq: sReqs} + + go func() { + time.Sleep(100 * time.Millisecond) + _, _, _ = cConn.SendRequest("keepalive@openssh.com", false, nil) + time.Sleep(100 * time.Millisecond) + _, _, _ = cConn.SendRequest("hostkeys-00@openssh.com", false, nil) + }() + + start := time.Now() + req := s.waitForTCPIPForward() + elapsed := time.Since(start) + + if req != nil { + t.Error("expected nil request after timeout") + } + if elapsed < 400*time.Millisecond { + t.Errorf("expected timeout ~500ms after last request, got %v", elapsed) + } }) t.Run("Channel Closed", func(t *testing.T) { diff --git a/internal/session/slug/slug.go b/internal/session/slug/slug.go index b9684d1..7fedc59 100644 --- a/internal/session/slug/slug.go +++ b/internal/session/slug/slug.go @@ -1,11 +1,14 @@ package slug +import "sync" + type Slug interface { String() string Set(slug string) } type slug struct { + mu sync.RWMutex slug string } @@ -16,9 +19,13 @@ func New() Slug { } func (s *slug) String() string { + s.mu.RLock() + defer s.mu.RUnlock() return s.slug } func (s *slug) Set(slug string) { + s.mu.Lock() + defer s.mu.Unlock() s.slug = slug }