From 1159dd6982c9a669639bd1356abd109892e56d2b Mon Sep 17 00:00:00 2001
From: Dave <dave@mac-mini.local>
Date: Sun, 21 Jan 2024 21:28:21 -0500
Subject: [PATCH] Add `addDynamicIframe` rewrite function.

Add unit tests for `add_dynamic_iframe` rewrite.
---
 internal/reader/rewrite/rewrite_functions.go | 37 +++++++++++++++
 internal/reader/rewrite/rewriter.go          |  2 +
 internal/reader/rewrite/rewriter_test.go     | 48 ++++++++++++++++++++
 3 files changed, 87 insertions(+)

diff --git a/internal/reader/rewrite/rewrite_functions.go b/internal/reader/rewrite/rewrite_functions.go
index b5d63464..e84f3d3c 100644
--- a/internal/reader/rewrite/rewrite_functions.go
+++ b/internal/reader/rewrite/rewrite_functions.go
@@ -167,6 +167,43 @@ func addDynamicImage(entryURL, entryContent string) string {
 	return entryContent
 }
 
+func addDynamicIframe(entryURL, entryContent string) string {
+	doc, err := goquery.NewDocumentFromReader(strings.NewReader(entryContent))
+	if err != nil {
+		return entryContent
+	}
+
+	// Ordered most preferred to least preferred.
+	candidateAttrs := []string{
+		"data-src",
+		"data-original",
+		"data-orig",
+		"data-url",
+		"data-lazy-src",
+	}
+
+	changed := false
+
+	doc.Find("iframe").Each(func(i int, iframe *goquery.Selection) {
+		for _, candidateAttr := range candidateAttrs {
+			if srcAttr, found := iframe.Attr(candidateAttr); found {
+				changed = true
+
+				iframe.SetAttr("src", srcAttr)
+
+				break
+			}
+		}
+	})
+
+	if changed {
+		output, _ := doc.Find("body").First().Html()
+		return output
+	}
+
+	return entryContent
+}
+
 func fixMediumImages(entryURL, entryContent string) string {
 	doc, err := goquery.NewDocumentFromReader(strings.NewReader(entryContent))
 	if err != nil {
diff --git a/internal/reader/rewrite/rewriter.go b/internal/reader/rewrite/rewriter.go
index 9885890f..08fa1fa9 100644
--- a/internal/reader/rewrite/rewriter.go
+++ b/internal/reader/rewrite/rewriter.go
@@ -69,6 +69,8 @@ func applyRule(entryURL string, entry *model.Entry, rule rule) {
 		entry.Content = addMailtoSubject(entryURL, entry.Content)
 	case "add_dynamic_image":
 		entry.Content = addDynamicImage(entryURL, entry.Content)
+	case "add_dynamic_iframe":
+		entry.Content = addDynamicIframe(entryURL, entry.Content)
 	case "add_youtube_video":
 		entry.Content = addYoutubeVideo(entryURL, entry.Content)
 	case "add_invidious_video":
diff --git a/internal/reader/rewrite/rewriter_test.go b/internal/reader/rewrite/rewriter_test.go
index bbf5bfdd..3910ccd4 100644
--- a/internal/reader/rewrite/rewriter_test.go
+++ b/internal/reader/rewrite/rewriter_test.go
@@ -333,6 +333,54 @@ func TestRewriteWithImageAndLazySrcset(t *testing.T) {
 	}
 }
 
+func TestRewriteWithNoLazyIframe(t *testing.T) {
+	controlEntry := &model.Entry{
+		Title:   `A title`,
+		Content: `<iframe src="https://example.org/embed" allowfullscreen></iframe>`,
+	}
+	testEntry := &model.Entry{
+		Title:   `A title`,
+		Content: `<iframe src="https://example.org/embed" allowfullscreen></iframe>`,
+	}
+	Rewriter("https://example.org/article", testEntry, "add_dynamic_iframe")
+
+	if !reflect.DeepEqual(testEntry, controlEntry) {
+		t.Errorf(`Not expected output: got "%+v" instead of "%+v"`, testEntry, controlEntry)
+	}
+}
+
+func TestRewriteWithLazyIframe(t *testing.T) {
+	controlEntry := &model.Entry{
+		Title:   `A title`,
+		Content: `<iframe data-src="https://example.org/embed" allowfullscreen="" src="https://example.org/embed"></iframe>`,
+	}
+	testEntry := &model.Entry{
+		Title:   `A title`,
+		Content: `<iframe data-src="https://example.org/embed" allowfullscreen></iframe>`,
+	}
+	Rewriter("https://example.org/article", testEntry, "add_dynamic_iframe")
+
+	if !reflect.DeepEqual(testEntry, controlEntry) {
+		t.Errorf(`Not expected output: got "%+v" instead of "%+v"`, testEntry, controlEntry)
+	}
+}
+
+func TestRewriteWithLazyIframeAndSrc(t *testing.T) {
+	controlEntry := &model.Entry{
+		Title:   `A title`,
+		Content: `<iframe src="https://example.org/embed" data-src="https://example.org/embed" allowfullscreen=""></iframe>`,
+	}
+	testEntry := &model.Entry{
+		Title:   `A title`,
+		Content: `<iframe src="about:blank" data-src="https://example.org/embed" allowfullscreen></iframe>`,
+	}
+	Rewriter("https://example.org/article", testEntry, "add_dynamic_iframe")
+
+	if !reflect.DeepEqual(testEntry, controlEntry) {
+		t.Errorf(`Not expected output: got "%+v" instead of "%+v"`, testEntry, controlEntry)
+	}
+}
+
 func TestNewLineRewriteRule(t *testing.T) {
 	controlEntry := &model.Entry{
 		Title:   `A title`,