fix(session): resolve race conditions #151

Merged
bagas merged 4 commits from fix/session-race-conditions into staging 2026-07-17 22:22:25 +07:00
4 changed files with 95 additions and 22 deletions
+16 -3
View File
@@ -28,6 +28,7 @@ type Forwarder interface {
Close() error Close() error
} }
type forwarder struct { type forwarder struct {
mu sync.RWMutex
listener net.Listener listener net.Listener
tunnelType types.TunnelType tunnelType types.TunnelType
forwardedPort uint16 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) { 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 { type channelResult struct {
channel ssh.Channel channel ssh.Channel
reqs <-chan *ssh.Request 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) { func (f *forwarder) SetType(tunnelType types.TunnelType) {
f.mu.Lock()
defer f.mu.Unlock()
f.tunnelType = tunnelType f.tunnelType = tunnelType
} }
func (f *forwarder) TunnelType() types.TunnelType { func (f *forwarder) TunnelType() types.TunnelType {
f.mu.RLock()
defer f.mu.RUnlock()
return f.tunnelType return f.tunnelType
} }
func (f *forwarder) ForwardedPort() uint16 { func (f *forwarder) ForwardedPort() uint16 {
f.mu.RLock()
defer f.mu.RUnlock()
return f.forwardedPort return f.forwardedPort
} }
func (f *forwarder) SetForwardedPort(port uint16) { func (f *forwarder) SetForwardedPort(port uint16) {
f.mu.Lock()
defer f.mu.Unlock()
f.forwardedPort = port f.forwardedPort = port
} }
func (f *forwarder) SetListener(listener net.Listener) { func (f *forwarder) SetListener(listener net.Listener) {
f.mu.Lock()
defer f.mu.Unlock()
f.listener = listener f.listener = listener
} }
func (f *forwarder) Listener() net.Listener { func (f *forwarder) Listener() net.Listener {
f.mu.RLock()
defer f.mu.RUnlock()
return f.listener return f.listener
} }
func (f *forwarder) Close() error { func (f *forwarder) Close() error {
if f.Listener() != nil { if listener := f.Listener(); listener != nil {
return f.listener.Close() return listener.Close()
} }
return nil return nil
} }
+7 -9
View File
@@ -158,7 +158,7 @@ func (s *session) setupInteractiveMode(channel ssh.NewChannel) error {
} }
go func() { go func() {
err = s.HandleGlobalRequest(reqs) err := s.HandleGlobalRequest(reqs)
if err != nil { if err != nil {
log.Printf("global request handler error: %v", err) log.Printf("global request handler error: %v", err)
} }
@@ -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,15 +210,13 @@ 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 {
@@ -379,8 +378,7 @@ func (s *session) HandleTCPForward(req *ssh.Request, addr string, portToBind uin
} }
go func() { go func() {
err = tcpServer.Serve(listener) if err := tcpServer.Serve(listener); err != nil {
if err != nil {
log.Printf("Failed serving tcp server: %s\n", err) log.Printf("Failed serving tcp server: %s\n", err)
} }
}() }()
+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) {
+7
View File
@@ -1,11 +1,14 @@
package slug package slug
import "sync"
type Slug interface { type Slug interface {
String() string String() string
Set(slug string) Set(slug string)
} }
type slug struct { type slug struct {
mu sync.RWMutex
slug string slug string
} }
@@ -16,9 +19,13 @@ func New() Slug {
} }
func (s *slug) String() string { func (s *slug) String() string {
s.mu.RLock()
defer s.mu.RUnlock()
return s.slug return s.slug
} }
func (s *slug) Set(slug string) { func (s *slug) Set(slug string) {
s.mu.Lock()
defer s.mu.Unlock()
s.slug = slug s.slug = slug
} }