refactor: restructure session initialization to avoid circular references

This commit is contained in:
2025-12-04 19:32:00 +07:00
parent 82050a738f
commit 039e979142
10 changed files with 448 additions and 330 deletions

View File

@ -9,92 +9,24 @@ import (
"log"
"net"
"strconv"
"sync"
"time"
portUtil "tunnel_pls/internal/port"
"tunnel_pls/types"
"tunnel_pls/utils"
"golang.org/x/crypto/ssh"
)
type Status string
var forbiddenSlug = []string{
"ping",
}
type UserConnection struct {
Reader io.Reader
Writer net.Conn
}
var (
clientsMutex sync.RWMutex
Clients = make(map[string]*SSHSession)
)
func registerClient(slug string, session *SSHSession) bool {
clientsMutex.Lock()
defer clientsMutex.Unlock()
if _, exists := Clients[slug]; exists {
return false
}
Clients[slug] = session
return true
}
func unregisterClient(slug string) {
clientsMutex.Lock()
defer clientsMutex.Unlock()
delete(Clients, slug)
}
func (s *SSHSession) Close() error {
if s.Forwarder.Listener != nil {
err := s.Forwarder.Listener.Close()
if err != nil && !errors.Is(err, net.ErrClosed) {
return err
}
}
if s.channel != nil {
err := s.channel.Close()
if err != nil && !errors.Is(err, io.EOF) {
return err
}
}
if s.Conn != nil {
err := s.Conn.Close()
if err != nil && !errors.Is(err, net.ErrClosed) {
return err
}
}
slug := s.Forwarder.getSlug()
if slug != "" {
unregisterClient(slug)
}
if s.Forwarder.TunnelType == TCP && s.Forwarder.Listener != nil {
err := portUtil.Manager.SetPortStatus(s.Forwarder.ForwardedPort, false)
if err != nil {
return err
}
}
return nil
}
func (s *SSHSession) HandleGlobalRequest(GlobalRequest <-chan *ssh.Request) {
for req := range GlobalRequest {
switch req.Type {
case "tcpip-forward":
s.handleTCPIPForward(req)
s.HandleTCPIPForward(req)
return
case "shell", "pty-req", "window-change":
err := req.Reply(true, nil)
@ -113,7 +45,7 @@ func (s *SSHSession) HandleGlobalRequest(GlobalRequest <-chan *ssh.Request) {
}
}
func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
func (s *SSHSession) HandleTCPIPForward(req *ssh.Request) {
log.Println("Port forwarding request detected")
reader := bytes.NewReader(req.Payload)
@ -126,7 +58,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
log.Println("Failed to reply to request:", err)
return
}
err = s.Close()
err = s.Lifecycle.Close()
if err != nil {
log.Printf("failed to close session: %v", err)
}
@ -142,7 +74,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
log.Println("Failed to reply to request:", err)
return
}
err = s.Close()
err = s.Lifecycle.Close()
if err != nil {
log.Printf("failed to close session: %v", err)
}
@ -156,7 +88,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
log.Println("Failed to reply to request:", err)
return
}
err = s.Close()
err = s.Lifecycle.Close()
if err != nil {
log.Printf("failed to close session: %v", err)
}
@ -172,7 +104,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
log.Println("Failed to reply to request:", err)
return
}
err = s.Close()
err = s.Lifecycle.Close()
if err != nil {
log.Printf("failed to close session: %v", err)
}
@ -180,11 +112,11 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
}
s.Interaction.SendMessage("\033[H\033[2J")
s.Lifecycle.Status = RUNNING
s.Lifecycle.SetStatus(types.RUNNING)
go s.Interaction.HandleUserInput()
if portToBind == 80 || portToBind == 443 {
s.handleHTTPForward(req, portToBind)
s.HandleHTTPForward(req, portToBind)
return
} else {
if portToBind == 0 {
@ -197,7 +129,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
log.Println("Failed to reply to request:", err)
return
}
err = s.Close()
err = s.Lifecycle.Close()
if err != nil {
log.Printf("failed to close session: %v", err)
}
@ -210,7 +142,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
log.Println("Failed to reply to request:", err)
return
}
err = s.Close()
err = s.Lifecycle.Close()
if err != nil {
log.Printf("failed to close session: %v", err)
}
@ -222,7 +154,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
return
}
}
s.handleTCPForward(req, addr, portToBind)
s.HandleTCPForward(req, addr, portToBind)
}
var blockedReservedPorts = []uint16{1080, 1433, 1521, 1900, 2049, 3306, 3389, 5432, 5900, 6379, 8080, 8443, 9000, 9200, 27017}
@ -242,9 +174,9 @@ func isBlockedPort(port uint16) bool {
return false
}
func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) {
s.Forwarder.TunnelType = HTTP
s.Forwarder.ForwardedPort = portToBind
func (s *SSHSession) HandleHTTPForward(req *ssh.Request, portToBind uint16) {
s.Forwarder.SetType(types.HTTP)
s.Forwarder.SetForwardedPort(portToBind)
slug := generateUniqueSlug()
if slug == "" {
@ -256,7 +188,7 @@ func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) {
return
}
s.Forwarder.setSlug(slug)
s.SlugManager.Set(slug)
registerClient(slug, s)
buf := new(bytes.Buffer)
@ -282,8 +214,8 @@ func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) {
}
}
func (s *SSHSession) handleTCPForward(req *ssh.Request, addr string, portToBind uint16) {
s.Forwarder.TunnelType = TCP
func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind uint16) {
s.Forwarder.SetType(types.TCP)
log.Printf("Requested forwarding on %s:%d", addr, portToBind)
listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", portToBind))
@ -294,16 +226,16 @@ func (s *SSHSession) handleTCPForward(req *ssh.Request, addr string, portToBind
log.Println("Failed to reply to request:", err)
return
}
err = s.Close()
err = s.Lifecycle.Close()
if err != nil {
log.Printf("failed to close session: %v", err)
}
return
}
s.Forwarder.Listener = listener
s.Forwarder.ForwardedPort = portToBind
s.Forwarder.SetListener(listener)
s.Forwarder.SetForwardedPort(portToBind)
s.Interaction.ShowWelcomeMessage()
s.Interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to %s://%s:%d \r\n", s.Forwarder.TunnelType, utils.Getenv("domain"), s.Forwarder.ForwardedPort))
s.Interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to %s://%s:%d \r\n", s.Forwarder.GetTunnelType(), utils.Getenv("domain"), s.Forwarder.GetForwardedPort()))
go s.acceptTCPConnections()
@ -323,7 +255,7 @@ func (s *SSHSession) handleTCPForward(req *ssh.Request, addr string, portToBind
func (s *SSHSession) acceptTCPConnections() {
for {
conn, err := s.Forwarder.Listener.Accept()
conn, err := s.Forwarder.GetListener().Accept()
if err != nil {
if errors.Is(err, net.ErrClosed) {
return
@ -333,7 +265,7 @@ func (s *SSHSession) acceptTCPConnections() {
}
originHost, originPort := ParseAddr(conn.RemoteAddr().String())
payload := createForwardedTCPIPPayload(originHost, uint16(originPort), s.Forwarder.GetForwardedPort())
channel, reqs, err := s.Conn.OpenChannel("forwarded-tcpip", payload)
channel, reqs, err := s.Lifecycle.GetConnection().OpenChannel("forwarded-tcpip", payload)
if err != nil {
log.Printf("Failed to open forwarded-tcpip channel: %v", err)
return
@ -371,71 +303,6 @@ func generateUniqueSlug() string {
return ""
}
func (s *SSHSession) waitForRunningStatus() {
timeout := time.After(3 * time.Second)
ticker := time.NewTicker(150 * time.Millisecond)
defer ticker.Stop()
frames := []string{"-", "\\", "|", "/"}
i := 0
for {
select {
case <-ticker.C:
s.Interaction.SendMessage(fmt.Sprintf("\rLoading %s", frames[i]))
i = (i + 1) % len(frames)
if s.Lifecycle.Status == RUNNING {
s.Interaction.SendMessage("\r\033[K")
return
}
case <-timeout:
s.Interaction.SendMessage("\r\033[K")
s.Interaction.SendMessage("TCP/IP request not received in time.\r\nCheck your internet connection and confirm the server responds within 3000ms.\r\nEnsure you ran the correct command. For more details, visit https://tunnl.live.\r\n\r\n")
err := s.Close()
if err != nil {
log.Printf("failed to close session: %v", err)
}
log.Println("Timeout waiting for session to start running")
return
}
}
}
func isForbiddenSlug(slug string) bool {
for _, s := range forbiddenSlug {
if slug == s {
return true
}
}
return false
}
func isValidSlug(slug string) bool {
if len(slug) < 3 || len(slug) > 20 {
return false
}
if slug[0] == '-' || slug[len(slug)-1] == '-' {
return false
}
for _, c := range slug {
if !((c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-') {
return false
}
}
return true
}
func waitForKeyPress(connection ssh.Channel) {
keyBuf := make([]byte, 1)
for {
_, err := connection.Read(keyBuf)
if err == nil {
break
}
}
}
func (s *SSHSession) HandleForwardedConnection(dst io.ReadWriter, src ssh.Channel, remoteAddr net.Addr) {
defer func(src ssh.Channel) {
err := src.Close()