fix: wait for initial payload

This commit is contained in:
2026-07-16 19:23:09 +07:00
parent e8be839dac
commit 3c6e44cd0a
2 changed files with 70 additions and 16 deletions
+5 -6
View File
@@ -200,6 +200,7 @@ func (s *session) waitForSessionEnd() error {
} }
func (s *session) waitForTCPIPForward() *ssh.Request { func (s *session) waitForTCPIPForward() *ssh.Request {
for {
select { select {
case req, ok := <-s.initialReq: case req, ok := <-s.initialReq:
if !ok { if !ok {
@@ -209,16 +210,14 @@ func (s *session) waitForTCPIPForward() *ssh.Request {
if req.Type == "tcpip-forward" { if req.Type == "tcpip-forward" {
return req return req
} }
if err := req.Reply(false, nil); err != nil { log.Printf("Ignoring unexpected global request: %s", req.Type)
log.Printf("Failed to reply to request: %v", err) _ = req.Reply(false, nil)
}
log.Printf("Expected tcpip-forward request, got: %s", req.Type)
return nil
case <-time.After(500 * time.Millisecond): case <-time.After(500 * time.Millisecond):
log.Println("No forwarding request received") log.Println("No tcpip-forward request received within timeout")
return nil return nil
} }
} }
}
func (s *session) handleWindowChange(req *ssh.Request) error { func (s *session) handleWindowChange(req *ssh.Request) error {
p := req.Payload p := req.Payload
+56 -1
View File
@@ -807,7 +807,7 @@ func (m *mockNewChanFail) Accept() (ssh.Channel, <-chan *ssh.Request, error) {
} }
func TestWaitForTCPIPForward_EdgeCases(t *testing.T) { 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) _, sReqs, _, cConn, cleanup := setupSSH(t)
defer cleanup() defer cleanup()
@@ -817,10 +817,65 @@ func TestWaitForTCPIPForward_EdgeCases(t *testing.T) {
_, _, _ = cConn.SendRequest("not-tcpip-forward", true, nil) _, _, _ = cConn.SendRequest("not-tcpip-forward", true, nil)
}() }()
start := time.Now()
req := s.waitForTCPIPForward() req := s.waitForTCPIPForward()
elapsed := time.Since(start)
if req != nil { if req != nil {
t.Error("expected nil request") 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) { t.Run("Channel Closed", func(t *testing.T) {