JSON Unmarshalling of the Initial Authentication Message, Security and Input Validation, timeout

Authentication:
The module now reads a single JSON message containing both secret and host_port values.

Timeouts:
The Telnet dial uses a 10‑second timeout, and the WebSocket connection is closed after 120 seconds of inactivity. Adjust these values as needed.

Security:
The upgrader currently allows all origins. In a production environment, restrict this by verifying r.Origin.

Error Handling & Goroutine Cleanup:
Each goroutine writes to a buffered error channel so that the first error will cancel the connection. In a more complex scenario, you might use a context with cancellation or a sync.WaitGroup to better manage goroutine lifetimes.

telnetproxy.go

@@ -2,14 +2,12 @@ package telnetproxy
 
 import (
 	"bufio"
-    "context"
 	"errors"
 	"fmt"
 	"io"
 	"log"
 	"net"
 	"net/http"
-    "strings"
 	"time"
 
 	"github.com/caddyserver/caddy/v2"
@@ -18,7 +16,7 @@ import (
 	"github.com/gorilla/websocket"
 )
 
-// TelnetProxy is a Caddy module that proxies websocket connections to telnet servers.
+// TelnetProxy is a Caddy module that proxies WebSocket connections to Telnet servers.
 type TelnetProxy struct {
 	Secret string `json:"secret,omitempty"`
 }
@@ -41,11 +39,17 @@ func (tp *TelnetProxy) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 	return nil
 }
 
+// authMessage is used for the initial authentication and connection details.
+type authMessage struct {
+	Secret   string `json:"secret"`
+	HostPort string `json:"host_port"`
+}
+
 // ServeHTTP implements caddyhttp.MiddlewareHandler.
 func (tp TelnetProxy) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
 	upgrader := websocket.Upgrader{
+		// In production, restrict allowed origins to trusted hosts.
 		CheckOrigin: func(r *http.Request) bool {
-            // Allow all origins (you might want to restrict this in production)
 			return true
 		},
 	}
@@ -55,109 +59,84 @@ func (tp TelnetProxy) ServeHTTP(w http.ResponseWriter, r *http.Request, next cad
 		log.Printf("WebSocket upgrade error: %v", err)
 		return err
 	}
+	// The connection will be closed in handleWebSocket.
 	defer ws.Close()
 
-    // Handle WebSocket connection in a goroutine
+	// Handle the WebSocket connection in a separate goroutine.
-    tp.handleWebSocket(ws)
+	go tp.handleWebSocket(ws)
 
