From 43506ab6e10cc9dd0c9c67a980cea7af96c0bb79 Mon Sep 17 00:00:00 2001
From: XOF <xofine@noreply.localhost>
Date: Mon, 15 Dec 2025 18:00:49 +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 | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/proxy/handler.go b/proxy/handler.go
index 1469212..093135d 100644
--- a/proxy/handler.go
+++ b/proxy/handler.go
@@ -6,6 +6,7 @@ import (
     "io"
     "log"
     "net/http"
+    "net/url"
     "strings"
     "time"
     
@@ -143,6 +144,35 @@ func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
     }
     defer resp.Body.Close()
     
+    // 处理 404：尝试从 Referer 提取正确路径
+    if resp.StatusCode == 404 {
+        if referer := r.Header.Get("Referer"); referer != "" && strings.Contains(referer, "/p/"+token) {
+            if u, err := url.Parse(targetURL); err == nil && u.Host != "" {
+                // 直接使用原始 targetURL 的 host + path
+                baseURL := u.Scheme + "://" + u.Host
+                correctPath := u.Path
+                if u.RawQuery != "" {
+                    correctPath += "?" + u.RawQuery
+                }
+                
+                newTargetURL := baseURL + correctPath
+                
+                // 避免重复请求相同 URL
+                if newTargetURL != targetURL {
+                    retryReq, _ := http.NewRequest(r.Method, newTargetURL, nil)
+                    h.setProxyHeaders(retryReq, r)
+                    
+                    if retryResp, err := client.Do(retryReq); err == nil && retryResp.StatusCode == 200 {
+                        resp.Body.Close()
+                        resp = retryResp
+                        defer resp.Body.Close()
+                        log.Printf("404 retry success: %s -> %s", targetURL, newTargetURL)
+                    }
+                }
+            }
+        }
+    }
+    
     body, err := h.readResponseBody(resp)
     if err != nil {
         log.Printf("Failed to read response: %v", err)
@@ -169,7 +199,6 @@ func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
     h.sendResponse(w, resp, body)
 }
 
-
 func (h *ProxyHandler) rewriteContent(body []byte, targetURL, contentType, token string) []byte {
     rewriter, err := NewContentRewriter(targetURL, token)
     if err != nil {