From 8442fedef19cd44890d6c4131281a11944b9638f Mon Sep 17 00:00:00 2001 From: bagas Date: Fri, 28 Nov 2025 16:56:03 +0700 Subject: [PATCH] refactor: handle error --- server/http.go | 54 ++++++++-- server/https.go | 24 ++++- session/handler.go | 223 +++++++++++++++++++++++++++++++++++------ session/interaction.go | 132 ++++++++++++++++++++---- 4 files changed, 367 insertions(+), 66 deletions(-) diff --git a/server/http.go b/server/http.go index c632641..49e23dd 100644 --- a/server/http.go +++ b/server/http.go @@ -38,9 +38,21 @@ func (w *connResponseWriter) WriteHeader(statusCode int) { return } w.wrote = true - fmt.Fprintf(w.conn, "HTTP/1.1 %d %s\r\n", statusCode, http.StatusText(statusCode)) - w.header.Write(w.conn) - fmt.Fprint(w.conn, "\r\n") + _, err := fmt.Fprintf(w.conn, "HTTP/1.1 %d %s\r\n", statusCode, http.StatusText(statusCode)) + if err != nil { + log.Printf("Error writing HTTP response: %v", err) + return + } + err = w.header.Write(w.conn) + if err != nil { + log.Printf("Error writing HTTP header: %v", err) + return + } + _, err = fmt.Fprint(w.conn, "\r\n") + if err != nil { + log.Printf("Error writing HTTP header: %v", err) + return + } } func (w *connResponseWriter) Write(b []byte) (int, error) { @@ -120,20 +132,36 @@ func Handler(conn net.Conn) { host := strings.Split(parseHostFromHeader(headers), ".") if len(host) < 1 { - conn.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n")) - conn.Close() + _, err := conn.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n")) + if err != nil { + log.Println("Failed to write 400 Bad Request:", err) + return + } + err = conn.Close() + if err != nil { + log.Println("Failed to close connection:", err) + return + } return } slug := host[0] if redirectTLS { - conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" + + _, err := conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" + fmt.Sprintf("Location: https://%s.%s/\r\n", slug, utils.Getenv("domain")) + "Content-Length: 0\r\n" + "Connection: close\r\n" + "\r\n")) - conn.Close() + if err != nil { + log.Println("Failed to write 301 Moved Permanently:", err) + return + } + err = conn.Close() + if err != nil { + log.Println("Failed to close connection:", err) + return + } return } @@ -173,12 +201,20 @@ func Handler(conn net.Conn) { sshSession, ok := session.Clients[slug] if !ok { - conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" + + _, err := conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" + fmt.Sprintf("Location: https://tunnl.live/tunnel-not-found?slug=%s\r\n", slug) + "Content-Length: 0\r\n" + "Connection: close\r\n" + "\r\n")) - conn.Close() + if err != nil { + log.Println("Failed to write 301 Moved Permanently:", err) + return + } + err = conn.Close() + if err != nil { + log.Println("Failed to close connection:", err) + return + } return } diff --git a/server/https.go b/server/https.go index a86eb84..649287d 100644 --- a/server/https.go +++ b/server/https.go @@ -54,8 +54,16 @@ func HandlerTLS(conn net.Conn) { host := strings.Split(parseHostFromHeader(headers), ".") if len(host) < 1 { - conn.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n")) - conn.Close() + _, err := conn.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n")) + if err != nil { + log.Println("Failed to write 400 Bad Request:", err) + return + } + err = conn.Close() + if err != nil { + log.Println("Failed to close connection:", err) + return + } return } @@ -97,12 +105,20 @@ func HandlerTLS(conn net.Conn) { sshSession, ok := session.Clients[slug] if !ok { - conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" + + _, err := conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" + fmt.Sprintf("Location: https://tunnl.live/tunnel-not-found?slug=%s\r\n", slug) + "Content-Length: 0\r\n" + "Connection: close\r\n" + "\r\n")) - conn.Close() + if err != nil { + log.Println("Failed to write 301 Moved Permanently:", err) + return + } + err = conn.Close() + if err != nil { + log.Println("Failed to close connection:", err) + return + } return } diff --git a/session/handler.go b/session/handler.go index 34f5069..b1c4545 100644 --- a/session/handler.go +++ b/session/handler.go @@ -11,6 +11,7 @@ import ( "log" "net" "strconv" + "strings" "sync" "time" portUtil "tunnel_pls/internal/port" @@ -36,6 +37,71 @@ var ( Clients = make(map[string]*SSHSession) ) +type HeaderModifier struct { + r io.Reader + headerBuf []byte + headerDone bool + state int +} + +func (hm *HeaderModifier) Read(p []byte) (int, error) { + n, err := hm.r.Read(p) + if n > 0 && !hm.headerDone { + for i := 0; i < n; i++ { + b := p[i] + hm.headerBuf = append(hm.headerBuf, b) + + switch hm.state { + case 0: + if b == '\r' { + hm.state = 1 + } + case 1: + if b == '\n' { + hm.state = 2 + } else { + hm.state = 0 + } + case 2: + if b == '\r' { + hm.state = 3 + } else { + hm.state = 0 + } + case 3: + if b == '\n' { + hm.headerDone = true + modifiedHeader := hm.modifyHeader(hm.headerBuf) + copy(p, modifiedHeader) + return len(modifiedHeader), nil + } else { + hm.state = 0 + } + } + } + } + + return n, err +} +func (hm *HeaderModifier) modifyHeader(header []byte) []byte { + lines := strings.Split(string(header), "\r\n") + found := false + + for i, line := range lines { + if strings.HasPrefix(strings.ToLower(line), "server:") { + lines[i] = "Server: tunnel_please" + found = true + } + } + + if !found { + lines = append(lines[:len(lines)-2], "Server: tunnel_please", "", "") + } + + modified := strings.Join(lines, "\r\n") + return []byte(modified) +} + func registerClient(slug string, session *SSHSession) bool { clientsMutex.Lock() defer clientsMutex.Unlock() @@ -99,10 +165,18 @@ func (s *SSHSession) HandleGlobalRequest(GlobalRequest <-chan *ssh.Request) { s.handleTCPIPForward(req) return case "shell", "pty-req", "window-change": - req.Reply(true, nil) + err := req.Reply(true, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } default: log.Println("Unknown request type:", req.Type) - req.Reply(false, nil) + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } } } } @@ -115,8 +189,12 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) { addr, err := readSSHString(reader) if err != nil { log.Println("Failed to read address from payload:", err) - req.Reply(false, nil) - err := s.Close() + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } + err = s.Close() if err != nil { log.Printf("failed to close session: %v", err) } @@ -127,8 +205,12 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) { 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)) - req.Reply(false, nil) - err := s.Close() + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } + err = s.Close() if err != nil { log.Printf("failed to close session: %v", err) } @@ -137,8 +219,12 @@ 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)) - req.Reply(false, nil) - err := s.Close() + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } + err = s.Close() if err != nil { log.Printf("failed to close session: %v", err) } @@ -149,8 +235,12 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) { 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)) - req.Reply(false, nil) - err := s.Close() + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } + err = s.Close() if err != nil { log.Printf("failed to close session: %v", err) } @@ -170,8 +260,12 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) { portToBind = unassign if !success { s.interaction.SendMessage(fmt.Sprintf("No available port\r\n", portToBind)) - req.Reply(false, nil) - err := s.Close() + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } + err = s.Close() if err != nil { log.Printf("failed to close session: %v", err) } @@ -179,14 +273,22 @@ func (s *SSHSession) handleTCPIPForward(req *ssh.Request) { } } 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)) - req.Reply(false, nil) - err := s.Close() + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } + err = s.Close() if err != nil { log.Printf("failed to close session: %v", err) } return } - portUtil.Manager.SetPortStatus(portToBind, true) + err := portUtil.Manager.SetPortStatus(portToBind, true) + if err != nil { + log.Println("Failed to set port status:", err) + return + } } s.handleTCPForward(req, addr, portToBind) } @@ -214,7 +316,11 @@ func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) { slug := generateUniqueSlug() if slug == "" { - req.Reply(false, nil) + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } return } @@ -222,7 +328,11 @@ func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) { registerClient(slug, s) buf := new(bytes.Buffer) - binary.Write(buf, binary.BigEndian, uint32(80)) + err := binary.Write(buf, binary.BigEndian, uint32(80)) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } log.Printf("HTTP forwarding approved on port: %d", 80) domain := utils.Getenv("domain") @@ -233,7 +343,11 @@ func (s *SSHSession) handleHTTPForward(req *ssh.Request, portToBind uint16) { s.interaction.ShowWelcomeMessage() s.interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to %s://%s.%s\r\n", protocol, slug, domain)) - req.Reply(true, buf.Bytes()) + err = req.Reply(true, buf.Bytes()) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } } func (s *SSHSession) handleTCPForward(req *ssh.Request, addr string, portToBind uint16) { @@ -243,8 +357,12 @@ func (s *SSHSession) handleTCPForward(req *ssh.Request, addr string, 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)) - req.Reply(false, nil) - err := s.Close() + err := req.Reply(false, nil) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } + err = s.Close() if err != nil { log.Printf("failed to close session: %v", err) } @@ -258,9 +376,17 @@ func (s *SSHSession) handleTCPForward(req *ssh.Request, addr string, portToBind go s.acceptTCPConnections() buf := new(bytes.Buffer) - binary.Write(buf, binary.BigEndian, uint32(portToBind)) + err = binary.Write(buf, binary.BigEndian, uint32(portToBind)) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } log.Printf("TCP forwarding approved on port: %d", portToBind) - req.Reply(true, buf.Bytes()) + err = req.Reply(true, buf.Bytes()) + if err != nil { + log.Println("Failed to reply to request:", err) + return + } } func (s *SSHSession) acceptTCPConnections() { @@ -366,7 +492,12 @@ func waitForKeyPress(connection ssh.Channel) { } func (s *SSHSession) HandleForwardedConnection(conn UserConnection, sshConn *ssh.ServerConn) { - defer conn.Writer.Close() + defer func(Writer net.Conn) { + err := Writer.Close() + if err != nil { + log.Println("Failed to close connection:", err) + } + }(conn.Writer) log.Printf("Handling new forwarded connection from %s", conn.Writer.RemoteAddr()) host, originPort := ParseAddr(conn.Writer.RemoteAddr().String()) @@ -380,7 +511,12 @@ func (s *SSHSession) HandleForwardedConnection(conn UserConnection, sshConn *ssh sendBadGatewayResponse(conn.Writer) return } - defer channel.Close() + defer func(channel ssh.Channel) { + err := channel.Close() + if err != nil { + log.Println("Failed to close connection:", err) + } + }(channel) go func() { defer func() { @@ -389,7 +525,11 @@ func (s *SSHSession) HandleForwardedConnection(conn UserConnection, sshConn *ssh } }() for req := range reqs { - req.Reply(false, nil) + err := req.Reply(false, nil) + if err != nil { + log.Printf("Failed to reply to request: %v", err) + return + } } }() @@ -407,7 +547,6 @@ func (s *SSHSession) HandleForwardedConnection(conn UserConnection, sshConn *ssh } cancel() }() - _, err := io.Copy(channel, conn.Reader) if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) { log.Printf("Error copying from conn.Reader to channel: %v", err) @@ -447,7 +586,13 @@ func (s *SSHSession) HandleForwardedConnection(conn UserConnection, sshConn *ssh 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) + if s.forwarder.GetTunnelType() == HTTP { + ir := &HeaderModifier{r: reader} + _, err = io.Copy(conn.Writer, ir) + } else { + _, err = io.Copy(conn.Writer, reader) + } + if err != nil && !errors.Is(err, io.EOF) { log.Printf("Error copying from channel to conn.Writer: %v", err) } @@ -458,11 +603,19 @@ func sendBadGatewayResponse(writer io.Writer) { "Content-Length: 11\r\n" + "Content-Type: text/plain\r\n\r\n" + "Bad Gateway" - io.Copy(writer, bytes.NewReader([]byte(response))) + _, err := io.Copy(writer, bytes.NewReader([]byte(response))) + if err != nil { + log.Printf("failed to write Bad Gateway response: %v", err) + return + } } func writeSSHString(buffer *bytes.Buffer, str string) { - binary.Write(buffer, binary.BigEndian, uint32(len(str))) + err := binary.Write(buffer, binary.BigEndian, uint32(len(str))) + if err != nil { + log.Printf("Failed to write string to buffer: %v", err) + return + } buffer.WriteString(str) } @@ -480,9 +633,17 @@ func createForwardedTCPIPPayload(host string, originPort, port uint16) []byte { var buf bytes.Buffer writeSSHString(&buf, "localhost") - binary.Write(&buf, binary.BigEndian, uint32(port)) + err := binary.Write(&buf, binary.BigEndian, uint32(port)) + if err != nil { + log.Printf("Failed to write string to buffer: %v", err) + return nil + } writeSSHString(&buf, host) - binary.Write(&buf, binary.BigEndian, uint32(originPort)) + err = binary.Write(&buf, binary.BigEndian, uint32(originPort)) + if err != nil { + log.Printf("Failed to write string to buffer: %v", err) + return nil + } return buf.Bytes() } diff --git a/session/interaction.go b/session/interaction.go index 07d0ead..3ef056a 100644 --- a/session/interaction.go +++ b/session/interaction.go @@ -79,14 +79,30 @@ func (i *Interaction) HandleSlugEditMode(connection ssh.Channel, char byte, comm } else if char == 8 || char == 127 { if len(i.EditSlug) > 0 { i.EditSlug = (i.EditSlug)[:len(i.EditSlug)-1] - connection.Write([]byte("\r\033[K")) - connection.Write([]byte("➤ " + i.EditSlug + "." + utils.Getenv("domain"))) + _, err := connection.Write([]byte("\r\033[K")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("➤ " + i.EditSlug + "." + utils.Getenv("domain"))) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } } } else if char >= 32 && char <= 126 { if (char >= 'a' && char <= 'z') || (char >= '0' && char <= '9') || char == '-' { i.EditSlug += string(char) - connection.Write([]byte("\r\033[K")) - connection.Write([]byte("➤ " + i.EditSlug + "." + utils.Getenv("domain"))) + _, err := connection.Write([]byte("\r\033[K")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("➤ " + i.EditSlug + "." + utils.Getenv("domain"))) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } } } } @@ -94,7 +110,11 @@ func (i *Interaction) HandleSlugEditMode(connection ssh.Channel, char byte, comm func (i *Interaction) HandleSlugSave(connection ssh.Channel) { isValid := isValidSlug(i.EditSlug) - connection.Write([]byte("\033[H\033[2J")) + _, err := connection.Write([]byte("\033[H\033[2J")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } if isValid { oldSlug := i.getSlug() newSlug := i.EditSlug @@ -104,24 +124,72 @@ func (i *Interaction) HandleSlugSave(connection ssh.Channel) { return } - connection.Write([]byte("\r\n\r\n✅ SUBDOMAIN UPDATED ✅\r\n\r\n")) - connection.Write([]byte("Your new address is: " + newSlug + "." + utils.Getenv("domain") + "\r\n\r\n")) - connection.Write([]byte("Press any key to continue...\r\n")) + _, err := connection.Write([]byte("\r\n\r\n✅ SUBDOMAIN UPDATED ✅\r\n\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("Your new address is: " + newSlug + "." + utils.Getenv("domain") + "\r\n\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("Press any key to continue...\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } } else if isForbiddenSlug(i.EditSlug) { - connection.Write([]byte("\r\n\r\n❌ FORBIDDEN SUBDOMAIN ❌\r\n\r\n")) - connection.Write([]byte("This subdomain is not allowed.\r\n")) - connection.Write([]byte("Please try a different subdomain.\r\n\r\n")) - connection.Write([]byte("Press any key to continue...\r\n")) + _, err := connection.Write([]byte("\r\n\r\n❌ FORBIDDEN SUBDOMAIN ❌\r\n\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("This subdomain is not allowed.\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("Please try a different subdomain.\r\n\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("Press any key to continue...\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } } else { - connection.Write([]byte("\r\n\r\n❌ INVALID SUBDOMAIN ❌\r\n\r\n")) - connection.Write([]byte("Use only lowercase letters, numbers, and hyphens.\r\n")) - connection.Write([]byte("Length must be 3-20 characters and cannot start or end with a hyphen.\r\n\r\n")) - connection.Write([]byte("Press any key to continue...\r\n")) + _, err := connection.Write([]byte("\r\n\r\n❌ INVALID SUBDOMAIN ❌\r\n\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("Use only lowercase letters, numbers, and hyphens.\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("Length must be 3-20 characters and cannot start or end with a hyphen.\r\n\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("Press any key to continue...\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } } waitForKeyPress(connection) - connection.Write([]byte("\033[H\033[2J")) + _, err = connection.Write([]byte("\033[H\033[2J")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } i.ShowWelcomeMessage() domain := utils.Getenv("domain") @@ -129,7 +197,11 @@ func (i *Interaction) HandleSlugSave(connection ssh.Channel) { if utils.Getenv("tls_enabled") == "true" { protocol = "https" } - connection.Write([]byte(fmt.Sprintf("Forwarding your traffic to %s://%s.%s \r\n", protocol, i.getSlug(), domain))) + _, err = connection.Write([]byte(fmt.Sprintf("Forwarding your traffic to %s://%s.%s \r\n", protocol, i.getSlug(), domain))) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } i.EditMode = false i.CommandBuffer.Reset() @@ -137,13 +209,29 @@ func (i *Interaction) HandleSlugSave(connection ssh.Channel) { func (i *Interaction) HandleSlugCancel(connection ssh.Channel, commandBuffer *bytes.Buffer) { i.EditMode = false - connection.Write([]byte("\033[H\033[2J")) - connection.Write([]byte("\r\n\r\n⚠️ SUBDOMAIN EDIT CANCELLED ⚠️\r\n\r\n")) - connection.Write([]byte("Press any key to continue...\r\n")) + _, err := connection.Write([]byte("\033[H\033[2J")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("\r\n\r\n⚠️ SUBDOMAIN EDIT CANCELLED ⚠️\r\n\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } + _, err = connection.Write([]byte("Press any key to continue...\r\n")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } waitForKeyPress(connection) - connection.Write([]byte("\033[H\033[2J")) + _, err = connection.Write([]byte("\033[H\033[2J")) + if err != nil { + log.Printf("failed to write to channel: %v", err) + return + } i.ShowWelcomeMessage() commandBuffer.Reset()