From f9b7593e6588d77a12c6169932aa061b86bd4f97 Mon Sep 17 00:00:00 2001
From: Martin Probst <martin@probst.io>
Date: Thu, 1 May 2014 20:55:53 +0200
Subject: [PATCH] fix: Handle all different token types that the parser can
 emit (d'oh).

---
 inline_test.go | 10 ++++++++++
 sanitize.go    |  9 ++++++++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/inline_test.go b/inline_test.go
index b7fe3ec..3f9da80 100644
--- a/inline_test.go
+++ b/inline_test.go
@@ -201,6 +201,16 @@ func TestRawHtmlTag(t *testing.T) {
 		"<iframe src=http://ha.ckers.org/scriptlet.html <",
 		// The hyperlink gets linkified, the <iframe> gets escaped
 		"<p>&lt;iframe src=<a href=\"http://ha.ckers.org/scriptlet.html\">http://ha.ckers.org/scriptlet.html</a> &lt;</p>\n",
+
+		// Additonal token types: SelfClosing, Comment, DocType.
+		"<br/>",
+		"<p><br></p>\n",
+
+		"<!-- Comment -->",
+		"<!-- Comment -->\n",
+
+		"<!DOCTYPE test>",
+		"<p>&lt;!DOCTYPE test&gt;</p>\n",
 	}
 	doTestsInlineParam(t, tests, 0, HTML_SKIP_STYLE|HTML_SANITIZE_OUTPUT)
 }
diff --git a/sanitize.go b/sanitize.go
index e8fd13c..c6cb622 100644
--- a/sanitize.go
+++ b/sanitize.go
@@ -64,7 +64,7 @@ func sanitizeHtmlSafe(input []byte) []byte {
 		case html.TextToken:
 			// Text is written escaped.
 			wr.WriteString(tokenizer.Token().String())
-		case html.StartTagToken:
+		case html.SelfClosingTagToken, html.StartTagToken:
 			// HTML tags are escaped unless whitelisted.
 			tag, hasAttributes := tokenizer.TagName()
 			tagName := string(tag)
@@ -105,7 +105,14 @@ func sanitizeHtmlSafe(input []byte) []byte {
 			} else {
 				wr.WriteString(html.EscapeString(string(tokenizer.Raw())))
 			}
+		case html.CommentToken:
+			// Comments are not really expected, but harmless.
+			wr.Write(tokenizer.Raw())
+		case html.DoctypeToken:
+			// Escape DOCTYPES, entities etc can be dangerous
+			wr.WriteString(html.EscapeString(string(tokenizer.Raw())))
 		default:
+			tokenizer.Token()
 			panic(fmt.Errorf("Unexpected token type %v", t))
 		}
 	}