Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 198147c676 | |||
| 1c591882cc | |||
| d9f3737797 | |||
| 3c6e44cd0a | |||
| e8be839dac |
@@ -13,7 +13,6 @@ import (
|
|||||||
"tunnel_pls/internal/session/slug"
|
"tunnel_pls/internal/session/slug"
|
||||||
"tunnel_pls/internal/types"
|
"tunnel_pls/internal/types"
|
||||||
|
|
||||||
"tunnel_pls/internal/port"
|
|
||||||
"tunnel_pls/internal/registry"
|
"tunnel_pls/internal/registry"
|
||||||
|
|
||||||
proto "git.fossy.my.id/bagas/tunnel-please-grpc/gen"
|
proto "git.fossy.my.id/bagas/tunnel-please-grpc/gen"
|
||||||
@@ -885,16 +884,16 @@ func (m *mockLifecycle) Connection() ssh.Conn {
|
|||||||
return args.Get(0).(ssh.Conn)
|
return args.Get(0).(ssh.Conn)
|
||||||
}
|
}
|
||||||
func (m *mockLifecycle) User() string { return m.Called().String(0) }
|
func (m *mockLifecycle) User() string { return m.Called().String(0) }
|
||||||
func (m *mockLifecycle) SetChannel(channel ssh.Channel) { m.Called(channel) }
|
func (m *mockLifecycle) SetChannel(channel ssh.Channel) error { return m.Called(channel).Error(0) }
|
||||||
func (m *mockLifecycle) SetStatus(status types.SessionStatus) { m.Called(status) }
|
func (m *mockLifecycle) SetStatus(status types.SessionStatus) { m.Called(status) }
|
||||||
func (m *mockLifecycle) IsActive() bool { return m.Called().Bool(0) }
|
func (m *mockLifecycle) IsActive() bool { return m.Called().Bool(0) }
|
||||||
func (m *mockLifecycle) StartedAt() time.Time { return m.Called().Get(0).(time.Time) }
|
func (m *mockLifecycle) StartedAt() time.Time { return m.Called().Get(0).(time.Time) }
|
||||||
func (m *mockLifecycle) PortRegistry() port.Port {
|
func (m *mockLifecycle) PortRegistry() lifecycle.PortRegistry {
|
||||||
args := m.Called()
|
args := m.Called()
|
||||||
if args.Get(0) == nil {
|
if args.Get(0) == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
return args.Get(0).(port.Port)
|
return args.Get(0).(lifecycle.PortRegistry)
|
||||||
}
|
}
|
||||||
|
|
||||||
type mockEventServiceClient struct {
|
type mockEventServiceClient struct {
|
||||||
|
|||||||
@@ -4,7 +4,6 @@ import (
|
|||||||
"sync"
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
"tunnel_pls/internal/port"
|
|
||||||
"tunnel_pls/internal/session/forwarder"
|
"tunnel_pls/internal/session/forwarder"
|
||||||
"tunnel_pls/internal/session/interaction"
|
"tunnel_pls/internal/session/interaction"
|
||||||
"tunnel_pls/internal/session/lifecycle"
|
"tunnel_pls/internal/session/lifecycle"
|
||||||
@@ -78,15 +77,15 @@ func (ml *mockLifecycle) Connection() ssh.Conn {
|
|||||||
return args.Get(0).(ssh.Conn)
|
return args.Get(0).(ssh.Conn)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (ml *mockLifecycle) PortRegistry() port.Port {
|
func (ml *mockLifecycle) PortRegistry() lifecycle.PortRegistry {
|
||||||
args := ml.Called()
|
args := ml.Called()
|
||||||
if args.Get(0) == nil {
|
if args.Get(0) == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
return args.Get(0).(port.Port)
|
return args.Get(0).(lifecycle.PortRegistry)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (ml *mockLifecycle) SetChannel(channel ssh.Channel) { ml.Called(channel) }
|
func (ml *mockLifecycle) SetChannel(channel ssh.Channel) error { return ml.Called(channel).Error(0) }
|
||||||
func (ml *mockLifecycle) SetStatus(status types.SessionStatus) { ml.Called(status) }
|
func (ml *mockLifecycle) SetStatus(status types.SessionStatus) { ml.Called(status) }
|
||||||
func (ml *mockLifecycle) IsActive() bool { return ml.Called().Bool(0) }
|
func (ml *mockLifecycle) IsActive() bool { return ml.Called().Bool(0) }
|
||||||
func (ml *mockLifecycle) StartedAt() time.Time { return ml.Called().Get(0).(time.Time) }
|
func (ml *mockLifecycle) StartedAt() time.Time { return ml.Called().Get(0).(time.Time) }
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1922,10 +1922,6 @@ func TestInteraction_Start_ProtocolSelection(t *testing.T) {
|
|||||||
time.Sleep(50 * time.Millisecond)
|
time.Sleep(50 * time.Millisecond)
|
||||||
|
|
||||||
i := mockInteraction.(*interaction)
|
i := mockInteraction.(*interaction)
|
||||||
if i.program != nil {
|
|
||||||
assert.NotNil(t, i.program, "program should be initialized")
|
|
||||||
}
|
|
||||||
|
|
||||||
i.Stop()
|
i.Stop()
|
||||||
|
|
||||||
mockConfig.AssertExpectations(t)
|
mockConfig.AssertExpectations(t)
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package lifecycle
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"errors"
|
"errors"
|
||||||
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net"
|
"net"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -9,8 +10,6 @@ import (
|
|||||||
"tunnel_pls/internal/session/slug"
|
"tunnel_pls/internal/session/slug"
|
||||||
"tunnel_pls/internal/types"
|
"tunnel_pls/internal/types"
|
||||||
|
|
||||||
portUtil "tunnel_pls/internal/port"
|
|
||||||
|
|
||||||
"golang.org/x/crypto/ssh"
|
"golang.org/x/crypto/ssh"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -24,6 +23,12 @@ type SessionRegistry interface {
|
|||||||
Remove(key types.SessionKey)
|
Remove(key types.SessionKey)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type PortRegistry interface {
|
||||||
|
Unassigned() (uint16, bool)
|
||||||
|
Claim(port uint16) bool
|
||||||
|
SetStatus(port uint16, assigned bool) error
|
||||||
|
}
|
||||||
|
|
||||||
type lifecycle struct {
|
type lifecycle struct {
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
status types.SessionStatus
|
status types.SessionStatus
|
||||||
@@ -34,18 +39,18 @@ type lifecycle struct {
|
|||||||
slug slug.Slug
|
slug slug.Slug
|
||||||
startedAt time.Time
|
startedAt time.Time
|
||||||
sessionRegistry SessionRegistry
|
sessionRegistry SessionRegistry
|
||||||
portRegistry portUtil.Port
|
portRegistry PortRegistry
|
||||||
user string
|
user string
|
||||||
}
|
}
|
||||||
|
|
||||||
func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port portUtil.Port, sessionRegistry SessionRegistry, user string) Lifecycle {
|
func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port PortRegistry, sessionRegistry SessionRegistry, user string) Lifecycle {
|
||||||
return &lifecycle{
|
return &lifecycle{
|
||||||
status: types.SessionStatusINITIALIZING,
|
status: types.SessionStatusINITIALIZING,
|
||||||
conn: conn,
|
conn: conn,
|
||||||
channel: nil,
|
channel: nil,
|
||||||
forwarder: forwarder,
|
forwarder: forwarder,
|
||||||
slug: slugManager,
|
slug: slugManager,
|
||||||
startedAt: time.Now(),
|
startedAt: time.Time{},
|
||||||
sessionRegistry: sessionRegistry,
|
sessionRegistry: sessionRegistry,
|
||||||
portRegistry: port,
|
portRegistry: port,
|
||||||
user: user,
|
user: user,
|
||||||
@@ -55,16 +60,16 @@ func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port portUti
|
|||||||
type Lifecycle interface {
|
type Lifecycle interface {
|
||||||
Connection() ssh.Conn
|
Connection() ssh.Conn
|
||||||
Channel() ssh.Channel
|
Channel() ssh.Channel
|
||||||
PortRegistry() portUtil.Port
|
PortRegistry() PortRegistry
|
||||||
User() string
|
User() string
|
||||||
SetChannel(channel ssh.Channel)
|
SetChannel(channel ssh.Channel) error
|
||||||
SetStatus(status types.SessionStatus)
|
SetStatus(status types.SessionStatus)
|
||||||
IsActive() bool
|
IsActive() bool
|
||||||
StartedAt() time.Time
|
StartedAt() time.Time
|
||||||
Close() error
|
Close() error
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l *lifecycle) PortRegistry() portUtil.Port {
|
func (l *lifecycle) PortRegistry() PortRegistry {
|
||||||
return l.portRegistry
|
return l.portRegistry
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -72,11 +77,25 @@ func (l *lifecycle) User() string {
|
|||||||
return l.user
|
return l.user
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l *lifecycle) SetChannel(channel ssh.Channel) {
|
func (l *lifecycle) SetChannel(channel ssh.Channel) error {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
if l.status == types.SessionStatusCLOSED {
|
||||||
|
return fmt.Errorf("lifecycle is closed")
|
||||||
|
}
|
||||||
|
if channel == nil {
|
||||||
|
return fmt.Errorf("channel cannot be nil")
|
||||||
|
}
|
||||||
|
if l.channel != nil {
|
||||||
|
return fmt.Errorf("channel already set")
|
||||||
|
}
|
||||||
l.channel = channel
|
l.channel = channel
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l *lifecycle) Channel() ssh.Channel {
|
func (l *lifecycle) Channel() ssh.Channel {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
return l.channel
|
return l.channel
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -87,7 +106,13 @@ func (l *lifecycle) Connection() ssh.Conn {
|
|||||||
func (l *lifecycle) SetStatus(status types.SessionStatus) {
|
func (l *lifecycle) SetStatus(status types.SessionStatus) {
|
||||||
l.mu.Lock()
|
l.mu.Lock()
|
||||||
defer l.mu.Unlock()
|
defer l.mu.Unlock()
|
||||||
|
if l.status == types.SessionStatusCLOSED {
|
||||||
|
return
|
||||||
|
}
|
||||||
l.status = status
|
l.status = status
|
||||||
|
if status == types.SessionStatusRUNNING && l.startedAt.IsZero() {
|
||||||
|
l.startedAt = time.Now()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l *lifecycle) IsActive() bool {
|
func (l *lifecycle) IsActive() bool {
|
||||||
@@ -98,50 +123,74 @@ func (l *lifecycle) IsActive() bool {
|
|||||||
|
|
||||||
func (l *lifecycle) Close() error {
|
func (l *lifecycle) Close() error {
|
||||||
l.mu.Lock()
|
l.mu.Lock()
|
||||||
defer l.mu.Unlock()
|
|
||||||
if l.status == types.SessionStatusCLOSED {
|
if l.status == types.SessionStatusCLOSED {
|
||||||
return l.closeErr
|
closeErr := l.closeErr
|
||||||
|
l.mu.Unlock()
|
||||||
|
return closeErr
|
||||||
}
|
}
|
||||||
l.status = types.SessionStatusCLOSED
|
l.status = types.SessionStatusCLOSED
|
||||||
|
|
||||||
|
channel := l.channel
|
||||||
|
conn := l.conn
|
||||||
|
l.mu.Unlock()
|
||||||
|
|
||||||
var errs []error
|
var errs []error
|
||||||
tunnelType := l.forwarder.TunnelType()
|
if channel != nil {
|
||||||
|
if err := channel.Close(); err != nil && !isClosedError(err) {
|
||||||
if l.channel != nil {
|
errs = append(errs, err)
|
||||||
if err := l.channel.Close(); err != nil && !isClosedError(err) {
|
}
|
||||||
|
}
|
||||||
|
if conn != nil {
|
||||||
|
if err := conn.Close(); err != nil && !isClosedError(err) {
|
||||||
errs = append(errs, err)
|
errs = append(errs, err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if l.conn != nil {
|
l.cleanupRegistry()
|
||||||
if err := l.conn.Close(); err != nil && !isClosedError(err) {
|
if err := l.cleanupForwarder(); err != nil {
|
||||||
errs = append(errs, err)
|
errs = append(errs, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
closeErr := errors.Join(errs...)
|
||||||
|
|
||||||
|
l.mu.Lock()
|
||||||
|
l.closeErr = closeErr
|
||||||
|
l.mu.Unlock()
|
||||||
|
|
||||||
|
return closeErr
|
||||||
}
|
}
|
||||||
|
|
||||||
clientSlug := l.slug.String()
|
func (l *lifecycle) cleanupRegistry() {
|
||||||
|
slugStr := l.slug.String()
|
||||||
|
if slugStr == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
key := types.SessionKey{
|
key := types.SessionKey{
|
||||||
Id: clientSlug,
|
Id: slugStr,
|
||||||
Type: tunnelType,
|
Type: l.forwarder.TunnelType(),
|
||||||
}
|
}
|
||||||
l.sessionRegistry.Remove(key)
|
l.sessionRegistry.Remove(key)
|
||||||
|
|
||||||
if tunnelType == types.TunnelTypeTCP {
|
|
||||||
errs = append(errs, l.PortRegistry().SetStatus(l.forwarder.ForwardedPort(), false))
|
|
||||||
errs = append(errs, l.forwarder.Close())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
l.closeErr = errors.Join(errs...)
|
func (l *lifecycle) cleanupForwarder() error {
|
||||||
return l.closeErr
|
if l.forwarder.TunnelType() != types.TunnelTypeTCP {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var errs []error
|
||||||
|
errs = append(errs, l.portRegistry.SetStatus(l.forwarder.ForwardedPort(), false))
|
||||||
|
errs = append(errs, l.forwarder.Close())
|
||||||
|
return errors.Join(errs...)
|
||||||
}
|
}
|
||||||
|
|
||||||
func isClosedError(err error) bool {
|
func isClosedError(err error) bool {
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
return errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || err.Error() == "EOF"
|
return errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l *lifecycle) StartedAt() time.Time {
|
func (l *lifecycle) StartedAt() time.Time {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
return l.startedAt
|
return l.startedAt
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import (
|
|||||||
"errors"
|
"errors"
|
||||||
"io"
|
"io"
|
||||||
"net"
|
"net"
|
||||||
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
"tunnel_pls/internal/types"
|
"tunnel_pls/internal/types"
|
||||||
|
|
||||||
@@ -177,8 +178,14 @@ func TestLifecycle_SetChannel(t *testing.T) {
|
|||||||
|
|
||||||
mockSSHChannel := &MockSSHChannel{}
|
mockSSHChannel := &MockSSHChannel{}
|
||||||
|
|
||||||
mockLifecycle.SetChannel(mockSSHChannel)
|
err := mockLifecycle.SetChannel(mockSSHChannel)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
assert.Equal(t, mockSSHChannel, mockLifecycle.Channel())
|
||||||
|
|
||||||
|
anotherChannel := &MockSSHChannel{}
|
||||||
|
err = mockLifecycle.SetChannel(anotherChannel)
|
||||||
|
assert.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "channel already set")
|
||||||
assert.Equal(t, mockSSHChannel, mockLifecycle.Channel())
|
assert.Equal(t, mockSSHChannel, mockLifecycle.Channel())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -276,14 +283,15 @@ func TestLifecycle_Close(t *testing.T) {
|
|||||||
mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
|
mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
|
||||||
|
|
||||||
mockLifecycle.SetStatus(types.SessionStatusRUNNING)
|
mockLifecycle.SetStatus(types.SessionStatusRUNNING)
|
||||||
mockLifecycle.SetChannel(mockSSHChannel)
|
err := mockLifecycle.SetChannel(mockSSHChannel)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
if tt.alreadyClosed {
|
if tt.alreadyClosed {
|
||||||
err := mockLifecycle.Close()
|
err = mockLifecycle.Close()
|
||||||
assert.NoError(t, err)
|
assert.NoError(t, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
err := mockLifecycle.Close()
|
err = mockLifecycle.Close()
|
||||||
|
|
||||||
if tt.expectErr {
|
if tt.expectErr {
|
||||||
assert.Error(t, err)
|
assert.Error(t, err)
|
||||||
@@ -301,3 +309,116 @@ func TestLifecycle_Close(t *testing.T) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestLifecycle_ConcurrentClose(t *testing.T) {
|
||||||
|
mockSSHConn := &MockSSHConn{}
|
||||||
|
mockSSHConn.On("Close").Return(nil)
|
||||||
|
|
||||||
|
mockForwarder := &MockForwarder{}
|
||||||
|
mockForwarder.On("TunnelType").Return(types.TunnelTypeHTTP)
|
||||||
|
|
||||||
|
mockSlug := &MockSlug{}
|
||||||
|
mockSlug.On("String").Return("test-slug")
|
||||||
|
|
||||||
|
mockPort := &MockPort{}
|
||||||
|
|
||||||
|
mockSessionRegistry := &MockSessionRegistry{}
|
||||||
|
mockSessionRegistry.On("Remove", mock.Anything).Return()
|
||||||
|
|
||||||
|
mockSSHChannel := &MockSSHChannel{}
|
||||||
|
mockSSHChannel.On("Close").Return(nil)
|
||||||
|
|
||||||
|
mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
|
||||||
|
mockLifecycle.SetStatus(types.SessionStatusRUNNING)
|
||||||
|
err := mockLifecycle.SetChannel(mockSSHChannel)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
const numGoroutines = 10
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
errChan := make(chan error, numGoroutines)
|
||||||
|
|
||||||
|
for i := 0; i < numGoroutines; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
err := mockLifecycle.Close()
|
||||||
|
errChan <- err
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
close(errChan)
|
||||||
|
|
||||||
|
for err := range errChan {
|
||||||
|
assert.NoError(t, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
assert.False(t, mockLifecycle.IsActive())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLifecycle_SetChannel_AfterClose(t *testing.T) {
|
||||||
|
mockSSHConn := new(MockSSHConn)
|
||||||
|
mockSSHConn.On("Close").Return(nil)
|
||||||
|
mockForwarder := &MockForwarder{}
|
||||||
|
mockForwarder.On("TunnelType").Return(types.TunnelTypeHTTP)
|
||||||
|
mockSlug := &MockSlug{}
|
||||||
|
mockSlug.On("String").Return("test-slug")
|
||||||
|
mockPort := &MockPort{}
|
||||||
|
mockSessionRegistry := &MockSessionRegistry{}
|
||||||
|
mockSessionRegistry.On("Remove", mock.Anything).Return()
|
||||||
|
mockSSHChannel := &MockSSHChannel{}
|
||||||
|
mockSSHChannel.On("Close").Return(nil)
|
||||||
|
|
||||||
|
mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
|
||||||
|
mockLifecycle.SetStatus(types.SessionStatusRUNNING)
|
||||||
|
err := mockLifecycle.SetChannel(mockSSHChannel)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
err = mockLifecycle.Close()
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
anotherChannel := &MockSSHChannel{}
|
||||||
|
err = mockLifecycle.SetChannel(anotherChannel)
|
||||||
|
assert.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "lifecycle is closed")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLifecycle_SetChannel_Nil(t *testing.T) {
|
||||||
|
mockSSHConn := new(MockSSHConn)
|
||||||
|
mockForwarder := &MockForwarder{}
|
||||||
|
mockSlug := &MockSlug{}
|
||||||
|
mockPort := &MockPort{}
|
||||||
|
mockSessionRegistry := &MockSessionRegistry{}
|
||||||
|
|
||||||
|
mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
|
||||||
|
|
||||||
|
err := mockLifecycle.SetChannel(nil)
|
||||||
|
assert.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "channel cannot be nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLifecycle_SetStatus_AfterClose(t *testing.T) {
|
||||||
|
mockSSHConn := new(MockSSHConn)
|
||||||
|
mockSSHConn.On("Close").Return(nil)
|
||||||
|
mockForwarder := &MockForwarder{}
|
||||||
|
mockForwarder.On("TunnelType").Return(types.TunnelTypeHTTP)
|
||||||
|
mockSlug := &MockSlug{}
|
||||||
|
mockSlug.On("String").Return("test-slug")
|
||||||
|
mockPort := &MockPort{}
|
||||||
|
mockSessionRegistry := &MockSessionRegistry{}
|
||||||
|
mockSessionRegistry.On("Remove", mock.Anything).Return()
|
||||||
|
mockSSHChannel := &MockSSHChannel{}
|
||||||
|
mockSSHChannel.On("Close").Return(nil)
|
||||||
|
|
||||||
|
mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
|
||||||
|
mockLifecycle.SetStatus(types.SessionStatusRUNNING)
|
||||||
|
err := mockLifecycle.SetChannel(mockSSHChannel)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
err = mockLifecycle.Close()
|
||||||
|
assert.NoError(t, err)
|
||||||
|
assert.False(t, mockLifecycle.IsActive())
|
||||||
|
|
||||||
|
mockLifecycle.SetStatus(types.SessionStatusRUNNING)
|
||||||
|
assert.False(t, mockLifecycle.IsActive(), "SetStatus should be ignored after Close")
|
||||||
|
}
|
||||||
|
|||||||
+10
-10
@@ -158,13 +158,15 @@ 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)
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
|
|
||||||
s.lifecycle.SetChannel(ch)
|
if err = s.lifecycle.SetChannel(ch); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
s.interaction.SetChannel(ch)
|
s.interaction.SetChannel(ch)
|
||||||
s.interaction.SetMode(types.InteractiveModeINTERACTIVE)
|
s.interaction.SetMode(types.InteractiveModeINTERACTIVE)
|
||||||
|
|
||||||
@@ -198,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 {
|
||||||
@@ -207,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
|
||||||
@@ -377,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)
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
|
|||||||
@@ -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) {
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -55,8 +55,7 @@ func TestHTTPServer_Serve(t *testing.T) {
|
|||||||
|
|
||||||
go func() {
|
go func() {
|
||||||
time.Sleep(100 * time.Millisecond)
|
time.Sleep(100 * time.Millisecond)
|
||||||
err = listener.Close()
|
_ = listener.Close()
|
||||||
assert.NoError(t, err)
|
|
||||||
}()
|
}()
|
||||||
|
|
||||||
err = srv.Serve(listener)
|
err = srv.Serve(listener)
|
||||||
|
|||||||
@@ -63,8 +63,7 @@ func TestHTTPSServer_Serve(t *testing.T) {
|
|||||||
|
|
||||||
go func() {
|
go func() {
|
||||||
time.Sleep(100 * time.Millisecond)
|
time.Sleep(100 * time.Millisecond)
|
||||||
err = listener.Close()
|
_ = listener.Close()
|
||||||
assert.NoError(t, err)
|
|
||||||
}()
|
}()
|
||||||
|
|
||||||
err = srv.Serve(listener)
|
err = srv.Serve(listener)
|
||||||
|
|||||||
@@ -45,8 +45,7 @@ func TestTCPServer_Serve(t *testing.T) {
|
|||||||
|
|
||||||
go func() {
|
go func() {
|
||||||
time.Sleep(100 * time.Millisecond)
|
time.Sleep(100 * time.Millisecond)
|
||||||
err = listener.Close()
|
_ = listener.Close()
|
||||||
assert.NoError(t, err)
|
|
||||||
}()
|
}()
|
||||||
|
|
||||||
err = srv.Serve(listener)
|
err = srv.Serve(listener)
|
||||||
|
|||||||
Reference in New Issue
Block a user