添加 proxy/handler.go

proxy/handler.go

@@ -0,0 +1,214 @@
+// proxy/handler.go
+package proxy
+
+import (
+    "bytes"
+    "fmt"
+    "io"
+    "log"
+    "net/http"
+    "net/url"
+    "strings"
+    "time"
+    
+    "siteproxy/cache"
+    "siteproxy/security"
+)
+
+type ProxyHandler struct {
+    validator   *security.RequestValidator
+    rateLimiter *security.RateLimiter
+    cache       *cache.MemoryCache
+    client      *http.Client
+    userAgent   string
+    maxSize     int64
+}
+
+func NewHandler(validator *security.RequestValidator, rateLimiter *security.RateLimiter, memCache *cache.MemoryCache, userAgent string, maxSize int64) *ProxyHandler {
+    return &ProxyHandler{
+        validator:   validator,
+        rateLimiter: rateLimiter,
+        cache:       memCache,
+        userAgent:   userAgent,
+        maxSize:     maxSize,
+        client: &http.Client{
+            Timeout: 30 * time.Second,
+            CheckRedirect: func(req *http.Request, via []*http.Request) error {
+                if len(via) >= 10 {
+                    return fmt.Errorf("too many redirects")
+                }
+                return nil
+            },
+        },
+    }
+}
+
+func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+    // 获取 session ID
+    cookie, err := r.Cookie("session_id")
+    if err != nil {
+        http.Error(w, "Unauthorized", http.StatusUnauthorized)
+        return
+    }
+    
+    sessionID := cookie.Value
+    
+    // 速率限制检查
+    if !h.rateLimiter.Allow(sessionID) {
+        http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
+        return
+    }
+    
+    // 获取目标 URL
+    targetURL := r.URL.Query().Get("url")
+    if targetURL == "" {
+        http.Error(w, "Missing url parameter", http.StatusBadRequest)
+        return
+    }
+    
+    // 验证 URL
+    if err := h.validator.ValidateURL(targetURL); err != nil {
+        log.Printf("URL validation failed: %v", err)
+        http.Error(w, "Invalid or blocked URL", http.StatusForbidden)
+        return
+    }
+    
+    // 检查缓存
+    cacheKey := h.cache.GenerateKey(targetURL)
+    if cached, ok := h.cache.Get(cacheKey); ok {
+        log.Printf("Cache hit: %s", targetURL)
+        h.serveCached(w, cached)
+        return
+    }
+    
+    // 发起代理请求
+    log.Printf("Proxying: %s", targetURL)
+    if err := h.proxyRequest(w, targetURL, cacheKey); err != nil {
+        log.Printf("Proxy error: %v", err)
+        http.Error(w, "Proxy request failed", http.StatusBadGateway)
+        return
+    }
+}
+
+func (h *ProxyHandler) proxyRequest(w http.ResponseWriter, targetURL, cacheKey string) error {
+    // 创建请求
+    req, err := http.NewRequest("GET", targetURL, nil)
+    if err != nil {
+        return err
+    }
+    
+    // 设置请求头
+    req.Header.Set("User-Agent", h.userAgent)
+    req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
+    req.Header.Set("Accept-Language", "en-US,en;q=0.9")
+    req.Header.Set("Accept-Encoding", "gzip, deflate, br")
+    req.Header.Set("DNT", "1")
+    req.Header.Set("Connection", "keep-alive")
+    req.Header.Set("Upgrade-Insecure-Requests", "1")
+    
+    // 发送请求
+    resp, err := h.client.Do(req)
+    if err != nil {
+        return err
+    }
+    defer resp.Body.Close()
+    
+    // 限制响应大小
+    limitedReader := io.LimitReader(resp.Body, h.maxSize)
+    body, err := io.ReadAll(limitedReader)
+    if err != nil {
+        return err
+    }
+    
+    // 检查内容类型
+    contentType := resp.Header.Get("Content-Type")
+    
+    // 如果是 HTML，进行 URL 重写
+    if strings.Contains(contentType, "text/html") {
+        body = h.rewriteHTML(body, targetURL)
+    } else if strings.Contains(contentType, "text/css") {
+        body = h.rewriteCSS(body, targetURL)
+    }
+    
+    // 缓存静态资源
+    if h.shouldCache(contentType) {
+        headers := make(map[string]string)
+        headers["Content-Type"] = contentType
+        h.cache.Set(cacheKey, body, headers)
+    }
+    
+    // 设置响应头
+    w.Header().Set("Content-Type", contentType)
+    w.Header().Set("X-Proxied-By", "SiteProxy")
+    
+    // 移除可能泄露的头
+    for _, header := range []string{"Server", "X-Powered-By", "Set-Cookie"} {
+        w.Header().Del(header)
+    }
+    
+    w.WriteHeader(resp.StatusCode)
+    w.Write(body)
+    
+    return nil
+}
+
+func (h *ProxyHandler) rewriteHTML(body []byte, baseURL string) []byte {
+    content := string(body)
+    
+    // 解析基础 URL
+    base, err := url.Parse(baseURL)
+    if err != nil {
+        return body
+    }
+    
+    // 重写绝对 URL
+    content = strings.ReplaceAll(content, `href="`+base.Scheme+`://`+base.Host, `href="/proxy?url=`+url.QueryEscape(base.Scheme+`://`+base.Host))
+    content = strings.ReplaceAll(content, `src="`+base.Scheme+`://`+base.Host, `src="/proxy?url=`+url.QueryEscape(base.Scheme+`://`+base.Host))
+    
+    // 重写相对 URL（简化版本）
+    // 注意：完整实现需要 HTML 解析器
+    
+    return []byte(content)
+}
+
+func (h *ProxyHandler) rewriteCSS(body []byte, baseURL string) []byte {
+    content := string(body)
+    
+    base, err := url.Parse(baseURL)
+    if err != nil {
+        return body
+    }
+    
+    // 重写 CSS 中的 url()
+    content = strings.ReplaceAll(content, `url(`+base.Scheme+`://`+base.Host, `url(/proxy?url=`+url.QueryEscape(base.Scheme+`://`+base.Host))
+    
+    return []byte(content)
+}
+
+func (h *ProxyHandler) shouldCache(contentType string) bool {
+    cacheableTypes := []string{
+        "image/",
+        "text/css",
+        "application/javascript",
+        "application/json",
+        "font/",
+    }
+    
+    for _, t := range cacheableTypes {
+        if strings.Contains(contentType, t) {
+            return true
+        }
+    }
+    
+    return false
+}
+
+func (h *ProxyHandler) serveCached(w http.ResponseWriter, entry *cache.CacheEntry) {
+    for key, value := range entry.Headers {
+        w.Header().Set(key, value)
+    }
+    w.Header().Set("X-Cache-Status", "HIT")
+    w.Header().Set("X-Proxied-By", "SiteProxy")
+    w.WriteHeader(http.StatusOK)
+    w.Write(entry.Data)
+}