2 files changed, 26 insertions, 16 deletions

diff --git a/plugins/viewhtml/sanitize.go b/plugins/viewhtml/sanitize.go
index 8cee481..d5c94a7 100644
--- a/plugins/viewhtml/sanitize.go
+++ b/plugins/viewhtml/sanitize.go
@@ -6,6 +6,7 @@ import (
 	"regexp"
 	"strings"
 
+	koushinbase "git.sr.ht/~emersion/koushin/plugins/base"
 	"github.com/aymerick/douceur/css"
 	cssparser "github.com/chris-ramon/douceur/parser"
 	"github.com/microcosm-cc/bluemonday"
@@ -68,7 +69,15 @@ var allowedStyles = map[string]bool{
 	"list-style-position": true,
 }
 
-func sanitizeCSSDecls(decls []*css.Declaration) []*css.Declaration {
+type sanitizer struct {
+	msg *koushinbase.IMAPMessage
+}
+
+func (san *sanitizer) sanitizeResourceURL(src string) string {
+	return "about:blank"
+}
+
+func (san *sanitizer) sanitizeCSSDecls(decls []*css.Declaration) []*css.Declaration {
 	sanitized := make([]*css.Declaration, 0, len(decls))
 	for _, decl := range decls {
 		if !allowedStyles[decl.Property] {
@@ -86,26 +95,26 @@ func sanitizeCSSDecls(decls []*css.Declaration) []*css.Declaration {
 	return sanitized
 }
 
-func sanitizeCSSRule(rule *css.Rule) {
+func (san *sanitizer) sanitizeCSSRule(rule *css.Rule) {
 	// Disallow @import
 	if rule.Kind == css.AtRule && strings.EqualFold(rule.Name, "@import") {
 		rule.Prelude = "url(about:blank)"
 	}
 
-	rule.Declarations = sanitizeCSSDecls(rule.Declarations)
+	rule.Declarations = san.sanitizeCSSDecls(rule.Declarations)
 
 	for _, child := range rule.Rules {
-		sanitizeCSSRule(child)
+		san.sanitizeCSSRule(child)
 	}
 }
 
-func sanitizeNode(n *html.Node) {
+func (san *sanitizer) sanitizeNode(n *html.Node) {
 	if n.Type == html.ElementNode {
 		if strings.EqualFold(n.Data, "img") {
 			for i := range n.Attr {
 				attr := &n.Attr[i]
 				if strings.EqualFold(attr.Key, "src") {
-					attr.Val = "about:blank"
+					attr.Val = san.sanitizeResourceURL(attr.Val)
 				}
 			}
 		} else if strings.EqualFold(n.Data, "style") {
@@ -126,7 +135,7 @@ func sanitizeNode(n *html.Node) {
 				s = ""
 			} else {
 				for _, rule := range stylesheet.Rules {
-					sanitizeCSSRule(rule)
+					san.sanitizeCSSRule(rule)
 				}
 
 				s = stylesheet.String()
@@ -149,7 +158,7 @@ func sanitizeNode(n *html.Node) {
 					continue
 				}
 
-				decls = sanitizeCSSDecls(decls)
+				decls = san.sanitizeCSSDecls(decls)
 
 				attr.Val = ""
 				for _, d := range decls {
@@ -160,17 +169,17 @@ func sanitizeNode(n *html.Node) {
 	}
 
 	for c := n.FirstChild; c != nil; c = c.NextSibling {
-		sanitizeNode(c)
+		san.sanitizeNode(c)
 	}
 }
 
-func sanitizeHTML(b []byte) ([]byte, error) {
+func (san *sanitizer) sanitizeHTML(b []byte) ([]byte, error) {
 	doc, err := html.Parse(bytes.NewReader(b))
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse HTML: %v", err)
 	}
 
-	sanitizeNode(doc)
+	san.sanitizeNode(doc)
 
 	var buf bytes.Buffer
 	if err := html.Render(&buf, doc); err != nil {
diff --git a/plugins/viewhtml/viewer.go b/plugins/viewhtml/viewer.go
index 9734c37..47f5eea 100644
--- a/plugins/viewhtml/viewer.go
+++ b/plugins/viewhtml/viewer.go
@@ -12,7 +12,7 @@ import (
 	"github.com/emersion/go-message"
 )
 
-const tpl = `
+const tplSrc = `
 <!-- allow-same-origin is required to resize the frame with its content -->
 <!-- allow-popups is required for target="_blank" links -->
 <iframe id="email-frame" srcdoc="{{.}}" sandbox="allow-same-origin allow-popups"></iframe>
@@ -20,6 +20,8 @@ const tpl = `
 <link rel="stylesheet" href="/plugins/viewhtml/assets/style.css">
 `
 
+var tpl = template.Must(template.New("view-html.html").Parse(tplSrc))
+
 type viewer struct{}
 
 func (viewer) ViewMessagePart(ctx *koushin.Context, msg *koushinbase.IMAPMessage, part *message.Entity) (interface{}, error) {
@@ -36,15 +38,14 @@ func (viewer) ViewMessagePart(ctx *koushin.Context, msg *koushinbase.IMAPMessage
 		return nil, fmt.Errorf("failed to read part body: %v", err)
 	}
 
-	body, err = sanitizeHTML(body)
+	san := sanitizer{msg}
+	body, err = san.sanitizeHTML(body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to sanitize HTML part: %v", err)
 	}
 
-	t := template.Must(template.New("view-html.html").Parse(tpl))
-
 	var buf bytes.Buffer
-	err = t.Execute(&buf, string(body))
+	err = tpl.Execute(&buf, string(body))
 	if err != nil {
 		return nil, err
 	}