更新 proxy/handler.go

proxy/handler.go

@@ -63,10 +63,25 @@ func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
     
     session := h.sessionManager.Get(token)
     if session == nil {
-        http.Error(w, "Session expired or invalid", http.StatusUnauthorized)
-        return
+        // 尝试从 cookie 恢复
+        if cookie, err := r.Cookie("ORIGINALHOST_" + token); err == nil {
+            session = &ProxySession{TargetURL: cookie.Value}
+        } else {
+            http.Error(w, "Session expired or invalid", http.StatusUnauthorized)
+            return
+        }
     }
     
+    // 设置 cookie（10分钟有效期）
+    http.SetCookie(w, &http.Cookie{
+        Name:     "ORIGINALHOST_" + token,
+        Value:    session.TargetURL,
+        Path:     "/p/" + token,
+        MaxAge:   600,
+        HttpOnly: true,
+        SameSite: http.SameSiteLaxMode,
+    })
+    
     var targetURL string
     
     if strings.HasPrefix(targetPath, "http:/") && !strings.HasPrefix(targetPath, "http://") {