refactor: separate session responsibilities and inject dependencies

session/session.go

@@ -1,8 +1,10 @@
 package session
 
 import (
+	"bytes"
 	"log"
 	"net"
+	"sync"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -16,38 +18,140 @@ const (
 type TunnelType string
 
 const (
-	HTTP    TunnelType = "http"
-	TCP     TunnelType = "tcp"
-	UNKNOWN TunnelType = "unknown"
+	HTTP TunnelType = "http"
+	TCP  TunnelType = "tcp"
 )
 
-type Session struct {
-	Connection    *ssh.ServerConn
-	ConnChannel   ssh.Channel
+type SessionLifecycle interface {
+	Close() error
+	WaitForRunningStatus()
+}
+
+type SessionCloser interface {
+	Close() error
+}
+
+type InteractionController interface {
+	SendMessage(message string)
+	HandleUserInput()
+	HandleCommand(conn ssh.Channel, command string, inSlugEditMode *bool, editSlug *string, buf *bytes.Buffer)
+	HandleSlugEditMode(conn ssh.Channel, inSlugEditMode *bool, editSlug *string, char byte, buf *bytes.Buffer)
+	HandleSlugSave(conn ssh.Channel, inSlugEditMode *bool, editSlug *string, buf *bytes.Buffer)
+	HandleSlugCancel(conn ssh.Channel, inSlugEditMode *bool, buf *bytes.Buffer)
+	HandleSlugUpdateError()
+	ShowWelcomeMessage()
+	DisplaySlugEditor()
+}
+
+type ForwardingController interface {
+	HandleGlobalRequest(ch <-chan *ssh.Request)
+	HandleTCPIPForward(req *ssh.Request)
+	HandleHTTPForward(req *ssh.Request, port uint16)
+	HandleTCPForward(req *ssh.Request, addr string, port uint16)
+	AcceptTCPConnections()
+	HandleForwardedConnection(conn UserConnection, sshConn *ssh.ServerConn)
+}
+
+type Session interface {
+	SessionLifecycle
+	InteractionController
+	ForwardingController
+}
+
+type Lifecycle struct {
+	Status SessionStatus
+}
+
+type Forwarder struct {
 	Listener      net.Listener
 	TunnelType    TunnelType
 	ForwardedPort uint16
-	Status        SessionStatus
-	Slug          string
+
+	getSlug func() string
+	setSlug func(string)
+}
+
+type ForwarderInfo interface {
+	GetTunnelType() TunnelType
+	GetForwardedPort() uint16
+}
+
+func (f *Forwarder) GetTunnelType() TunnelType {
+	return f.TunnelType
+}
+
+func (f *Forwarder) GetForwardedPort() uint16 {
+	return f.ForwardedPort
+}
+
+type Interaction struct {
+	CommandBuffer *bytes.Buffer
+	EditMode      bool
+	EditSlug      string
+	channel       ssh.Channel
+
+	getSlug func() string
+	setSlug func(string)
+
+	session SessionCloser
+
+	forwarder ForwarderInfo
+}
+type SSHSession struct {
+	lifecycle   *Lifecycle
+	interaction *Interaction
+	forwarder   *Forwarder
+
+	Conn    *ssh.ServerConn
+	channel ssh.Channel
+
+	slug   string
+	slugMu sync.RWMutex
 }
 
 func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan ssh.NewChannel) {
-	session := &Session{
-		Status:      INITIALIZING,
-		Slug:        "",
-		ConnChannel: nil,
-		Connection:  conn,
-		TunnelType:  UNKNOWN,
+	session := SSHSession{
+		lifecycle: &Lifecycle{
+			Status: INITIALIZING,
+		},
+		interaction: &Interaction{
+			CommandBuffer: new(bytes.Buffer),
+			EditMode:      false,
+			EditSlug:      "",
+			channel:       nil,
+			getSlug:       nil,
+			setSlug:       nil,
+			session:       nil,
+			forwarder:     nil,
+		},
+		forwarder: &Forwarder{
+			Listener:      nil,
+			TunnelType:    "",
+			ForwardedPort: 0,
+			getSlug:       nil,
+			setSlug:       nil,
+		},
+		Conn:    conn,
+		channel: nil,
+		slug:    "",
 	}
 
+	session.forwarder.getSlug = session.GetSlug
+	session.forwarder.setSlug = session.SetSlug
+	session.interaction.getSlug = session.GetSlug
+	session.interaction.setSlug = session.SetSlug
+	session.interaction.session = &session
+	session.interaction.forwarder = session.forwarder
+
 	go func() {
 		go session.waitForRunningStatus()
 
 		for channel := range sshChan {
 			ch, reqs, _ := channel.Accept()
-			if session.ConnChannel == nil {
-				session.ConnChannel = ch
-				session.Status = SETUP
+			if session.channel == nil {
+				session.channel = ch
+				session.interaction.channel = ch
+				session.lifecycle.Status = SETUP
 				go session.HandleGlobalRequest(forwardingReq)
 			}
 			go session.HandleGlobalRequest(reqs)
@@ -56,5 +160,18 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
 		if err != nil {
 			log.Printf("failed to close session: %v", err)
 		}
+		return
 	}()
 }
+
+func (s *SSHSession) GetSlug() string {
+	s.slugMu.RLock()
+	defer s.slugMu.RUnlock()
+	return s.slug
+}
+
+func (s *SSHSession) SetSlug(slug string) {
+	s.slugMu.Lock()
+	s.slug = slug
+	s.slugMu.Unlock()
+}