fix: correct logic when checking tcpip-forward request

session/session.go

@@ -2,13 +2,15 @@ package session
 
 import (
 	"bytes"
+	"fmt"
 	"log"
 	"sync"
+	"time"
 	"tunnel_pls/session/forwarder"
 	"tunnel_pls/session/interaction"
 	"tunnel_pls/session/lifecycle"
 	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/utils"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -30,8 +32,6 @@ type SSHSession struct {
 	Interaction interaction.Controller
 	Forwarder   forwarder.ForwardingController
 	SlugManager slug.Manager
-
-	channelOnce sync.Once
 }
 
 func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan ssh.NewChannel) {
@@ -71,20 +71,27 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
 		SlugManager: slugManager,
 	}
 
+	var once sync.Once
 	for channel := range sshChan {
 		ch, reqs, err := channel.Accept()
 		if err != nil {
 			log.Printf("failed to accept channel: %v", err)
 			continue
 		}
-		session.channelOnce.Do(func() {
+		once.Do(func() {
 			session.Lifecycle.SetChannel(ch)
 			session.Interaction.SetChannel(ch)
-			session.Lifecycle.SetStatus(types.SETUP)
-			go session.HandleGlobalRequest(forwardingReq)
-			session.Lifecycle.WaitForRunningStatus()
-		})
 
+			tcpipReq := session.waitForTCPIPForward(forwardingReq)
+			if tcpipReq == nil {
+				session.Interaction.SendMessage(fmt.Sprintf("Port forwarding request not received.\r\nEnsure you ran the correct command with -R flag.\r\nExample: ssh %s -p %s -R 80:localhost:3000\r\nFor more details, visit https://tunnl.live.\r\n\r\n", utils.Getenv("DOMAIN", "localhost"), utils.Getenv("PORT", "2200")))
+				if err := session.Lifecycle.Close(); err != nil {
+					log.Printf("failed to close session: %v", err)
+				}
+				return
+			}
+			session.HandleTCPIPForward(tcpipReq)
+		})
 		go session.HandleGlobalRequest(reqs)
 	}
 	if err := session.Lifecycle.Close(); err != nil {
@@ -92,6 +99,27 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
 	}
 }
 
+func (s *SSHSession) waitForTCPIPForward(forwardingReq <-chan *ssh.Request) *ssh.Request {
+	select {
+	case req, ok := <-forwardingReq:
+		if !ok {
+			log.Println("Forwarding request channel closed")
+			return nil
+		}
+		if req.Type == "tcpip-forward" {
+			return req
+		}
+		if err := req.Reply(false, nil); err != nil {
+			log.Printf("Failed to reply to request: %v", err)
+		}
+		log.Printf("Expected tcpip-forward request, got: %s", req.Type)
+		return nil
+	case <-time.After(500 * time.Millisecond):
+		log.Println("No forwarding request received")
+		return nil
+	}
+}
+
 func updateClientSlug(oldSlug, newSlug string) bool {
 	clientsMutex.Lock()
 	defer clientsMutex.Unlock()