refactor: restructure session initialization to avoid circular references
This commit is contained in:
@ -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()
|
||||
|
||||
Reference in New Issue
Block a user