diff --git a/internal/session/session.go b/internal/session/session.go index b31e829..59eb2ad 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -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 } } 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) {