添加 proxy/rewriter.go

proxy/rewriter.go

@@ -0,0 +1,306 @@
+//proxy/rewriter.go
+package proxy
+
+import (
+    "bytes"
+    "golang.org/x/net/html"
+    "net/url"
+    "strings"
+)
+
+type ContentRewriter struct {
+    baseURL *url.URL
+}
+
+func NewContentRewriter(baseURL string) (*ContentRewriter, error) {
+    u, err := url.Parse(baseURL)
+    if err != nil {
+        return nil, err
+    }
+    
+    return &ContentRewriter{
+        baseURL: u,
+    }, 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
+    }
+    
+    r.rewriteNode(doc)
+    
+    var buf bytes.Buffer
+    if err := html.Render(&buf, doc); err != nil {
+        return r.simpleRewriteHTML(body), nil
+    }
+    
+    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,
+        }
+        
+        for i, attr := range n.Attr {
+            if attrs[attr.Key] {
+                if rewritten := r.rewriteURL(attr.Val); rewritten != attr.Val {
+                    n.Attr[i].Val = rewritten
+                }
+            }
+            
+            // 处理 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)
+            }
+        }
+        
+        // 处理 <base> 标签
+        if n.Data == "base" {
+            for i, attr := range n.Attr {
+                if attr.Key == "href" {
+                    n.Attr[i].Val = r.baseURL.String()
+                }
+            }
+        }
+        
+        // 处理 <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
+                        }
+                    }
+                }
+            }
+        }
+    }
+    
+    // 递归处理子节点
+    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:") ||
+        strings.HasPrefix(urlStr, "tel:") ||
+        strings.HasPrefix(urlStr, "#") ||
+        urlStr == "" {
+        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())
+}
+
+// rewriteSrcset 重写 srcset 属性
+func (r *ContentRewriter) rewriteSrcset(srcset string) string {
+    if srcset == "" {
+        return srcset
+    }
+    
+    parts := strings.Split(srcset, ",")
+    var rewritten []string
+    
+    for _, part := range parts {
+        part = strings.TrimSpace(part)
+        fields := strings.Fields(part)
+        
+        if len(fields) > 0 {
+            fields[0] = r.rewriteURL(fields[0])
+            rewritten = append(rewritten, strings.Join(fields, " "))
+        }
+    }
+    
+    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
+    }{
+        {`url("`, `")`},
+        {`url('`, `')`},
+        {`url(`, `)`},
+    }
+    
+    for _, pattern := range patterns {
+        start := 0
+        for {
+            idx := strings.Index(result[start:], pattern.prefix)
+            if idx == -1 {
+                break
+            }
+            
+            idx += start
+            urlStart := idx + len(pattern.prefix)
+            urlEnd := strings.Index(result[urlStart:], pattern.suffix)
+            
+            if urlEnd == -1 {
+                break
+            }
+            
+            urlEnd += urlStart
+            originalURL := result[urlStart:urlEnd]
+            rewrittenURL := r.rewriteURL(originalURL)
+            
+            result = result[:urlStart] + rewrittenURL + result[urlEnd:]
+            start = urlStart + len(rewrittenURL)
+        }
+    }
+    
+    // 处理 @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('`,
+    }
+    
+    for _, pattern := range patterns {
+        start := 0
+        for {
+            idx := strings.Index(result[start:], pattern)
+            if idx == -1 {
+                break
+            }
+            
+            idx += start
+            urlStart := idx + len(pattern)
+            
+            var endChar string
+            if strings.Contains(pattern, `"`) {
+                endChar = `"`
+            } else {
+                endChar = `'`
+            }
+            
+            urlEnd := strings.Index(result[urlStart:], endChar)
+            if urlEnd == -1 {
+                break
+            }
+            
+            urlEnd += urlStart
+            originalURL := result[urlStart:urlEnd]
+            rewrittenURL := r.rewriteURL(originalURL)
+            
+            result = result[:urlStart] + rewrittenURL + result[urlEnd:]
+            start = urlStart + len(rewrittenURL)
+        }
+    }
+    
+    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)},
+    }
+    
+    for _, r := range replacements {
+        content = strings.ReplaceAll(content, r.old, r.new)
+    }
+    
+    return []byte(content)
+}
+
+// isTrackingScript 检查是否是跟踪脚本
+func (r *ContentRewriter) isTrackingScript(src string) bool {
+    trackingDomains := []string{
+        "google-analytics.com",
+        "googletagmanager.com",
+        "facebook.net",
+        "doubleclick.net",
+        "analytics.js",
+        "ga.js",
+        "gtag.js",
+    }
+    
+    srcLower := strings.ToLower(src)
+    for _, domain := range trackingDomains {
+        if strings.Contains(srcLower, domain) {
+            return true
+        }
+    }
+    
+    return false
+}