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: #151
Co-authored-by: Bagas <bagas@fossy.my.id>
Co-committed-by: Bagas <bagas@fossy.my.id>

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
 }

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)
 		}
 	}()

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) {

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
 }