Enhance tunnel ID extraction and routing logic in server

- Implement fallback mechanisms for tunnel ID retrieval from cookies and Referer headers.
- Update request handling to set cookies for path-based routing.
- Introduce a new function to extract tunnel ID from Referer URLs.
- Ensure proper error handling for missing or invalid tunnel IDs.

Author: SonnyTaylor

internal/server/server.go

@@ -210,26 +210,66 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) {
 // handleRequest handles incoming HTTP requests and routes them to the appropriate tunnel.
 func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 	var tunnelID string
+	var fromPath bool
 
 	if s.config.RoutingMode == protocol.RoutingModePath {
 		// Extract tunnel ID from the first path segment
 		tunnelID = protocol.ExtractTunnelIDFromPath(r.URL.Path)
-		if tunnelID == "" {
-			http.Error(w, "Missing tunnel ID in path", http.StatusNotFound)
-			return
+		if tunnelID != "" && s.registry.Exists(tunnelID) {
+			fromPath = true
+		}
+
+		// If no valid tunnel ID in path, try cookie
+		if !fromPath {
+			if cookie, err := r.Cookie("exio_tunnel"); err == nil && cookie.Value != "" {
+				if s.registry.Exists(cookie.Value) {
+					tunnelID = cookie.Value
+					s.logger.Printf("Cookie routing: %s (tunnel: %s)", r.URL.Path, tunnelID)
+				}
+			}
 		}
 
-		// Rewrite the path to strip the tunnel ID prefix
-		originalPath := r.URL.Path
-		r.URL.Path = protocol.StripTunnelIDPrefix(r.URL.Path, tunnelID)
-		r.RequestURI = r.URL.RequestURI()
+		// If still no tunnel, try Referer header as fallback
+		if tunnelID == "" || !s.registry.Exists(tunnelID) {
+			referer := r.Header.Get("Referer")
+			if referer != "" {
+				refererTunnelID := protocol.ExtractTunnelIDFromReferer(referer)
+				if refererTunnelID != "" && s.registry.Exists(refererTunnelID) {
+					tunnelID = refererTunnelID
+					s.logger.Printf("Referer routing: %s (tunnel: %s)", r.URL.Path, tunnelID)
+				}
+			}
+		}
 
-		// Also update RawPath if set
-		if r.URL.RawPath != "" {
-			r.URL.RawPath = protocol.StripTunnelIDPrefix(r.URL.RawPath, tunnelID)
+		if tunnelID == "" || !s.registry.Exists(tunnelID) {
+			http.Error(w, "Tunnel not found", http.StatusNotFound)
+			return
 		}
 
-		s.logger.Printf("Path routing: %s -> %s (tunnel: %s)", originalPath, r.URL.Path, tunnelID)
+		// Set cookie for future requests (when accessing via path with tunnel ID)
+		if fromPath {
+			http.SetCookie(w, &http.Cookie{
+				Name:     "exio_tunnel",
+				Value:    tunnelID,
+				Path:     "/",
+				MaxAge:   3600, // 1 hour
+				HttpOnly: true,
+				SameSite: http.SameSiteLaxMode,
+			})
+		}
+
+		// Only strip tunnel ID prefix if it was in the path
+		if fromPath {
+			originalPath := r.URL.Path
+			r.URL.Path = protocol.StripTunnelIDPrefix(r.URL.Path, tunnelID)
+			r.RequestURI = r.URL.RequestURI()
+
+			if r.URL.RawPath != "" {
+				r.URL.RawPath = protocol.StripTunnelIDPrefix(r.URL.RawPath, tunnelID)
+			}
+
+			s.logger.Printf("Path routing: %s -> %s (tunnel: %s)", originalPath, r.URL.Path, tunnelID)
+		}
 	} else {
 		// Extract subdomain from Host header (existing behavior)
 		tunnelID = protocol.ExtractSubdomain(r, s.config.BaseDomain)

pkg/protocol/protocol.go

@@ -178,3 +178,34 @@ func StripTunnelIDPrefix(path, tunnelID string) string {
 
 	return path
 }
+
+// ExtractTunnelIDFromReferer extracts the tunnel ID from a Referer URL.
+// For example, "https://tunnel.example.com/bold-owl-716/page" returns "bold-owl-716".
+func ExtractTunnelIDFromReferer(referer string) string {
+	if referer == "" {
+		return ""
+	}
+
+	// Find the path portion after the host
+	// Look for :// then find the next /
+	schemeEnd := 0
+	for i := 0; i < len(referer)-2; i++ {
+		if referer[i] == ':' && referer[i+1] == '/' && referer[i+2] == '/' {
+			schemeEnd = i + 3
+			break
+		}
+	}
+
+	// Find the start of the path (first / after the host)
+	pathStart := schemeEnd
+	for pathStart < len(referer) && referer[pathStart] != '/' {
+		pathStart++
+	}
+
+	if pathStart >= len(referer) {
+		return ""
+	}
+
+	// Extract tunnel ID from the path
+	return ExtractTunnelIDFromPath(referer[pathStart:])
+}