From 576f7989f6fab5db7a26d576b66d47bd9197cae6 Mon Sep 17 00:00:00 2001
From: XOF <xofine@noreply.localhost>
Date: Mon, 15 Dec 2025 19:47:50 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20proxy/handler.go?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 proxy/handler.go | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/proxy/handler.go b/proxy/handler.go
index 21b33c3..5b80953 100644
--- a/proxy/handler.go
+++ b/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://") {