Add SSH keepalive and auto-reconnection to prevent idle disconnections

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

cmd/server/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bufio"
 	"fmt"
 	"io"
 	"log"
@@ -91,9 +92,13 @@ func startHttpProxy(port string) {
 		}
 		slug := parts[0]
 
+		log.Printf("Request Host: %s, Extracted Slug: %s", host, slug)
+
 		tunnel, ok := manager.Get(slug)
 		if !ok {
+			log.Printf("Tunnel not found for slug: %s", slug)
 			// Serve 404 page
+			w.WriteHeader(http.StatusNotFound)
 			http.ServeFile(w, r, "./cmd/server/static/404.html")
 			return
 		}
@@ -125,21 +130,15 @@ func startHttpProxy(port string) {
 		// But typically `gliderlabs/ssh` is for allowing the server to be a jump host.
 		// We want to be an HTTP Gateway.
 
-		// 1. Hijack the connection to handle bidirectional traffic (WebSockets)
-		hijacker, ok := w.(http.Hijacker)
-		if !ok {
-			http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
-			return
+		// Panic recovery
+		defer func() {
+			if r := recover(); r != nil {
+				log.Printf("Recovered from panic in handler: %v", r)
+				http.Error(w, "Internal Server Error", http.StatusInternalServerError)
 			}
-		clientConn, bufrw, err := hijacker.Hijack()
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusServiceUnavailable)
-			return
-		}
-		defer clientConn.Close()
+		}()
 
 		// 2. Open channel to client
-		// "forwarded-tcpip" arguments:
 		destHost := "0.0.0.0"
 		destPort := tunnel.LocalPort
 		srcHost := "127.0.0.1"
@@ -161,36 +160,103 @@ func startHttpProxy(port string) {
 		defer ch.Close()
 		go gossh.DiscardRequests(reqs)
 
+		// Check if it is a WebSocket Upgrade
+		isWebSocket := false
+		if strings.ToLower(r.Header.Get("Upgrade")) == "websocket" {
+			isWebSocket = true
+		}
+
+		if isWebSocket {
+			// WEBSOCKET STRATEGY: Hijack and bidirectional copy
+			hijacker, ok := w.(http.Hijacker)
+			if !ok {
+				http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
+				return
+			}
+			clientConn, bufrw, err := hijacker.Hijack()
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusServiceUnavailable)
+				return
+			}
+			defer clientConn.Close()
+
+			// Manual Request writing to avoid touching Body after Hijack/Panic
+			// Request Line
+			reqLine := fmt.Sprintf("%s %s %s\r\n", r.Method, r.RequestURI, r.Proto)
+			if _, err := io.WriteString(ch, reqLine); err != nil {
+				log.Printf("Error writing websocket request line: %v", err)
+				return
+			}
+			// Headers
+			if err := r.Header.Write(ch); err != nil {
+				log.Printf("Error writing websocket headers: %v", err)
+				return
+			}
+			// End of headers
+			if _, err := io.WriteString(ch, "\r\n"); err != nil {
+				log.Printf("Error writing websocket header terminator: %v", err)
+				return
+			}
+
 			var wg sync.WaitGroup
 			wg.Add(2)
 
-		// 3. Browser -> Backend (Write request + Copy raw stream)
+			// Copy existing buffer from hijack + future reads -> backend
 			go func() {
-			defer wg.Done()
-			// Write the initial request (Method, Path, Headers)
-			// This sets up the handshake or request.
-			// Note: We use r.Write to reconstruct the request line and headers.
-			// For WebSockets, the Body is empty, so this writes headers and returns.
-			// For POSTs, it writes headers and tries to copy Body.
-			if err := r.Write(ch); err != nil {
-				log.Printf("Error writing request to backend: %v", err)
-				return
+				defer func() {
+					if r := recover(); r != nil {
+						log.Printf("Recovered from panic in WS writer: %v", r)
 					}
-			// Important: Continue copying any subsequent data (like WebSocket frames)
-			// from the hijacked buffer/connection to the channel.
-			io.Copy(ch, bufrw)
-			// e.g. when browser closes or stops sending, we are done here.
+					wg.Done()
+				}()
+				if bufrw.Reader.Buffered() > 0 {
+					io.CopyN(ch, bufrw, int64(bufrw.Reader.Buffered()))
+				}
+				io.Copy(ch, clientConn)
 			}()
 
-		// 4. Backend -> Browser (Copy raw stream)
+			// Backend -> Browser
 			go func() {
-			defer wg.Done()
+				defer func() {
+					if r := recover(); r != nil {
+						log.Printf("Recovered from panic in WS reader: %v", r)
+					}
+					wg.Done()
+				}()
 				io.Copy(clientConn, ch)
-			// When backend closes connection, close browser connection
-			clientConn.Close()
 			}()
 
 			wg.Wait()
+			return
+		}
+
+		// STANDARD HTTP STRATEGY: Request/Response Tunneling
+		// 1. Write Request to Channel (simulating wire)
+		if err := r.Write(ch); err != nil {
+			log.Printf("Error writing request to backend: %v", err)
+			http.Error(w, "Gateway Error", http.StatusBadGateway)
+			return
+		}
+
+		// 2. Read Response from Channel (parsing wire)
+		resp, err := http.ReadResponse(bufio.NewReader(ch), r)
+		if err != nil {
+			log.Printf("Error reading response from backend: %v", err)
+			http.Error(w, "Bad Gateway", http.StatusBadGateway)
+			return
+		}
+		defer resp.Body.Close()
+
+		// 3. Copy Headers to `w`
+		for k, vv := range resp.Header {
+			for _, v := range vv {
+				w.Header().Add(k, v)
+			}
+		}
+		w.WriteHeader(resp.StatusCode)
+
+		// 4. Copy Body to `w`
+		io.Copy(w, resp.Body)
 
 	})
 

