refactor: move session package to internal

2026-03-30 11:37:59 +07:00
parent 7ecc56aa3c
commit ae71b46482
22 changed files with 28 additions and 28 deletions
@@ -0,0 +1,147 @@
+package lifecycle
+
+import (
+	"errors"
+	"io"
+	"net"
+	"sync"
+	"time"
+	"tunnel_pls/internal/session/slug"
+
+	portUtil "tunnel_pls/internal/port"
+	"tunnel_pls/types"
+
+	"golang.org/x/crypto/ssh"
+)
+
+type Forwarder interface {
+	Close() error
+	TunnelType() types.TunnelType
+	ForwardedPort() uint16
+}
+
+type SessionRegistry interface {
+	Remove(key types.SessionKey)
+}
+
+type lifecycle struct {
+	mu              sync.Mutex
+	status          types.SessionStatus
+	closeErr        error
+	conn            ssh.Conn
+	channel         ssh.Channel
+	forwarder       Forwarder
+	slug            slug.Slug
+	startedAt       time.Time
+	sessionRegistry SessionRegistry
+	portRegistry    portUtil.Port
+	user            string
+}
+
+func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port portUtil.Port, sessionRegistry SessionRegistry, user string) Lifecycle {
+	return &lifecycle{
+		status:          types.SessionStatusINITIALIZING,
+		conn:            conn,
+		channel:         nil,
+		forwarder:       forwarder,
+		slug:            slugManager,
+		startedAt:       time.Now(),
+		sessionRegistry: sessionRegistry,
+		portRegistry:    port,
+		user:            user,
+	}
+}
+
+type Lifecycle interface {
+	Connection() ssh.Conn
+	Channel() ssh.Channel
+	PortRegistry() portUtil.Port
+	User() string
+	SetChannel(channel ssh.Channel)
+	SetStatus(status types.SessionStatus)
+	IsActive() bool
+	StartedAt() time.Time
+	Close() error
+}
+
+func (l *lifecycle) PortRegistry() portUtil.Port {
+	return l.portRegistry
+}
+
+func (l *lifecycle) User() string {
+	return l.user
+}
+
+func (l *lifecycle) SetChannel(channel ssh.Channel) {
+	l.channel = channel
+}
+
+func (l *lifecycle) Channel() ssh.Channel {
+	return l.channel
+}
+
+func (l *lifecycle) Connection() ssh.Conn {
+	return l.conn
+}
+
+func (l *lifecycle) SetStatus(status types.SessionStatus) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	l.status = status
+}
+
+func (l *lifecycle) IsActive() bool {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.status == types.SessionStatusRUNNING
+}
+
+func (l *lifecycle) Close() error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if l.status == types.SessionStatusCLOSED {
+		return l.closeErr
+	}
+	l.status = types.SessionStatusCLOSED
+
+	var errs []error
+	tunnelType := l.forwarder.TunnelType()
+
+	if l.channel != nil {
+		if err := l.channel.Close(); err != nil && !isClosedError(err) {
+			errs = append(errs, err)
+		}
+	}
+
+	if l.conn != nil {
+		if err := l.conn.Close(); err != nil && !isClosedError(err) {
+			errs = append(errs, err)
+		}
+	}
+
+	clientSlug := l.slug.String()
+	key := types.SessionKey{
+		Id:   clientSlug,
+		Type: tunnelType,
+	}
+	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...)
+	return l.closeErr
+}
+
+func isClosedError(err error) bool {
+	if err == nil {
+		return false
+	}
+	return errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || err.Error() == "EOF"
+}
+
+func (l *lifecycle) StartedAt() time.Time {
+	return l.startedAt
+}
@@ -0,0 +1,303 @@
+package lifecycle
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"testing"
+	"tunnel_pls/types"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"golang.org/x/crypto/ssh"
+)
+
+type MockSessionRegistry struct {
+	mock.Mock
+}
+
+func (m *MockSessionRegistry) Remove(key types.SessionKey) {
+	m.Called(key)
+}
+
+type MockForwarder struct {
+	mock.Mock
+}
+
+func (m *MockForwarder) CreateForwardedTCPIPPayload(origin net.Addr) []byte {
+	args := m.Called(origin)
+	return args.Get(0).([]byte)
+}
+
+func (m *MockForwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel) {
+	m.Called(dst, src)
+}
+
+func (m *MockForwarder) Close() error {
+	args := m.Called()
+	return args.Error(0)
+}
+
+func (m *MockForwarder) TunnelType() types.TunnelType {
+	args := m.Called()
+	return args.Get(0).(types.TunnelType)
+}
+
+func (m *MockForwarder) ForwardedPort() uint16 {
+	args := m.Called()
+	return args.Get(0).(uint16)
+}
+
+func (m *MockForwarder) SetType(tunnelType types.TunnelType) {
+	m.Called(tunnelType)
+}
+
+func (m *MockForwarder) SetForwardedPort(port uint16) {
+	m.Called(port)
+}
+
+func (m *MockForwarder) SetListener(listener net.Listener) {
+	m.Called(listener)
+}
+
+func (m *MockForwarder) Listener() net.Listener {
+	args := m.Called()
+	return args.Get(0).(net.Listener)
+}
+
+func (m *MockForwarder) OpenForwardedChannel(ctx context.Context, origin net.Addr) (ssh.Channel, <-chan *ssh.Request, error) {
+	args := m.Called(ctx, origin)
+	if args.Get(0) == nil {
+		return nil, nil, args.Error(2)
+	}
+	return args.Get(0).(ssh.Channel), args.Get(1).(<-chan *ssh.Request), args.Error(2)
+}
+
+type MockPort struct {
+	mock.Mock
+}
+
+func (m *MockPort) AddRange(startPort, endPort uint16) error {
+	return m.Called(startPort, endPort).Error(0)
+}
+func (m *MockPort) Unassigned() (uint16, bool) {
+	args := m.Called()
+	var port uint16
+	if args.Get(0) != nil {
+		switch v := args.Get(0).(type) {
+		case int:
+			port = uint16(v)
+		case uint16:
+			port = v
+		case uint32:
+			port = uint16(v)
+		case int32:
+			port = uint16(v)
+		case float64:
+			port = uint16(v)
+		default:
+			port = uint16(args.Int(0))
+		}
+	}
+	return port, args.Bool(1)
+}
+func (m *MockPort) SetStatus(port uint16, assigned bool) error {
+	return m.Called(port, assigned).Error(0)
+}
+func (m *MockPort) Claim(port uint16) bool {
+	return m.Called(port).Bool(0)
+}
+
+type MockSlug struct {
+	mock.Mock
+}
+
+func (ms *MockSlug) Set(slug string) {
+	ms.Called(slug)
+}
+func (ms *MockSlug) String() string {
+	return ms.Called().String(0)
+}
+
+type MockSSHConn struct {
+	ssh.Conn
+	mock.Mock
+}
+
+func (m *MockSSHConn) Close() error {
+	args := m.Called()
+	return args.Error(0)
+}
+
+type MockSSHChannel struct {
+	ssh.Channel
+	mock.Mock
+}
+
+func (m *MockSSHChannel) Close() error {
+	return m.Called().Error(0)
+}
+
+func TestNew(t *testing.T) {
+	mockSSHConn := new(MockSSHConn)
+	mockForwarder := &MockForwarder{}
+	mockSlug := &MockSlug{}
+	mockPort := &MockPort{}
+	mockSessionRegistry := &MockSessionRegistry{}
+
+	mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
+
+	assert.NotNil(t, mockLifecycle.Connection())
+	assert.NotNil(t, mockLifecycle.User())
+	assert.NotNil(t, mockLifecycle.PortRegistry())
+	assert.NotNil(t, mockLifecycle.StartedAt())
+}
+
+func TestLifecycle_User(t *testing.T) {
+	mockSSHConn := new(MockSSHConn)
+	mockForwarder := &MockForwarder{}
+	mockSlug := &MockSlug{}
+	mockPort := &MockPort{}
+	mockSessionRegistry := &MockSessionRegistry{}
+
+	user := "mas-fuad"
+	mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, user)
+	assert.Equal(t, user, mockLifecycle.User())
+}
+
+func TestLifecycle_SetChannel(t *testing.T) {
+	mockSSHConn := new(MockSSHConn)
+	mockForwarder := &MockForwarder{}
+	mockSlug := &MockSlug{}
+	mockPort := &MockPort{}
+	mockSessionRegistry := &MockSessionRegistry{}
+
+	mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
+
+	mockSSHChannel := &MockSSHChannel{}
+
+	mockLifecycle.SetChannel(mockSSHChannel)
+
+	assert.Equal(t, mockSSHChannel, mockLifecycle.Channel())
+}
+
+func TestLifecycle_SetStatus(t *testing.T) {
+	mockSSHConn := new(MockSSHConn)
+	mockForwarder := &MockForwarder{}
+	mockSlug := &MockSlug{}
+	mockPort := &MockPort{}
+	mockSessionRegistry := &MockSessionRegistry{}
+
+	mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
+
+	mockLifecycle.SetStatus(types.SessionStatusRUNNING)
+	assert.True(t, mockLifecycle.IsActive())
+}
+
+func TestLifecycle_IsActive(t *testing.T) {
+	mockSSHConn := new(MockSSHConn)
+	mockForwarder := &MockForwarder{}
+	mockSlug := &MockSlug{}
+	mockPort := &MockPort{}
+	mockSessionRegistry := &MockSessionRegistry{}
+
+	mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
+
+	mockLifecycle.SetStatus(types.SessionStatusRUNNING)
+	assert.True(t, mockLifecycle.IsActive())
+}
+
+func TestLifecycle_Close(t *testing.T) {
+	tests := []struct {
+		name            string
+		tunnelType      types.TunnelType
+		connCloseErr    error
+		channelCloseErr error
+		expectErr       bool
+		alreadyClosed   bool
+	}{
+		{
+			name:       "Close HTTP forwarding success",
+			tunnelType: types.TunnelTypeHTTP,
+			expectErr:  false,
+		},
+		{
+			name:       "Close TCP forwarding success",
+			tunnelType: types.TunnelTypeTCP,
+			expectErr:  false,
+		},
+		{
+			name:         "Close with conn close error",
+			tunnelType:   types.TunnelTypeHTTP,
+			connCloseErr: errors.New("conn close error"),
+			expectErr:    true,
+		},
+		{
+			name:            "Close with channel close error",
+			tunnelType:      types.TunnelTypeHTTP,
+			channelCloseErr: errors.New("channel close error"),
+			expectErr:       true,
+		},
+		{
+			name:          "Close when already closed",
+			tunnelType:    types.TunnelTypeHTTP,
+			alreadyClosed: true,
+			expectErr:     false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockSSHConn := &MockSSHConn{}
+			mockSSHConn.On("Close").Return(tt.connCloseErr)
+
+			mockForwarder := &MockForwarder{}
+			mockForwarder.On("TunnelType").Return(tt.tunnelType)
+			if tt.tunnelType == types.TunnelTypeTCP {
+				mockForwarder.On("ForwardedPort").Return(uint16(8080))
+				mockForwarder.On("Close").Return(nil)
+			}
+
+			mockSlug := &MockSlug{}
+			mockSlug.On("String").Return("test-slug")
+
+			mockPort := &MockPort{}
+			if tt.tunnelType == types.TunnelTypeTCP {
+				mockPort.On("SetStatus", uint16(8080), false).Return(nil)
+			}
+
+			mockSessionRegistry := &MockSessionRegistry{}
+			mockSessionRegistry.On("Remove", mock.Anything).Return()
+
+			mockSSHChannel := &MockSSHChannel{}
+			mockSSHChannel.On("Close").Return(tt.channelCloseErr)
+
+			mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
+
+			mockLifecycle.SetStatus(types.SessionStatusRUNNING)
+			mockLifecycle.SetChannel(mockSSHChannel)
+
+			if tt.alreadyClosed {
+				err := mockLifecycle.Close()
+				assert.NoError(t, err)
+			}
+
+			err := mockLifecycle.Close()
+
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+			assert.False(t, mockLifecycle.IsActive())
+
+			mockSSHConn.AssertExpectations(t)
+			mockForwarder.AssertExpectations(t)
+			mockSlug.AssertExpectations(t)
+			mockPort.AssertExpectations(t)
+			mockSessionRegistry.AssertExpectations(t)
+			mockSSHChannel.AssertExpectations(t)
+		})
+	}
+}