goget/internal/linkrewrite/rewriter.go

//go:build linux || freebsd
// +build linux freebsd

// Package linkrewrite rewrites HTML links for local offline viewing.
package linkrewrite

import (
	"net/url"
	"path/filepath"
	"strings"
	"sync"

	"golang.org/x/net/html"
)

// Rewriter converts absolute URLs in HTML to local relative paths.
type Rewriter struct {
	baseURL   *url.URL
	outputDir string
	visited   map[string]string
	visitedMu sync.RWMutex
}

// New creates a link rewriter for a mirror or recursive download.
func New(baseURL *url.URL, outputDir string) *Rewriter {
	return &Rewriter{
		baseURL:   baseURL,
		outputDir: outputDir,
		visited:   make(map[string]string),
	}
}

// Register records a downloaded URL and its local path.
func (lr *Rewriter) Register(u *url.URL, localPath string) {
	lr.visitedMu.Lock()
	defer lr.visitedMu.Unlock()
	lr.visited[u.String()] = localPath
}

// RewriteLinks converts URLs in an HTML document to local paths.
func (lr *Rewriter) RewriteLinks(doc *html.Node, filePath string) int {
	converted := 0
	dir := filepath.Dir(filePath)

	var traverse func(*html.Node)
	traverse = func(n *html.Node) {
		if n.Type == html.ElementNode {
			switch n.Data {
			case "a":
				converted += lr.rewriteAttr(n, "href", dir)
			case "img", "script", "iframe", "video", "audio", "source", "track", "embed", "object":
				converted += lr.rewriteAttr(n, "src", dir)
			case "link":
				converted += lr.rewriteAttr(n, "href", dir)
			case "form":
				converted += lr.rewriteAttr(n, "action", dir)
			case "area":
				converted += lr.rewriteAttr(n, "href", dir)
			case "meta":
				for _, attr := range n.Attr {
					if attr.Key == "content" && strings.Contains(attr.Val, "url=") {
						converted += lr.rewriteMetaRefresh(n, dir)
					}
					if attr.Key == "property" && strings.Contains(attr.Val, "image") {
						converted += lr.rewriteAttr(n, "content", dir)
					}
				}
			}
		}
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			traverse(c)
		}
	}

	traverse(doc)
	return converted
}

func (lr *Rewriter) rewriteMetaRefresh(n *html.Node, baseDir string) int {
	for i, attr := range n.Attr {
		if attr.Key == "content" {
			content := attr.Val
			idx := strings.Index(strings.ToLower(content), "url=")
			if idx == -1 {
				return 0
			}
			urlStr := strings.TrimSpace(content[idx+4:])
			parsed, err := url.Parse(urlStr)
			if err != nil {
				return 0
			}
			resolved := lr.baseURL.ResolveReference(parsed)
			lr.visitedMu.RLock()
			localPath, exists := lr.visited[resolved.String()]
			lr.visitedMu.RUnlock()
			if exists {
				relPath, err := filepath.Rel(baseDir, localPath)
				if err != nil {
					relPath = localPath
				}
				n.Attr[i].Val = content[:idx+4] + filepath.ToSlash(relPath)
				return 1
			}
			return 0
		}
	}
	return 0
}

func (lr *Rewriter) rewriteAttr(n *html.Node, attrName, baseDir string) int {
	for i, attr := range n.Attr {
		if attr.Key != attrName {
			continue
		}
		value := attr.Val
		if strings.HasPrefix(value, "#") ||
			strings.HasPrefix(value, "javascript:") ||
			strings.HasPrefix(value, "mailto:") ||
			strings.HasPrefix(value, "tel:") ||
			strings.HasPrefix(value, "data:") ||
			strings.HasPrefix(value, "about:") {
			return 0
		}
		parsed, err := url.Parse(value)
		if err != nil {
			return 0
		}
		resolved := lr.baseURL.ResolveReference(parsed)
		lr.visitedMu.RLock()
		localPath, exists := lr.visited[resolved.String()]
		lr.visitedMu.RUnlock()
		if !exists {
			return 0
		}
		relPath, err := filepath.Rel(baseDir, localPath)
		if err != nil {
			relPath = localPath
		}
		n.Attr[i].Val = filepath.ToSlash(relPath)
		return 1
	}
	return 0
}