更新 proxy/rewriter.go

proxy/rewriter.go

@@ -10,9 +10,10 @@ import (
 
 type ContentRewriter struct {
     baseURL *url.URL
+    token   string
 }
 
-func NewContentRewriter(baseURL string) (*ContentRewriter, error) {
+func NewContentRewriter(baseURL, token string) (*ContentRewriter, error) {
     u, err := url.Parse(baseURL)
     if err != nil {
         return nil, err
@@ -20,14 +21,13 @@ func NewContentRewriter(baseURL string) (*ContentRewriter, error) {
     
     return &ContentRewriter{
         baseURL: u,
+        token:   token,
     }, nil
 }
 
-// RewriteHTML 重写 HTML 内容中的所有 URL
 func (r *ContentRewriter) RewriteHTML(body []byte) ([]byte, error) {
     doc, err := html.Parse(bytes.NewReader(body))
     if err != nil {
-        // 如果解析失败，使用简单的字符串替换
         return r.simpleRewriteHTML(body), nil
     }
     
@@ -41,16 +41,9 @@ func (r *ContentRewriter) RewriteHTML(body []byte) ([]byte, error) {
     return buf.Bytes(), nil
 }
 
-// rewriteNode 递归重写 HTML 节点
 func (r *ContentRewriter) rewriteNode(n *html.Node) {
     if n.Type == html.ElementNode {
-        // 重写需要处理的属性
-        attrs := map[string]bool{
-            "href":   true,
-            "src":    true,
-            "action": true,
-            "data":   true,
-        }
+        attrs := map[string]bool{"href": true, "src": true, "action": true, "data": true}
         
         for i, attr := range n.Attr {
             if attrs[attr.Key] {
@@ -59,31 +52,26 @@ func (r *ContentRewriter) rewriteNode(n *html.Node) {
                 }
             }
             
-            // 处理 srcset 属性
             if attr.Key == "srcset" {
                 n.Attr[i].Val = r.rewriteSrcset(attr.Val)
             }
             
-            // 处理 style 属性中的 URL
             if attr.Key == "style" {
                 n.Attr[i].Val = r.rewriteInlineCSS(attr.Val)
             }
         }
         
-            // 处理 <form> 标签 - 确保 action 属性存在
-            if n.Data == "form" {
-                hasAction := false
-                for i, attr := range n.Attr {
-                    if attr.Key == "action" {
-                        hasAction = true
-                        // 空 action 表示提交到当前页面
-                        if attr.Val == "" {
-                            n.Attr[i].Val = r.rewriteURL(r.baseURL.String())
-                        }
-                        break
+        if n.Data == "form" {
+            hasAction := false
+            for i, attr := range n.Attr {
+                if attr.Key == "action" {
+                    hasAction = true
+                    if attr.Val == "" {
+                        n.Attr[i].Val = r.rewriteURL(r.baseURL.String())
                     }
+                    break
                 }
-            // 如果没有 action 属性，添加一个
+            }
             if !hasAction {
                 n.Attr = append(n.Attr, html.Attribute{
                     Key: "action",
@@ -91,7 +79,7 @@ func (r *ContentRewriter) rewriteNode(n *html.Node) {
                 })
             }
         }
-        // 处理 <base> 标签
+        
         if n.Data == "base" {
             for i, attr := range n.Attr {
                 if attr.Key == "href" {
@@ -100,20 +88,16 @@ func (r *ContentRewriter) rewriteNode(n *html.Node) {
             }
         }
         
-        // 处理 <style> 标签内容
         if n.Data == "style" && n.FirstChild != nil {
             if n.FirstChild.Type == html.TextNode {
                 n.FirstChild.Data = r.rewriteInlineCSS(n.FirstChild.Data)
             }
         }
         
-        // 处理 <script> 标签，移除可能的跟踪脚本
         if n.Data == "script" {
             for _, attr := range n.Attr {
                 if attr.Key == "src" {
-                    // 可以在这里过滤掉已知的跟踪脚本
                     if r.isTrackingScript(attr.Val) {
-                        // 移除此节点
                         if n.Parent != nil {
                             n.Parent.RemoveChild(n)
                             return
@@ -124,17 +108,14 @@ func (r *ContentRewriter) rewriteNode(n *html.Node) {
         }
     }
     
-    // 递归处理子节点
     for c := n.FirstChild; c != nil; c = c.NextSibling {
         r.rewriteNode(c)
     }
 }
 
-// rewriteURL 重写单个 URL
 func (r *ContentRewriter) rewriteURL(urlStr string) string {
     urlStr = strings.TrimSpace(urlStr)
     
-    // 跳过特殊协议
     if strings.HasPrefix(urlStr, "javascript:") ||
         strings.HasPrefix(urlStr, "data:") ||
         strings.HasPrefix(urlStr, "mailto:") ||
@@ -144,22 +125,23 @@ func (r *ContentRewriter) rewriteURL(urlStr string) string {
         return urlStr
     }
     
-    // 解析 URL
     u, err := url.Parse(urlStr)
     if err != nil {
         return urlStr
     }
     
-    // 如果是相对 URL，转换为绝对 URL
     if !u.IsAbs() {
         u = r.baseURL.ResolveReference(u)
     }
     
-    // 生成代理 URL
-    return "/proxy?url=" + url.QueryEscape(u.String())
+    proxyPath := u.Path
+    if u.RawQuery != "" {
+        proxyPath += "?" + u.RawQuery
+    }
+    
+    return "/p/" + r.token + proxyPath
 }
 
-// rewriteSrcset 重写 srcset 属性
 func (r *ContentRewriter) rewriteSrcset(srcset string) string {
     if srcset == "" {
         return srcset
@@ -181,18 +163,14 @@ func (r *ContentRewriter) rewriteSrcset(srcset string) string {
     return strings.Join(rewritten, ", ")
 }
 
-// RewriteCSS 重写 CSS 内容
 func (r *ContentRewriter) RewriteCSS(body []byte) []byte {
     content := string(body)
     return []byte(r.rewriteInlineCSS(content))
 }
 
-// rewriteInlineCSS 重写内联 CSS 中的 URL
 func (r *ContentRewriter) rewriteInlineCSS(css string) string {
-    // 匹配 url(...) 模式
     result := css
     
-    // 处理 url("...") 和 url('...') 和 url(...)
     patterns := []struct {
         prefix string
         suffix string
@@ -227,22 +205,15 @@ func (r *ContentRewriter) rewriteInlineCSS(css string) string {
         }
     }
     
-    // 处理 @import
     result = r.rewriteImports(result)
     
     return result
 }
 
-// rewriteImports 重写 CSS @import 语句
 func (r *ContentRewriter) rewriteImports(css string) string {
     result := css
     
-    patterns := []string{
-        `@import "`,
-        `@import '`,
-        `@import url("`,
-        `@import url('`,
-    }
+    patterns := []string{`@import "`, `@import '`, `@import url("`, `@import url('`}
     
     for _, pattern := range patterns {
         start := 0
@@ -279,32 +250,28 @@ func (r *ContentRewriter) rewriteImports(css string) string {
     return result
 }
 
-// simpleRewriteHTML 简单的字符串替换重写（备用方案）
 func (r *ContentRewriter) simpleRewriteHTML(body []byte) []byte {
     content := string(body)
-    
-    // 重写绝对 URL
     baseStr := r.baseURL.Scheme + "://" + r.baseURL.Host
     
     replacements := []struct {
         old string
         new string
     }{
-        {`href="` + baseStr, `href="/proxy?url=` + url.QueryEscape(baseStr)},
-        {`src="` + baseStr, `src="/proxy?url=` + url.QueryEscape(baseStr)},
-        {`action="` + baseStr, `action="/proxy?url=` + url.QueryEscape(baseStr)},
-        {`href='` + baseStr, `href='/proxy?url=` + url.QueryEscape(baseStr)},
-        {`src='` + baseStr, `src='/proxy?url=` + url.QueryEscape(baseStr)},
+        {`href="` + baseStr, `href="/p/` + r.token},
+        {`src="` + baseStr, `src="/p/` + r.token},
+        {`action="` + baseStr, `action="/p/` + r.token},
+        {`href='` + baseStr, `href='/p/` + r.token},
+        {`src='` + baseStr, `src='/p/` + r.token},
     }
     
-    for _, r := range replacements {
-        content = strings.ReplaceAll(content, r.old, r.new)
+    for _, rep := range replacements {
+        content = strings.ReplaceAll(content, rep.old, rep.new)
     }
     
     return []byte(content)
 }
 
-// isTrackingScript 检查是否是跟踪脚本
 func (r *ContentRewriter) isTrackingScript(src string) bool {
     trackingDomains := []string{
         "google-analytics.com",
@@ -324,4 +291,4 @@ func (r *ContentRewriter) isTrackingScript(src string) bool {
     }
     
     return false
-}
+}