From fb590aceee50bf38efd71f5d5080ceb9a002ac7a Mon Sep 17 00:00:00 2001
From: Yewolf <yyewolf@gmail.com>
Date: Wed, 12 Jul 2023 09:50:52 +0200
Subject: [PATCH] Added src rewriter

---
 policy.go        | 14 ++++++++++++++
 sanitize.go      |  8 ++++++++
 sanitize_test.go | 15 +++++++++++++--
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/policy.go b/policy.go
index 995f46c..1b3f836 100644
--- a/policy.go
+++ b/policy.go
@@ -121,6 +121,11 @@ type Policy struct {
 	// if one would want to allow all URL schemes, they would add `.+`
 	allowURLSchemeRegexps []*regexp.Regexp
 
+	// If srcRewriter is not nil, it is used to rewrite the src attribute
+	// of tags that download resources, such as <img> and <script>.
+	// It requires that the URL is parsable by "net/url" url.Parse().
+	srcRewriter urlRewriter
+
 	// If an element has had all attributes removed as a result of a policy
 	// being applied, then the element would be removed from the output.
 	//
@@ -196,6 +201,8 @@ type stylePolicyBuilder struct {
 
 type urlPolicy func(url *url.URL) (allowUrl bool)
 
+type urlRewriter func(*url.URL)
+
 type SandboxValue int64
 
 const (
@@ -575,6 +582,13 @@ func (p *Policy) AllowURLSchemesMatching(r *regexp.Regexp) *Policy {
 	return p
 }
 
+// RewriteSrc will rewrite the src attribute of a resource downloading tag
+// (e.g. <img>, <script>, <iframe>) using the provided function.
+func (p *Policy) RewriteSrc(fn urlRewriter) *Policy {
+	p.srcRewriter = fn
+	return p
+}
+
 // RequireNoFollowOnLinks will result in all a, area, link tags having a
 // rel="nofollow"added to them if one does not already exist
 //
diff --git a/sanitize.go b/sanitize.go
index 9121aef..2b4485e 100644
--- a/sanitize.go
+++ b/sanitize.go
@@ -612,6 +612,14 @@ attrsLoop:
 				case "audio", "embed", "iframe", "img", "script", "source", "track", "video":
 					if htmlAttr.Key == "src" {
 						if u, ok := p.validURL(htmlAttr.Val); ok {
+							if p.srcRewriter != nil {
+								parsedURL, err := url.Parse(u)
+								if err != nil {
+									fmt.Println(err)
+								}
+								p.srcRewriter(parsedURL)
+								u = parsedURL.String()
+							}
 							htmlAttr.Val = u
 							tmpAttrs = append(tmpAttrs, htmlAttr)
 						}
diff --git a/sanitize_test.go b/sanitize_test.go
index 4e3a08f..eb7f094 100644
--- a/sanitize_test.go
+++ b/sanitize_test.go
@@ -150,16 +150,27 @@ func TestLinks(t *testing.T) {
 		},
 		{
 			in:       `<img src="giraffe.gif" />`,
-			expected: `<img src="giraffe.gif"/>`,
+			expected: `<img src="https://proxy.example.com/?u=giraffe.gif"/>`,
 		},
 		{
 			in:       `<img src="giraffe.gif?height=500&amp;width=500&amp;flag" />`,
-			expected: `<img src="giraffe.gif?height=500&amp;width=500&amp;flag"/>`,
+			expected: `<img src="https://proxy.example.com/?u=giraffe.gif?height=500&amp;width=500&amp;flag"/>`,
 		},
 	}
 
 	p := UGCPolicy()
 	p.RequireParseableURLs(true)
+	p.RewriteSrc(func(u *url.URL) {
+		// Proxify all requests to "https://proxy.example.com/?u=http://example.com/"
+		// This is a contrived example, but it shows how to rewrite URLs
+		// to proxy all requests through a single URL.
+
+		url := u.String()
+		u.Scheme = "https"
+		u.Host = "proxy.example.com"
+		u.Path = "/"
+		u.RawQuery = "u=" + url
+	})
 
 	// These tests are run concurrently to enable the race detector to pick up
 	// potential issues