From b3a778e9e8570f3b0d9fa126c06b984c1b2ba709 Mon Sep 17 00:00:00 2001
From: Fumiaki Kinoshita <fumiaki.kinoshita@herp.co.jp>
Date: Thu, 22 Jun 2023 14:14:24 +0900
Subject: [PATCH] extract PrefixList from the document

---
 src/Network/Wai/SAML2/C14N.hs       |  6 ++++--
 src/Network/Wai/SAML2/Response.hs   | 10 ++++++++++
 src/Network/Wai/SAML2/Validation.hs |  5 +++--
 src/Network/Wai/SAML2/XML.hs        |  7 +++++++
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/Network/Wai/SAML2/C14N.hs b/src/Network/Wai/SAML2/C14N.hs
index 5d249cb..5ca9bfe 100644
--- a/src/Network/Wai/SAML2/C14N.hs
+++ b/src/Network/Wai/SAML2/C14N.hs
@@ -14,6 +14,8 @@ module Network.Wai.SAML2.C14N (
 --------------------------------------------------------------------------------
 
 import qualified Data.ByteString as BS
+import Data.Text (Text)
+import qualified Data.Text.Encoding as T
 
 import Foreign.C.Types
 
@@ -22,8 +24,8 @@ import Text.XML.C14N
 --------------------------------------------------------------------------------
 
 -- | 'canonicalise' @xml@ produces a canonical representation of @xml@.
-canonicalise :: BS.ByteString -> IO BS.ByteString
-canonicalise xml = c14n c14nOpts c14n_exclusive_1_0 ["xs"] False Nothing xml
+canonicalise :: [Text] -> BS.ByteString -> IO BS.ByteString
+canonicalise prefixList xml = c14n c14nOpts c14n_exclusive_1_0 (map T.encodeUtf8 prefixList) False Nothing xml
 
 -- | The options we want to use for canonicalisation of XML documents.
 c14nOpts :: [CInt]
diff --git a/src/Network/Wai/SAML2/Response.hs b/src/Network/Wai/SAML2/Response.hs
index c9a197e..6bb3fa2 100644
--- a/src/Network/Wai/SAML2/Response.hs
+++ b/src/Network/Wai/SAML2/Response.hs
@@ -11,6 +11,7 @@ module Network.Wai.SAML2.Response (
     Response(..),
     removeSignature,
     extractSignedInfo,
+    extractPrefixList,
 
     -- * Re-exports
     module Network.Wai.SAML2.StatusCode,
@@ -133,6 +134,15 @@ extractSignedInfo cursor = do
                           ) >>= nodes
     pure signedInfo
 
+extractPrefixList :: Cursor -> [T.Text]
+extractPrefixList cursor = concatMap T.words
+    $ concatMap (attribute "PrefixList")
+    $ cursor
+    $/ element (dsName "Reference")
+    &/ element (dsName "Transforms")
+    &/ element (dsName "Transform")
+    &/ element (ecName "InclusiveNamespaces")
+
 --------------------------------------------------------------------------------
 
 -- Reference [StatusResponseType]
diff --git a/src/Network/Wai/SAML2/Validation.hs b/src/Network/Wai/SAML2/Validation.hs
index aaa4442..870b8e0 100644
--- a/src/Network/Wai/SAML2/Validation.hs
+++ b/src/Network/Wai/SAML2/Validation.hs
@@ -130,8 +130,9 @@ validateSAMLResponse cfg responseXmlDoc samlResponse now = do
     let signedInfoXml = XML.renderLBS def doc
 
     -- canonicalise the textual representation of the SignedInfo element
+    let prefixList = extractPrefixList (XML.fromDocument doc)
     signedInfoCanonResult <- liftIO $ try $
-        canonicalise (LBS.toStrict signedInfoXml)
+        canonicalise prefixList (LBS.toStrict signedInfoXml)
 
     normalisedSignedInfo <- case signedInfoCanonResult of
         Left err -> throwError $ CanonicalisationFailure err
@@ -161,7 +162,7 @@ validateSAMLResponse cfg responseXmlDoc samlResponse now = do
 
     -- then render the resulting document and canonicalise it
     let renderedXml = XML.renderLBS def docMinusSignature
-    refCanonResult <- liftIO $ try $ canonicalise (LBS.toStrict renderedXml)
+    refCanonResult <- liftIO $ try $ canonicalise prefixList (LBS.toStrict renderedXml)
 
     normalised <- case refCanonResult of
         Left err -> throwError $ CanonicalisationFailure err
diff --git a/src/Network/Wai/SAML2/XML.hs b/src/Network/Wai/SAML2/XML.hs
index 6694489..a45510b 100644
--- a/src/Network/Wai/SAML2/XML.hs
+++ b/src/Network/Wai/SAML2/XML.hs
@@ -13,6 +13,7 @@ module Network.Wai.SAML2.XML (
     xencName,
     dsName,
     mdName,
+    ecName,
 
     -- * Utility functions
     toMaybeText,
@@ -67,6 +68,12 @@ mdName name =
     Name name (Just "urn:oasis:names:tc:SAML:2.0:metadata") (Just "md")
 
 
+-- | 'ecName' @name@ constructs a 'Name' for @name@ in the
+-- http://www.w3.org/2001/10/xml-exc-c14n# namespace.
+ecName :: T.Text -> Name
+ecName name =
+    Name name (Just "http://www.w3.org/2001/10/xml-exc-c14n#") (Just "ec")
+
 -- | 'toMaybeText' @xs@ returns 'Nothing' if @xs@ is the empty list, or
 -- the result of concatenating @xs@ wrapped in 'Just' otherwise.
 toMaybeText :: [T.Text] -> Maybe T.Text