internal/tunnel/client.go

@@ -31,6 +31,7 @@ type Client struct {
 	Events     chan string // Channel to send logs/events to TUI
 	Metrics    chan int64  // Channel to send bytes transferred
 	PublicURL  string      // PublicURL is the URL accessible from the internet
+	stopKeepAlive chan struct{} // Signal to stop keepalive goroutine
 }
 
 func NewClient(serverAddr, localPort, authToken, hostHeader string, localHTTPS bool) *Client {
@@ -45,7 +46,7 @@ func NewClient(serverAddr, localPort, authToken, hostHeader string, localHTTPS b
 	}
 }
 
-func (c *Client) Start() error {
+func (c *Client) connect() error {
 	config := &ssh.ClientConfig{
 		User: "grokway",
 		Auth: []ssh.AuthMethod{
@@ -93,18 +94,86 @@ func (c *Client) Start() error {
 	c.PublicURL = fmt.Sprintf("https://%s.%s", slug, host)
 	c.Events <- fmt.Sprintf("Tunnel established! Public URL: %s", c.PublicURL)
 
-	go c.acceptLoop()
+	// Start SSH keepalive to prevent idle disconnection
+	c.stopKeepAlive = make(chan struct{})
+	go c.keepAlive()
 
 	return nil
 }
 
+func (c *Client) Start() error {
+	if err := c.connect(); err != nil {
+		return err
+	}
+	go c.acceptLoop()
+	return nil
+}
+
+// keepAlive sends periodic SSH keepalive requests to prevent
+// firewalls/NATs from dropping idle connections.
+func (c *Client) keepAlive() {
+	ticker := time.NewTicker(30 * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-c.stopKeepAlive:
+			return
+		case <-ticker.C:
+			if c.SSHClient != nil {
+				_, _, err := c.SSHClient.SendRequest("keepalive@openssh.com", true, nil)
+				if err != nil {
+					logToFile(fmt.Sprintf("Keepalive failed: %v", err))
+					return
+				}
+			}
+		}
+	}
+}
+
+func (c *Client) closeConnection() {
+	if c.stopKeepAlive != nil {
+		select {
+		case <-c.stopKeepAlive:
+			// Already closed
+		default:
+			close(c.stopKeepAlive)
+		}
+	}
+	if c.Listener != nil {
+		c.Listener.Close()
+	}
+	if c.SSHClient != nil {
+		c.SSHClient.Close()
+	}
+}
+
 func (c *Client) acceptLoop() {
 	for {
 		remoteConn, err := c.Listener.Accept()
 		if err != nil {
-			c.Events <- fmt.Sprintf("Accept error: %s", err)
+			c.Events <- fmt.Sprintf("Connection lost: %s", err)
+			c.closeConnection()
+
+			// Reconnect with backoff
+			delay := 2 * time.Second
+			maxDelay := 60 * time.Second
+			for {
+				c.Events <- fmt.Sprintf("Reconnecting in %s...", delay)
+				time.Sleep(delay)
+
+				if err := c.connect(); err != nil {
+					c.Events <- fmt.Sprintf("Reconnection failed: %s", err)
+					delay *= 2
+					if delay > maxDelay {
+						delay = maxDelay
+					}
+					continue
+				}
+				c.Events <- "Reconnected successfully!"
 				break
 			}
+			continue
+		}
 
 		c.Events <- "New Request received"
 		go c.handleConnection(remoteConn)