-    return nil // Don't let caddy return a 404 after us
+	// Returning nil prevents further Caddy middleware from interfering.
+	return nil
 }
 
 func (tp TelnetProxy) handleWebSocket(ws *websocket.Conn) {
+	// Recover from panics to avoid crashing the server.
 	defer func() {
 		if r := recover(); r != nil {
 			log.Printf("Recovered from panic: %v", r)
 		}
 	}()
 
-    var (
+	// Read and decode the initial authentication message.
-        hostPort string
+	var msg authMessage
-        clientSecret string
+	if err := ws.ReadJSON(&msg); err != nil {
-        telnetConn net.Conn
-        err error
-    )
-
-    // Initial Authentication and Host/Port
-    err = ws.ReadJSON(map[string]string{
-        "secret":     "",
-        "host_port":  "",
-    })
-    if err != nil {
 		log.Printf("Error reading initial JSON: %v", err)
 		ws.WriteMessage(websocket.TextMessage, []byte("Error reading initial JSON"))
 		return
 	}
 
-
+	// Check authentication.
-    err = ws.ReadJSON(&map[string]string{
+	if msg.Secret != tp.Secret {
-        "secret": &clientSecret,
+		log.Printf("Authentication failed. Received secret: %s", msg.Secret)
-        "host_port": &hostPort,
-    })
-    if err != nil {
-        log.Printf("Error reading initial JSON: %v", err)
-        ws.WriteMessage(websocket.TextMessage, []byte("Error reading initial JSON"))
-        return
-    }
-
-
-
-    // Authentication Check
-    if clientSecret != tp.Secret {
-        log.Printf("Authentication failed. Client secret: %s, server secret: %s", clientSecret, tp.Secret)
 		ws.WriteMessage(websocket.TextMessage, []byte("Authentication failed"))
 		return
 	}
 
-
+	// Connect to the Telnet server with a timeout.
-    // Telnet connection
+	telnetConn, err := net.DialTimeout("tcp", msg.HostPort, 10*time.Second)
-    telnetConn, err = net.Dial("tcp", hostPort)
 	if err != nil {
-        log.Printf("Error connecting to Telnet server: %v", err)
+		log.Printf("Error connecting to Telnet server at %s: %v", msg.HostPort, err)
 		ws.WriteMessage(websocket.TextMessage, []byte("Error connecting to Telnet server: "+err.Error()))
 		return
 	}
 	defer telnetConn.Close()
 
-    // Channel for errors
+	// Create a channel to signal errors from either direction.
-    errChan := make(chan error, 2) // Buffered channel to avoid blocking
+	errChan := make(chan error, 2)
 
-    // Copy from Telnet -> WebSocket (Read from telnet, send to client)
+	// Goroutine: Telnet -> WebSocket.
 	go func() {
-        // Using a buffered reader in case the server sends large chunks of text
 		reader := bufio.NewReader(telnetConn)
 		for {
-            data, err := reader.ReadBytes('\n') // Read until a newline
+			// Read until newline; adjust if the Telnet protocol doesn't send newline-terminated data.
+			data, err := reader.ReadBytes('\n')
 			if err != nil {
 				if !errors.Is(err, io.EOF) {
 					log.Printf("Telnet read error: %v", err)
 					errChan <- fmt.Errorf("Telnet read error: %v", err)
 				}
-                 return // End goroutine on EOF or error
-            }
-
-            if err = ws.WriteMessage(websocket.BinaryMessage, data); err != nil {
-                 log.Printf("Websocket write error (from telnet): %v", err)
-                 errChan <- fmt.Errorf("Websocket write error (from telnet): %v", err)
 				return
 			}
-
+			if err := ws.WriteMessage(websocket.BinaryMessage, data); err != nil {
+				log.Printf("WebSocket write error (from Telnet): %v", err)
+				errChan <- fmt.Errorf("WebSocket write error (from Telnet): %v", err)
+				return
+			}
 		}
 	}()
 
-    // Copy from WebSocket -> Telnet (Read from client, send to server)
+	// Goroutine: WebSocket -> Telnet.
 	go func() {
 		for {
-            _, msg, err := ws.ReadMessage()
+			_, data, err := ws.ReadMessage()
 			if err != nil {
 				if !websocket.IsCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
-                    log.Printf("Websocket read error (from browser): %v", err)
+					log.Printf("WebSocket read error (from client): %v", err)
-                     errChan <- fmt.Errorf("Websocket read error (from browser): %v", err)
+					errChan <- fmt.Errorf("WebSocket read error (from client): %v", err)
 				}
 				return
 			}
-
+			if _, err := telnetConn.Write(data); err != nil {
-            _, err = telnetConn.Write(msg)
-            if err != nil {
 				log.Printf("Telnet write error: %v", err)
 				errChan <- fmt.Errorf("Telnet write error: %v", err)
 				return
@@ -165,16 +144,15 @@ func (tp TelnetProxy) handleWebSocket(ws *websocket.Conn) {
 		}
 	}()
 
-    // Keep connection alive by waiting for an error on either read or write channel
+	// Wait until an error occurs or a timeout happens.
 	select {
-    case <-errChan: // First Error closes connections
+	case <-errChan:
-        log.Println("Closing connection")
+		log.Println("Closing connection due to an error.")
-        telnetConn.Close()
+	case <-time.After(120 * time.Second):
-        ws.Close()
+		log.Println("Closing connection due to inactivity.")
+	}
 
-    case <-time.After(120 * time.Second): // Optional timeout.
+	// Close both connections (deferred calls will handle this if not already closed).
-        log.Println("Closing connection due to inactivity")
 	telnetConn.Close()
 	ws.Close()
 }
-}