fix pypi

internal/proxy/pypi_rewrite.go

@@ -0,0 +1,126 @@
+package proxy
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
+
+	"golang.org/x/net/html"
+
+	"github.com/any-hub/any-hub/internal/server"
+)
+
+func (h *Handler) rewritePyPIResponse(route *server.HubRoute, resp *http.Response, path string) (*http.Response, error) {
+	if resp == nil {
+		return resp, nil
+	}
+	if !strings.HasPrefix(path, "/simple") && path != "/" {
+		return resp, nil
+	}
+	bodyBytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return resp, err
+	}
+	resp.Body.Close()
+
+	rewritten, contentType, err := rewritePyPIBody(bodyBytes, resp.Header.Get("Content-Type"), route.Config.Domain)
+	if err != nil {
+		resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
+		return resp, err
+	}
+
+	resp.Body = io.NopCloser(bytes.NewReader(rewritten))
+	resp.ContentLength = int64(len(rewritten))
+	resp.Header.Set("Content-Length", strconv.Itoa(len(rewritten)))
+	if contentType != "" {
+		resp.Header.Set("Content-Type", contentType)
+	}
+	resp.Header.Del("Content-Encoding")
+	return resp, nil
+}
+
+func rewritePyPIBody(body []byte, contentType string, domain string) ([]byte, string, error) {
+	lowerCT := strings.ToLower(contentType)
+	if strings.Contains(lowerCT, "application/vnd.pypi.simple.v1+json") || strings.HasPrefix(strings.TrimSpace(string(body)), "{") {
+		data := map[string]interface{}{}
+		if err := json.Unmarshal(body, &data); err != nil {
+			return body, contentType, err
+		}
+		if files, ok := data["files"].([]interface{}); ok {
+			for _, entry := range files {
+				if fileMap, ok := entry.(map[string]interface{}); ok {
+					if urlValue, ok := fileMap["url"].(string); ok {
+						fileMap["url"] = rewritePyPIFileURL(domain, urlValue)
+					}
+				}
+			}
+		}
+		rewriteBytes, err := json.Marshal(data)
+		if err != nil {
+			return body, contentType, err
+		}
+		return rewriteBytes, "application/vnd.pypi.simple.v1+json", nil
+	}
+
+	rewrittenHTML, err := rewritePyPIHTML(body, domain)
+	if err != nil {
+		return body, contentType, err
+	}
+	return rewrittenHTML, "text/html; charset=utf-8", nil
+}
+
+func rewritePyPIHTML(body []byte, domain string) ([]byte, error) {
+	node, err := html.Parse(bytes.NewReader(body))
+	if err != nil {
+		return nil, err
+	}
+	rewriteHTMLNode(node, domain)
+	var buf bytes.Buffer
+	if err := html.Render(&buf, node); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+func rewriteHTMLNode(n *html.Node, domain string) {
+	if n.Type == html.ElementNode {
+		rewriteHTMLAttributes(n, domain)
+	}
+	for child := n.FirstChild; child != nil; child = child.NextSibling {
+		rewriteHTMLNode(child, domain)
+	}
+}
+
+func rewriteHTMLAttributes(n *html.Node, domain string) {
+	for i, attr := range n.Attr {
+		switch attr.Key {
+		case "href", "data-dist-info-metadata", "data-core-metadata":
+			if strings.HasPrefix(attr.Val, "http://") || strings.HasPrefix(attr.Val, "https://") {
+				n.Attr[i].Val = rewritePyPIFileURL(domain, attr.Val)
+			}
+		}
+	}
+}
+
+func rewritePyPIFileURL(domain, original string) string {
+	parsed, err := url.Parse(original)
+	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
+		return original
+	}
+	prefix := "/files/" + parsed.Scheme + "/" + parsed.Host
+	newURL := url.URL{
+		Scheme:   "https",
+		Host:     domain,
+		Path:     prefix + parsed.Path,
+		RawQuery: parsed.RawQuery,
+		Fragment: parsed.Fragment,
+	}
+	if raw := parsed.RawPath; raw != "" {
+		newURL.RawPath = prefix + raw
+	}
+	return newURL.String()
+}