feat: add header factory

session/handler.go

@@ -1,9 +1,7 @@
 package session
 
 import (
-	"bufio"
 	"bytes"
-	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -56,8 +54,8 @@ func unregisterClient(slug string) {
 }
 
 func (s *SSHSession) Close() error {
-	if s.forwarder.Listener != nil {
-		err := s.forwarder.Listener.Close()
+	if s.Forwarder.Listener != nil {
+		err := s.Forwarder.Listener.Close()
 		if err != nil && !errors.Is(err, net.ErrClosed) {
 			return err
 		}
@@ -77,13 +75,13 @@ func (s *SSHSession) Close() error {
 		}
 	}
 
-	slug := s.forwarder.getSlug()
+	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 s.Forwarder.TunnelType == TCP && s.Forwarder.Listener != nil {
+		err := portUtil.Manager.SetPortStatus(s.Forwarder.ForwardedPort, false)
 		if err != nil {
 			return err
 		}
@@ -138,7 +136,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
 	var rawPortToBind uint32
 	if err := binary.Read(reader, binary.BigEndian, &rawPortToBind); err != nil {
 		log.Println("Failed to read port from payload:", err)
-		s.interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port. (02) \r\n", rawPortToBind))
+		s.Interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port. (02) \r\n", rawPortToBind))
 		err := req.Reply(false, nil)
 		if err != nil {
 			log.Println("Failed to reply to request:", err)
@@ -152,7 +150,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
 	}
 
 	if rawPortToBind > 65535 {
-		s.interaction.SendMessage(fmt.Sprintf("Port %d is larger then allowed port of 65535. (02)\r\n", rawPortToBind))
+		s.Interaction.SendMessage(fmt.Sprintf("Port %d is larger then allowed port of 65535. (02)\r\n", rawPortToBind))
 		err := req.Reply(false, nil)
 		if err != nil {
 			log.Println("Failed to reply to request:", err)
@@ -168,7 +166,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
 	portToBind := uint16(rawPortToBind)
 
 	if isBlockedPort(portToBind) {
-		s.interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port. (02)\r\n", portToBind))
+		s.Interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port. (02)\r\n", portToBind))
 		err := req.Reply(false, nil)
 		if err != nil {
 			log.Println("Failed to reply to request:", err)
@@ -181,9 +179,9 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
 		return
 	}
 
-	s.interaction.SendMessage("\033[H\033[2J")
-	s.lifecycle.Status = RUNNING
-	go s.interaction.HandleUserInput()
+	s.Interaction.SendMessage("\033[H\033[2J")
+	s.Lifecycle.Status = RUNNING
+	go s.Interaction.HandleUserInput()
 
 	if portToBind == 80 || portToBind == 443 {
 		s.handleHTTPForward(req, portToBind)
@@ -193,7 +191,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
 			unassign, success := portUtil.Manager.GetUnassignedPort()
 			portToBind = unassign
 			if !success {
-				s.interaction.SendMessage(fmt.Sprintf("No available port\r\n", portToBind))
+				s.Interaction.SendMessage(fmt.Sprintf("No available port\r\n", portToBind))
 				err := req.Reply(false, nil)
 				if err != nil {
 					log.Println("Failed to reply to request:", err)
@@ -206,7 +204,7 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) {
 				return
 			}
 		} else if isUse, isExist := portUtil.Manager.GetPortStatus(portToBind); isExist && isUse {
-			s.interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port. (03)\r\n", portToBind))
+			s.Interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port. (03)\r\n", portToBind))
 			err := req.Reply(false, nil)
 			if err != nil {
 				log.Println("Failed to reply to request:", err)
@@ -245,8 +243,8 @@ func isBlockedPort(port uint16) bool {
 }
 
 func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) {
-	s.forwarder.TunnelType = HTTP
-	s.forwarder.ForwardedPort = portToBind
+	s.Forwarder.TunnelType = HTTP
+	s.Forwarder.ForwardedPort = portToBind
 
 	slug := generateUniqueSlug()
 	if slug == "" {
@@ -258,7 +256,7 @@ func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) {
 		return
 	}
 
-	s.forwarder.setSlug(slug)
+	s.Forwarder.setSlug(slug)
 	registerClient(slug, s)
 
 	buf := new(bytes.Buffer)
@@ -275,8 +273,8 @@ func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) {
 		protocol = "https"
 	}
 
-	s.interaction.ShowWelcomeMessage()
-	s.interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to %s://%s.%s\r\n", protocol, slug, domain))
+	s.Interaction.ShowWelcomeMessage()
+	s.Interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to %s://%s.%s\r\n", protocol, slug, domain))
 	err = req.Reply(true, buf.Bytes())
 	if err != nil {
 		log.Println("Failed to reply to request:", err)
@@ -285,12 +283,12 @@ func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) {
 }
 
 func (s *SSHSession) handleTCPForward(req *ssh.Request, addr string, portToBind uint16) {
-	s.forwarder.TunnelType = TCP
+	s.Forwarder.TunnelType = TCP
 	log.Printf("Requested forwarding on %s:%d", addr, portToBind)
 
 	listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", portToBind))
 	if err != nil {
-		s.interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port.\r\n", portToBind))
+		s.Interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port.\r\n", portToBind))
 		err := req.Reply(false, nil)
 		if err != nil {
 			log.Println("Failed to reply to request:", err)
@@ -302,10 +300,10 @@ func (s *SSHSession) handleTCPForward(req *ssh.Request, addr string, portToBind
 		}
 		return
 	}
-	s.forwarder.Listener = listener
-	s.forwarder.ForwardedPort = 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.Forwarder.Listener = listener
+	s.Forwarder.ForwardedPort = 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))
 
 	go s.acceptTCPConnections()
 
@@ -325,7 +323,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.Listener.Accept()
 		if err != nil {
 			if errors.Is(err, net.ErrClosed) {
 				return
@@ -333,11 +331,35 @@ func (s *SSHSession) acceptTCPConnections() {
 			log.Printf("Error accepting connection: %v", err)
 			continue
 		}
+		originHost, originPort := ParseAddr(conn.RemoteAddr().String())
+		payload := createForwardedTCPIPPayload(originHost, uint16(originPort), s.Forwarder.GetForwardedPort())
+		channel, reqs, err := s.Conn.OpenChannel("forwarded-tcpip", payload)
+		if err != nil {
+			log.Printf("Failed to open forwarded-tcpip channel: %v", err)
+			return
+		}
+		defer func(channel ssh.Channel) {
+			err := channel.Close()
+			if err != nil {
+				log.Println("Failed to close connection:", err)
+			}
+		}(channel)
 
-		go s.HandleForwardedConnection(UserConnection{
-			Reader: nil,
-			Writer: conn,
-		}, s.Conn)
+		go func() {
+			defer func() {
+				if r := recover(); r != nil {
+					log.Printf("Panic in request handler: %v", r)
+				}
+			}()
+			for req := range reqs {
+				err := req.Reply(false, nil)
+				if err != nil {
+					log.Printf("Failed to reply to request: %v", err)
+					return
+				}
+			}
+		}()
+		go s.HandleForwardedConnection(conn, channel, conn.RemoteAddr())
 	}
 }
 
@@ -369,15 +391,15 @@ func (s *SSHSession) waitForRunningStatus() {
 	for {
 		select {
 		case <-ticker.C:
-			s.interaction.SendMessage(fmt.Sprintf("\rLoading %s", frames[i]))
+			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")
+			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")
+			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)
@@ -425,137 +447,40 @@ func waitForKeyPress(connection ssh.Channel) {
 	}
 }
 
-func (s *SSHSession) HandleForwardedConnection(conn UserConnection, sshConn *ssh.ServerConn) {
-	defer func(Writer net.Conn) {
-		err := Writer.Close()
-		if err != nil {
-			log.Println("Failed to close connection:", err)
+func (s *SSHSession) HandleForwardedConnection(dst io.ReadWriter, src ssh.Channel, remoteAddr net.Addr) {
+	defer func(src ssh.Channel) {
+		err := src.Close()
+		if err != nil && !errors.Is(err, io.EOF) {
+			log.Printf("Error closing connection: %v", err)
 		}
-	}(conn.Writer)
-
-	log.Printf("Handling new forwarded connection from %s", conn.Writer.RemoteAddr())
-	host, originPort := ParseAddr(conn.Writer.RemoteAddr().String())
-
-	timestamp := time.Now().Format("02/Jan/2006 15:04:05")
-
-	payload := createForwardedTCPIPPayload(host, uint16(originPort), s.forwarder.ForwardedPort)
-	channel, reqs, err := sshConn.OpenChannel("forwarded-tcpip", payload)
-	if err != nil {
-		log.Printf("Failed to open forwarded-tcpip channel: %v", err)
-		sendBadGatewayResponse(conn.Writer)
-		return
-	}
-	defer func(channel ssh.Channel) {
-		err := channel.Close()
-		if err != nil {
-			log.Println("Failed to close connection:", err)
-		}
-	}(channel)
+	}(src)
+	log.Printf("Handling new forwarded connection from %s", remoteAddr)
 
 	go func() {
-		defer func() {
-			if r := recover(); r != nil {
-				log.Printf("Panic in request handler: %v", r)
-			}
-		}()
-		for req := range reqs {
-			err := req.Reply(false, nil)
-			if err != nil {
-				log.Printf("Failed to reply to request: %v", err)
-				return
-			}
-		}
-	}()
-
-	if conn.Reader == nil {
-		conn.Reader = bufio.NewReader(conn.Writer)
-	}
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	go func() {
-		defer func() {
-			if r := recover(); r != nil {
-				log.Printf("Panic in reader copy: %v", r)
-			}
-			cancel()
-		}()
-		_, err := io.Copy(channel, conn.Reader)
+		_, err := io.Copy(src, dst)
 		if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
 			log.Printf("Error copying from conn.Reader to channel: %v", err)
 		}
-		cancel()
 	}()
 
-	reader := bufio.NewReader(channel)
-
-	peekChan := make(chan error, 1)
-	go func() {
-		_, err := reader.Peek(1)
-		peekChan <- err
-	}()
-
-	select {
-	case err := <-peekChan:
-		if err == io.EOF {
-			s.interaction.SendMessage(fmt.Sprintf("\033[33m%s -> [%s] WARNING -- \"Could not forward request to the tunnel addr\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.forwarder.TunnelType))
-			sendBadGatewayResponse(conn.Writer)
-			return
-		}
-		if err != nil {
-			log.Printf("Error peeking channel data: %v", err)
-			s.interaction.SendMessage(fmt.Sprintf("\033[33m%s -> [%s] WARNING -- \"Could not forward request to the tunnel addr\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.forwarder.TunnelType))
-			sendBadGatewayResponse(conn.Writer)
-			return
-		}
-	case <-time.After(5 * time.Second):
-		log.Printf("Timeout waiting for channel data from %s", conn.Writer.RemoteAddr())
-		s.interaction.SendMessage(fmt.Sprintf("\033[33m%s -> [%s] WARNING -- \"Could not forward request to the tunnel addr\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.forwarder.TunnelType))
-		sendBadGatewayResponse(conn.Writer)
-		return
-	case <-ctx.Done():
-		return
-	}
-
-	s.interaction.SendMessage(fmt.Sprintf("\033[32m%s -> [%s] TUNNEL ADDRESS -- \"%s\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.forwarder.TunnelType, timestamp))
-
-	_, err = io.Copy(conn.Writer, reader)
+	_, err := io.Copy(dst, src)
 
 	if err != nil && !errors.Is(err, io.EOF) {
 		log.Printf("Error copying from channel to conn.Writer: %v", err)
 	}
+	return
 }
 
-func sendBadGatewayResponse(writer io.Writer) {
-	response := "HTTP/1.1 502 Bad Gateway\r\n" +
-		"Content-Length: 11\r\n" +
-		"Content-Type: text/plain\r\n\r\n" +
-		"Bad Gateway"
-	_, err := io.Copy(writer, bytes.NewReader([]byte(response)))
-	if err != nil {
-		log.Printf("failed to write Bad Gateway response: %v", err)
-		return
+func readSSHString(reader *bytes.Reader) (string, error) {
+	var length uint32
+	if err := binary.Read(reader, binary.BigEndian, &length); err != nil {
+		return "", err
 	}
-}
-
-func writeSSHString(buffer *bytes.Buffer, str string) {
-	err := binary.Write(buffer, binary.BigEndian, uint32(len(str)))
-	if err != nil {
-		log.Printf("Failed to write string to buffer: %v", err)
-		return
+	strBytes := make([]byte, length)
+	if _, err := reader.Read(strBytes); err != nil {
+		return "", err
 	}
-	buffer.WriteString(str)
-}
-
-func ParseAddr(addr string) (string, uint32) {
-	host, portStr, err := net.SplitHostPort(addr)
-	if err != nil {
-		log.Printf("Failed to parse origin address: %s from address %s", err.Error(), addr)
-		return "0.0.0.0", uint32(0)
-	}
-	port, _ := strconv.Atoi(portStr)
-	return host, uint32(port)
+	return string(strBytes), nil
 }
 
 func createForwardedTCPIPPayload(host string, originPort, port uint16) []byte {
@@ -577,14 +502,21 @@ func createForwardedTCPIPPayload(host string, originPort, port uint16) []byte {
 	return buf.Bytes()
 }
 
-func readSSHString(reader *bytes.Reader) (string, error) {
-	var length uint32
-	if err := binary.Read(reader, binary.BigEndian, &length); err != nil {
-		return "", err
+func writeSSHString(buffer *bytes.Buffer, str string) {
+	err := binary.Write(buffer, binary.BigEndian, uint32(len(str)))
+	if err != nil {
+		log.Printf("Failed to write string to buffer: %v", err)
+		return
 	}
-	strBytes := make([]byte, length)
-	if _, err := reader.Read(strBytes); err != nil {
-		return "", err
-	}
-	return string(strBytes), nil
+	buffer.WriteString(str)
+}
+
+func ParseAddr(addr string) (string, uint32) {
+	host, portStr, err := net.SplitHostPort(addr)
+	if err != nil {
+		log.Printf("Failed to parse origin address: %s from address %s", err.Error(), addr)
+		return "0.0.0.0", uint32(0)
+	}
+	port, _ := strconv.Atoi(portStr)
+	return host, uint32(port)
 }

session/interaction.go

@@ -371,7 +371,7 @@ func updateClientSlug(oldSlug, newSlug string) bool {
 	}
 
 	delete(Clients, oldSlug)
-	client.forwarder.setSlug(newSlug)
+	client.Forwarder.setSlug(newSlug)
 	Clients[newSlug] = client
 	return true
 }

session/session.go

@@ -49,7 +49,6 @@ type ForwardingController interface {
 	HandleHTTPForward(req *ssh.Request, port uint16)
 	HandleTCPForward(req *ssh.Request, addr string, port uint16)
 	AcceptTCPConnections()
-	HandleForwardedConnection(conn UserConnection, sshConn *ssh.ServerConn)
 }
 
 type Session interface {
@@ -98,9 +97,9 @@ type Interaction struct {
 	forwarder ForwarderInfo
 }
 type SSHSession struct {
-	lifecycle   *Lifecycle
-	interaction *Interaction
-	forwarder   *Forwarder
+	Lifecycle   *Lifecycle
+	Interaction *Interaction
+	Forwarder   *Forwarder
 
 	Conn    *ssh.ServerConn
 	channel ssh.Channel
@@ -111,10 +110,10 @@ type SSHSession struct {
 
 func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan ssh.NewChannel) {
 	session := SSHSession{
-		lifecycle: &Lifecycle{
+		Lifecycle: &Lifecycle{
 			Status: INITIALIZING,
 		},
-		interaction: &Interaction{
+		Interaction: &Interaction{
 			CommandBuffer: new(bytes.Buffer),
 			EditMode:      false,
 			EditSlug:      "",
@@ -124,7 +123,7 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
 			session:       nil,
 			forwarder:     nil,
 		},
-		forwarder: &Forwarder{
+		Forwarder: &Forwarder{
 			Listener:      nil,
 			TunnelType:    "",
 			ForwardedPort: 0,
@@ -136,12 +135,12 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
 		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
+	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()
@@ -150,8 +149,8 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
 			ch, reqs, _ := channel.Accept()
 			if session.channel == nil {
 				session.channel = ch
-				session.interaction.channel = ch
-				session.lifecycle.Status = SETUP
+				session.Interaction.channel = ch
+				session.Lifecycle.Status = SETUP
 				go session.HandleGlobalRequest(forwardingReq)
 			}
 			go session.HandleGlobalRequest(reqs)