chore(scraper): improve data-src rewriting rule

ncarlier · Oct 2, 2024 · ab0f6b5 · ab0f6b5
1 parent 18202fa
commit ab0f6b5
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 24 deletions.
diff --git a/pkg/sanitizer/sanitizer.go b/pkg/sanitizer/sanitizer.go
@@ -27,7 +27,7 @@ func NewSanitizer(blockList *BlockList) *Sanitizer {
 	policy := bluemonday.UGCPolicy()
 	policy.AddTargetBlankToFullyQualifiedLinks(true)
 	policy.AllowAttrs("width", "height", "src", "allowfullscreen", "sandbox").OnElements("iframe")
-	policy.AllowAttrs("srcset", "sizes", "data-src").OnElements("img", "source")
+	policy.AllowAttrs("srcset", "sizes").OnElements("img", "source")
 	policy.AllowElements("picture", "source")
 
 	if blockList != nil {

diff --git a/pkg/scraper/html-rewriter/helper.go b/pkg/scraper/html-rewriter/helper.go
@@ -8,8 +8,8 @@ func newHTMLAttribute(key string) html.Attribute {
 	}
 }
 
-func findHTMLAttribute(attributes []html.Attribute, key string) *html.Attribute {
-	for _, attr := range attributes {
+func findHTMLAttribute(node *html.Node, key string) *html.Attribute {
+	for _, attr := range node.Attr {
 		if attr.Key == key {
 			return &attr
 		}

diff --git a/pkg/scraper/html-rewriter/img-rewriter.go b/pkg/scraper/html-rewriter/img-rewriter.go
@@ -6,26 +6,19 @@ import (
 )
 
 func RewriteDataSrcToSrcAttribute(node *html.Node) {
-	if node.DataAtom != atom.Img {
+	if node.DataAtom != atom.Img && node.DataAtom != atom.Source {
 		return
 	}
 
-	attrs := []html.Attribute{}
-	srcAttr := newHTMLAttribute("src")
-	dataSrcAttr := newHTMLAttribute("data-src")
-	for _, attr := range node.Attr {
-		switch attr.Key {
-		case srcAttr.Key:
-			srcAttr = attr
-		case dataSrcAttr.Key:
-			dataSrcAttr = attr
-		default:
-			attrs = append(attrs, attr)
+	keys := []string{"src", "srcset"}
+	for _, key := range keys {
+		realAttr := findHTMLAttribute(node, key)
+		if realAttr == nil || realAttr.Val == "" {
+			dataAttr := findHTMLAttribute(node, "data-"+key)
+			if dataAttr != nil && dataAttr.Val != "" {
+				dataAttr.Key = key
+				replaceHTMLAttribute(node, *dataAttr)
+			}
 		}
 	}
-	if dataSrcAttr.Val != "" && srcAttr.Val == "" {
-		srcAttr.Val = dataSrcAttr.Val
-		attrs = append(attrs, srcAttr)
-		node.Attr = attrs
-	}
 }
diff --git a/pkg/scraper/html-rewriter/picture-rewriter.go b/pkg/scraper/html-rewriter/picture-rewriter.go
@@ -25,11 +25,11 @@ func RewritePictureWithoutImgSrcAttribute(node *html.Node) {
 		return
 	}
 
-	srcAttr := findHTMLAttribute(imgNode.Attr, "src")
+	srcAttr := findHTMLAttribute(imgNode, "src")
 	if srcAttr != nil && srcAttr.Val != "" {
 		return
 	}
-	srcsetAttr := findHTMLAttribute(sourceNode.Attr, "srcset")
+	srcsetAttr := findHTMLAttribute(sourceNode, "srcset")
 	if srcsetAttr == nil || srcsetAttr.Val == "" {
 		return
 	}

diff --git a/pkg/scraper/html-rewriter/test/rewrite_test.go b/pkg/scraper/html-rewriter/test/rewrite_test.go
@@ -18,7 +18,7 @@ var testCases = []struct {
 }{
 	{
 		input:    "<html><body><img data-src=\"foo\" /></body></html>",
-		expected: "<html><head></head><body><img src=\"foo\"/></body></html>",
+		expected: "<html><head></head><body><img data-src=\"foo\" src=\"foo\"/></body></html>",
 	},
 	{
 		input:    "<html><body><img data-src=\"foo\" src=\"bar\" /></body></html>",

diff --git a/pkg/scraper/test/web-scraper_test.go b/pkg/scraper/test/web-scraper_test.go
@@ -22,7 +22,7 @@ func TestSimpleWebScraping(t *testing.T) {
 	assert.Equal(t, "readflow", page.Title)
 	assert.Equal(t, "read your Internet article flow in one place with complete peace of mind and freedom", page.Text)
 	assert.Contains(t, page.HTML, "relax.png")
-	assert.Equal(t, "https://about.readflow.app/images/readflow.png", page.Image)
+	assert.Equal(t, "https://about.readflow.app/img/readflow.png", page.Image)
 	assert.Equal(t, "https://about.readflow.app/favicon.png", page.Favicon)
 }