TokenReader for pre-sanitization transformations

sanitize.go

@@ -51,12 +51,12 @@ var (
 // It returns a HTML string that has been sanitized by the policy or an empty
 // string if an error has occurred (most likely as a consequence of extremely
 // malformed input)
-func (p *Policy) Sanitize(s string) string {
+func (p *Policy) Sanitize(s string, filters ...TokenReader) string {
 	if strings.TrimSpace(s) == "" {
 		return s
 	}
 
-	return p.sanitize(strings.NewReader(s)).String()
+	return p.sanitize(strings.NewReader(s), filters...).String()
 }
 
 // SanitizeBytes takes a []byte that contains a HTML fragment or document and applies
@@ -65,25 +65,25 @@ func (p *Policy) Sanitize(s string) string {
 // It returns a []byte containing the HTML that has been sanitized by the policy
 // or an empty []byte if an error has occurred (most likely as a consequence of
 // extremely malformed input)
-func (p *Policy) SanitizeBytes(b []byte) []byte {
+func (p *Policy) SanitizeBytes(b []byte, filters ...TokenReader) []byte {
 	if len(bytes.TrimSpace(b)) == 0 {
 		return b
 	}
 
-	return p.sanitize(bytes.NewReader(b)).Bytes()
+	return p.sanitize(bytes.NewReader(b), filters...).Bytes()
 }
 
 // SanitizeReader takes an io.Reader that contains a HTML fragment or document
 // and applies the given policy whitelist.
 //
 // It returns a bytes.Buffer containing the HTML that has been sanitized by the
 // policy. Errors during sanitization will merely return an empty result.
-func (p *Policy) SanitizeReader(r io.Reader) *bytes.Buffer {
-	return p.sanitize(r)
+func (p *Policy) SanitizeReader(r io.Reader, filters ...TokenReader) *bytes.Buffer {
+	return p.sanitize(r, filters...)
 }
 
 // Performs the actual sanitization process.
-func (p *Policy) sanitize(r io.Reader) *bytes.Buffer {
+func (p *Policy) sanitize(r io.Reader, filters ...TokenReader) *bytes.Buffer {
 
 	// It is possible that the developer has created the policy via:
 	//   p := bluemonday.Policy{}
@@ -100,22 +100,31 @@ func (p *Policy) sanitize(r io.Reader) *bytes.Buffer {
 		skipClosingTag           bool
 		closingTagToSkipStack    []string
 		mostRecentlyStartedToken string
+		reader                   TokenReader
 	)
 
-	tokenizer := html.NewTokenizer(r)
+	// Chain together TokenReader filters
+	reader = &tokenizerReader{html.NewTokenizer(r)}
+	for _, f := range filters {
+		f.Source(reader)
+		reader = f
+	}
+
 	for {
-		if tokenizer.Next() == html.ErrorToken {
-			err := tokenizer.Err()
+		token, err := reader.Token()
+		if token == nil {
 			if err == io.EOF {
 				// End of input means end of processing
 				return &buff
 			}
-
+			if err == nil {
+				// Skip nil token
+				continue
+			}
 			// Raw tokenizer error
 			return &bytes.Buffer{}
 		}
 
-		token := tokenizer.Token()
 		switch token.Type {
 		case html.DoctypeToken:
 

token.go

@@ -0,0 +1,55 @@
+// Copyright (c) 2014, David Kitchen <david@buro9.com>
+//
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice, this
+//   list of conditions and the following disclaimer.
+//
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// * Neither the name of the organisation (Microcosm) nor the names of its
+//   contributors may be used to endorse or promote products derived from
+//   this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+package bluemonday
+
+import "golang.org/x/net/html"
+
+type TokenReader interface {
+	Source(source TokenReader)
+	Token() (*html.Token, error)
+}
+
+type tokenizerReader struct {
+	*html.Tokenizer
+}
+
+var _ TokenReader = &tokenizerReader{}
+
+func (r *tokenizerReader) Token() (*html.Token, error) {
+	if r.Next() == html.ErrorToken {
+		return nil, r.Err()
+	}
+	token := r.Tokenizer.Token()
+	return &token, nil
+}
+
+// Source is a no-op for tokenizerReader
+func (r *tokenizerReader) Source(TokenReader) {
+}

token_test.go

@@ -0,0 +1,48 @@
+package bluemonday
+
+import (
+	"testing"
+
+	"golang.org/x/net/html"
+	"golang.org/x/net/html/atom"
+)
+
+type testRemoverReader struct {
+	source  TokenReader
+	tagAtom atom.Atom
+}
+
+func (r *testRemoverReader) Token() (*html.Token, error) {
+	t, err := r.source.Token()
+	if err != nil {
+		return t, err
+	}
+	if (t.Type == html.StartTagToken || t.Type == html.EndTagToken) && t.DataAtom == r.tagAtom {
+		// Skip bold, return next token
+		return r.source.Token()
+	}
+	return t, nil
+}
+
+func (r *testRemoverReader) Source(s TokenReader) {
+	r.source = s
+}
+
+func TestTokenReader(t *testing.T) {
+	p := UGCPolicy()
+
+	input := "<p><b>A bold statement.</b></p>"
+	want := "<p><b>A bold statement.</b></p>"
+	got := p.Sanitize(input)
+	if got != want {
+		t.Errorf("got: %q, want: %q", got, want)
+	}
+
+	removeBold := &testRemoverReader{tagAtom: atom.B}
+	input = "<p><b>A bold statement.</b></p>"
+	want = "<p>A bold statement.</p>"
+	got = p.Sanitize(input, removeBold)
+	if got != want {
+		t.Errorf("got: %q, want: %q", got, want)
+	}
+}