From 99a0388dc240b561ab2d779fad61280258e5248e Mon Sep 17 00:00:00 2001
From: Michael Snoyman <michael@snoyman.com>
Date: Sun, 31 Oct 2010 23:03:40 +0200
Subject: [PATCH] Ensure balanced tags

---
 Text/HTML/SanitizeXSS.hs | 36 +++++++++++++++++++++++++++---------
 test.hs                  |  4 ++--
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/Text/HTML/SanitizeXSS.hs b/Text/HTML/SanitizeXSS.hs
index 3fb7c25..0ad3e0b 100644
--- a/Text/HTML/SanitizeXSS.hs
+++ b/Text/HTML/SanitizeXSS.hs
@@ -9,20 +9,38 @@ import Network.URI ( parseURIReference, URI (..),
                      isAllowedInURI, escapeURIString, uriScheme )
 import Codec.Binary.UTF8.String ( encodeString )
 
+import qualified Data.Map as Map
+
 -- | santize the html to prevent XSS attacks. See README.md <http://github.com/gregwebs/haskell-xss-sanitize> for more details
 sanitizeXSS :: String -> String
 sanitizeXSS = renderTagsOptions renderOptions {
     optMinimize = \x -> x `elem` ["br","img"] -- <img><img> converts to <img />, <a/> converts to <a></a>
-  } .  safeTags . parseTags
+  } .  safeTags Map.empty . parseTags
   where
-    safeTags :: [Tag String] -> [Tag String]
-    safeTags [] = []
-    safeTags (t@(TagClose name):tags) | safeTagName name = t:(safeTags tags)
-                                      | otherwise = safeTags tags
-    safeTags (TagOpen name attributes:tags)
-      | safeTagName name = TagOpen name (filter safeAttribute attributes) : safeTags tags
-      | otherwise = safeTags tags
-    safeTags (t:tags) = t:safeTags tags
+    safeTags :: Map.Map String Int -> [Tag String] -> [Tag String]
+    safeTags m [] =
+        concatMap go $ Map.toList m
+      where
+        go (name, i) = replicate i $ TagClose name
+    safeTags m (t@(TagClose name):tags)
+        | safeTagName name =
+            case Map.lookup name m of
+                Nothing -> safeTags m tags
+                Just i ->
+                    let m' = if i == 1
+                                then Map.delete name m
+                                else Map.insert name (i - 1) m
+                     in t : safeTags m' tags
+        | otherwise = safeTags m tags
+    safeTags m (TagOpen name attributes:tags)
+      | safeTagName name =
+        let m' =
+                case Map.lookup name m of
+                    Nothing -> Map.insert name 1 m
+                    Just i -> Map.insert name (i + 1) m
+         in TagOpen name (filter safeAttribute attributes) : safeTags m' tags
+      | otherwise = safeTags m tags
+    safeTags m (t:tags) = t:safeTags m tags
 
 safeTagName :: String -> Bool
 safeTagName tagname = tagname `member` sanitaryTags
diff --git a/test.hs b/test.hs
index 44db219..840f59b 100644
--- a/test.hs
+++ b/test.hs
@@ -1,8 +1,8 @@
 import Text.HTML.SanitizeXSS
 
 main = do
-  let test = " <a href='http://safe.com'>safe</a><a href='unsafe://hack.com'>anchor</a> <img src='evil://evil.com' /> <unsafe></foo> <bar /> <br></br> "
+  let test = " <a href='http://safe.com'>safe</a><a href='unsafe://hack.com'>anchor</a> <img src='evil://evil.com' /> <unsafe></foo> <bar /> <br></br> <b>Unbalanced</div>"
   let actual = (sanitizeXSS test)
-  let expected = " <a href=\"http://safe.com\">safe</a><a>anchor</a> <img />   <br /> "
+  let expected = " <a href=\"http://safe.com\">safe</a><a>anchor</a> <img />   <br /> <b>Unbalanced</b>"
   putStrLn $ "testing: " ++ test
   putStrLn $ if actual == expected then "pass" else "failure\n" ++ "\nexpected:" ++ (show expected) ++ "\nactual:  " ++ (show actual)