bagas
fabfd96600
## Summary Comprehensive fixes for race conditions, concurrency bugs, and code quality improvements in the lifecycle module and related test files. ## Changes ### Race Condition Fixes - Fixed data race in `SetChannel()` and `Channel()` by adding mutex protection - Fixed data race in `StartedAt()` by adding mutex protection - Fixed test race conditions in `transport` and `interaction` packages where variables were shared between goroutines ### Concurrency Safety Improvements - Refactored `Close()` to release mutex before calling external code, preventing potential deadlocks - Made `SetChannel()` reject calls after `Close()` to prevent setting channel on torn-down lifecycle - Made `SetStatus()` ignore calls after `Close()` to ensure CLOSED is a terminal state - Added nil check to `SetChannel()` to prevent masking upstream bugs ### Code Quality Improvements - Refactored `Close()` to extract helper methods: `closeChannel()`, `closeConnection()`, `cleanupRegistry()`, `cleanupForwarder()` - Fixed `isClosedError()` to remove redundant string comparison - Added `PortRegistry` interface for better interface segregation (removed dependency on full `port.Port`) - Made `cleanupRegistry()` conditional on non-empty slug to avoid unnecessary registry calls - Fixed `startedAt` to be set when session starts running instead of at object creation Reviewed-on: #150 Co-authored-by: Bagas <bagas@fossy.my.id> Co-committed-by: Bagas <bagas@fossy.my.id> Reviewed-on: #152
146 lines
3.4 KiB
Go
146 lines
3.4 KiB
Go
package transport
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/mock"
|
|
"golang.org/x/crypto/ssh"
|
|
)
|
|
|
|
func TestNewTCPServer(t *testing.T) {
|
|
mf := new(MockForwarder)
|
|
port := uint16(9000)
|
|
|
|
srv := NewTCPServer(port, mf)
|
|
assert.NotNil(t, srv)
|
|
|
|
tcpSrv, ok := srv.(*tcp)
|
|
assert.True(t, ok)
|
|
assert.Equal(t, port, tcpSrv.port)
|
|
assert.Equal(t, mf, tcpSrv.forwarder)
|
|
}
|
|
|
|
func TestTCPServer_Listen(t *testing.T) {
|
|
mf := new(MockForwarder)
|
|
srv := NewTCPServer(0, mf)
|
|
|
|
listener, err := srv.Listen()
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, listener)
|
|
err = listener.Close()
|
|
assert.NoError(t, err)
|
|
}
|
|
|
|
func TestTCPServer_Serve(t *testing.T) {
|
|
mf := new(MockForwarder)
|
|
srv := NewTCPServer(0, mf)
|
|
|
|
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
|
assert.NoError(t, err)
|
|
|
|
go func() {
|
|
time.Sleep(100 * time.Millisecond)
|
|
_ = listener.Close()
|
|
}()
|
|
|
|
err = srv.Serve(listener)
|
|
assert.Nil(t, err)
|
|
}
|
|
|
|
func TestTCPServer_Serve_AcceptError(t *testing.T) {
|
|
mf := new(MockForwarder)
|
|
srv := NewTCPServer(0, mf)
|
|
|
|
ml := new(mockListener)
|
|
ml.On("Accept").Return(nil, errors.New("accept error")).Once()
|
|
ml.On("Accept").Return(nil, net.ErrClosed).Once()
|
|
|
|
err := srv.Serve(ml)
|
|
assert.Nil(t, err)
|
|
ml.AssertExpectations(t)
|
|
}
|
|
|
|
func TestTCPServer_Serve_Success(t *testing.T) {
|
|
mf := new(MockForwarder)
|
|
srv := NewTCPServer(0, mf)
|
|
|
|
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
|
assert.NoError(t, err)
|
|
port := listener.Addr().(*net.TCPAddr).Port
|
|
|
|
reqs := make(chan *ssh.Request)
|
|
mf.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(new(MockSSHChannel), (<-chan *ssh.Request)(reqs), nil)
|
|
mf.On("HandleConnection", mock.Anything, mock.Anything).Return()
|
|
|
|
go func() {
|
|
_ = srv.Serve(listener)
|
|
}()
|
|
|
|
conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", port))
|
|
assert.NoError(t, err)
|
|
|
|
time.Sleep(100 * time.Millisecond)
|
|
err = conn.Close()
|
|
assert.NoError(t, err)
|
|
err = listener.Close()
|
|
assert.NoError(t, err)
|
|
mf.AssertExpectations(t)
|
|
}
|
|
|
|
func TestTCPServer_handleTcp_Success(t *testing.T) {
|
|
mf := new(MockForwarder)
|
|
srv := NewTCPServer(0, mf).(*tcp)
|
|
|
|
serverConn, clientConn := net.Pipe()
|
|
defer func(clientConn net.Conn) {
|
|
err := clientConn.Close()
|
|
assert.NoError(t, err)
|
|
}(clientConn)
|
|
|
|
reqs := make(chan *ssh.Request)
|
|
mockChannel := new(MockSSHChannel)
|
|
mf.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(mockChannel, (<-chan *ssh.Request)(reqs), nil)
|
|
|
|
mf.On("HandleConnection", serverConn, mockChannel).Return()
|
|
|
|
srv.handleTcp(serverConn)
|
|
|
|
mf.AssertExpectations(t)
|
|
}
|
|
|
|
func TestTCPServer_handleTcp_CloseError(t *testing.T) {
|
|
mf := new(MockForwarder)
|
|
srv := NewTCPServer(0, mf).(*tcp)
|
|
|
|
mc := new(MockConn)
|
|
mc.On("Close").Return(errors.New("close error"))
|
|
mc.On("RemoteAddr").Return(&net.TCPAddr{})
|
|
|
|
mf.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(nil, (<-chan *ssh.Request)(nil), errors.New("open error"))
|
|
|
|
srv.handleTcp(mc)
|
|
mc.AssertExpectations(t)
|
|
}
|
|
|
|
func TestTCPServer_handleTcp_OpenChannelError(t *testing.T) {
|
|
mf := new(MockForwarder)
|
|
srv := NewTCPServer(0, mf).(*tcp)
|
|
|
|
serverConn, clientConn := net.Pipe()
|
|
defer func(clientConn net.Conn) {
|
|
err := clientConn.Close()
|
|
assert.NoError(t, err)
|
|
}(clientConn)
|
|
|
|
mf.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(nil, (<-chan *ssh.Request)(nil), errors.New("open error"))
|
|
|
|
srv.handleTcp(serverConn)
|
|
|
|
mf.AssertExpectations(t)
|
|
}
|