Support SCRAM-SHA-256

780df80cfc introduces support for the
OP_MSG protocol. Unfortunately, the upsert and multi options of the
update command still use flagBits to communicate the options, whereas
they must be provided directly into the command document,
alongside the "q" and "v" fields.

This commit:
 - introduces a test for a single-document upsert that, if isolated,
   succeeds against the reference MongoDB 3.6 container, but fails
   against an official 6.0 image.
 - provides a patch that sets the appropriate options.

The test is not perfect as the upsert operation is inherently racy and
this difficult to test. A comfortable threadDelay has been inserted as
a workaround to accomodate for medium workloads.

Database/MongoDB/Internal/Protocol.hs

@@ -494,6 +494,9 @@ data FlagBit =
     | ExhaustAllowed  -- ^ The client is prepared for multiple replies to this request using the moreToCome bit.
     deriving (Show, Eq, Enum)
 
+uOptDoc :: UpdateOption -> Document
+uOptDoc Upsert = ["upsert" =: True]
+uOptDoc MultiUpdate = ["multi" =: True]
 
 {-
   OP_MSG header == 16 byte
@@ -528,7 +531,7 @@ putOpMsg cmd requestId flagBit params = do
                 putCString "documents"               -- identifier
                 mapM_ putDocument iDocuments         -- payload
             Update{..} -> do
-                let doc = ["q" =: uSelector, "u" =: uUpdater]
+                let doc = ["q" =: uSelector, "u" =: uUpdater] <> concatMap uOptDoc uOptions
                     (sec0, sec1Size) =
                       prepSectionInfo
                           uFullCollection

Database/MongoDB/Query.hs

@@ -11,7 +11,7 @@ module Database.MongoDB.Query (
     -- * Database
     Database, allDatabases, useDb, thisDatabase,
     -- ** Authentication
-    Username, Password, auth, authMongoCR, authSCRAMSHA1,
+    Username, Password, auth, authMongoCR, authSCRAMSHA1, authSCRAMSHA256,
     -- * Collection
     Collection, allCollections,
     -- ** Selection
@@ -61,8 +61,10 @@ import Control.Monad
   )
 import Control.Monad.Reader (MonadReader, ReaderT, ask, asks, local, runReaderT)
 import Control.Monad.Trans (MonadIO, liftIO, lift)
+import Control.Monad.Trans.Except
 import qualified Crypto.Hash.MD5 as MD5
 import qualified Crypto.Hash.SHA1 as SHA1
+import qualified Crypto.Hash.SHA256 as SHA256
 import qualified Crypto.MAC.HMAC as HMAC
 import qualified Crypto.Nonce as Nonce
 import Data.Binary.Put (runPut)
@@ -285,62 +287,93 @@ authMongoCR usr pss = do
     n <- at "nonce" <$> runCommand ["getnonce" =: (1 :: Int)]
     true1 "ok" <$> runCommand ["authenticate" =: (1 :: Int), "user" =: usr, "nonce" =: n, "key" =: pwKey n usr pss]
 
+data HashAlgorithm = SHA1 | SHA256 deriving Show
+
+hash :: HashAlgorithm -> B.ByteString -> B.ByteString
+hash SHA1 = SHA1.hash
+hash SHA256 = SHA256.hash
+
 authSCRAMSHA1 :: MonadIO m => Username -> Password -> Action m Bool
+authSCRAMSHA1 = authSCRAMWith SHA1
+
+authSCRAMSHA256 :: MonadIO m => Username -> Password -> Action m Bool
+authSCRAMSHA256 = authSCRAMWith SHA256
+
+toAuthResult :: Functor m => ExceptT String (Action m) () -> Action m Bool
+toAuthResult = fmap (either (const False) (const True)) . runExceptT
+
+-- | It should technically perform SASLprep, but the implementation is currently id
+saslprep :: Text -> Text
+saslprep = id
+
+authSCRAMWith :: MonadIO m => HashAlgorithm -> Username -> Password -> Action m Bool
 -- ^ Authenticate with the current database, using the SCRAM-SHA-1 authentication mechanism (default in MongoDB server >= 3.0)
-authSCRAMSHA1 un pw = do
+authSCRAMWith algo un pw = toAuthResult $ do
-    let hmac = HMAC.hmac SHA1.hash 64
+    let hmac = HMAC.hmac (hash algo) 64
     nonce <- liftIO (Nonce.withGenerator Nonce.nonce128 <&> B64.encode)
     let firstBare = B.concat [B.pack $ "n=" ++ T.unpack un ++ ",r=", nonce]
-    let client1 = ["saslStart" =: (1 :: Int), "mechanism" =: ("SCRAM-SHA-1" :: String), "payload" =: (B.unpack . B64.encode $ B.concat [B.pack "n,,", firstBare]), "autoAuthorize" =: (1 :: Int)]
+    let client1 =
-    server1 <- runCommand client1
+          [ "saslStart" =: (1 :: Int)
+          , "mechanism" =: case algo of
+            SHA1 -> "SCRAM-SHA-1" :: String
+            SHA256 -> "SCRAM-SHA-256"
+          , "payload" =: (B.unpack . B64.encode $ B.concat [B.pack "n,,", firstBare])
+          , "autoAuthorize" =: (1 :: Int)
+          ]
+    server1 <- lift $ runCommand client1
 
-    shortcircuit (true1 "ok" server1) $ do
+    shortcircuit (true1 "ok" server1) (show server1)
-        let serverPayload1 = B64.decodeLenient . B.pack . at "payload" $ server1
+    let serverPayload1 = B64.decodeLenient . B.pack . at "payload" $ server1
-        let serverData1 = parseSCRAM serverPayload1
+    let serverData1 = parseSCRAM serverPayload1
-        let iterations = read . B.unpack $ Map.findWithDefault "1" "i" serverData1
+    let iterations = read . B.unpack $ Map.findWithDefault "1" "i" serverData1
-        let salt = B64.decodeLenient $ Map.findWithDefault "" "s" serverData1
+    let salt = B64.decodeLenient $ Map.findWithDefault "" "s" serverData1
-        let snonce = Map.findWithDefault "" "r" serverData1
+    let snonce = Map.findWithDefault "" "r" serverData1
 
-        shortcircuit (B.isInfixOf nonce snonce) $ do
+    shortcircuit (B.isInfixOf nonce snonce) "nonce"
-            let withoutProof = B.concat [B.pack "c=biws,r=", snonce]
+    let withoutProof = B.concat [B.pack "c=biws,r=", snonce]
-            let digestS = B.pack $ T.unpack un ++ ":mongo:" ++ T.unpack pw
+    let digest = case algo of
-            let digest = B16.encode $ MD5.hash digestS
+          SHA1 -> B16.encode $ MD5.hash $ B.pack $ T.unpack un ++ ":mongo:" ++ T.unpack pw
-            let saltedPass = scramHI digest salt iterations
+          SHA256 -> B.pack $ T.unpack $ saslprep pw
-            let clientKey = hmac saltedPass (B.pack "Client Key")
+    let saltedPass = scramHI algo digest salt iterations
-            let storedKey = SHA1.hash clientKey
+    let clientKey = hmac saltedPass (B.pack "Client Key")
-            let authMsg = B.concat [firstBare, B.pack ",", serverPayload1, B.pack ",", withoutProof]
+    let storedKey = hash algo clientKey
-            let clientSig = hmac storedKey authMsg
+    let authMsg = B.concat [firstBare, B.pack ",", serverPayload1, B.pack ",", withoutProof]
-            let pval = B64.encode . BS.pack $ BS.zipWith xor clientKey clientSig
+    let clientSig = hmac storedKey authMsg
-            let clientFinal = B.concat [withoutProof, B.pack ",p=", pval]
+    let pval = B64.encode . BS.pack $ BS.zipWith xor clientKey clientSig
-            let serverKey = hmac saltedPass (B.pack "Server Key")
+    let clientFinal = B.concat [withoutProof, B.pack ",p=", pval]
-            let serverSig = B64.encode $ hmac serverKey authMsg
-            let client2 = ["saslContinue" =: (1 :: Int), "conversationId" =: (at "conversationId" server1 :: Int), "payload" =: B.unpack (B64.encode clientFinal)]
-            server2 <- runCommand client2
 
-            shortcircuit (true1 "ok" server2) $ do
+    let client2 =
-                let serverPayload2 = B64.decodeLenient . B.pack $ at "payload" server2
+            [ "saslContinue" =: (1 :: Int)
-                let serverData2 = parseSCRAM serverPayload2
+            , "conversationId" =: (at "conversationId" server1 :: Int)
-                let serverSigComp = Map.findWithDefault "" "v" serverData2
+            , "payload" =: B.unpack (B64.encode clientFinal)
+            ]
+    server2 <- lift $ runCommand client2
+    shortcircuit (true1 "ok" server2) (show server2)
 
-                shortcircuit (serverSig == serverSigComp) $ do
+    let serverKey = hmac saltedPass (B.pack "Server Key")
-                  let done = true1 "done" server2
+    let serverSig = B64.encode $ hmac serverKey authMsg
-                  if done
+    let serverPayload2 = B64.decodeLenient . B.pack $ at "payload" server2
-                    then return True
+    let serverData2 = parseSCRAM serverPayload2
-                    else do
+    let serverSigComp = Map.findWithDefault "" "v" serverData2
-                      let client2Step2 = [ "saslContinue" =: (1 :: Int)
+
-                                         , "conversationId" =: (at "conversationId" server1 :: Int)
+    shortcircuit (serverSig == serverSigComp) "server signature does not match"
-                                         , "payload" =: String ""]
+    if true1 "done" server2
-                      server3 <- runCommand client2Step2
+      then return ()
-                      shortcircuit (true1 "ok" server3) $ do
+      else do
-                        return True
+        let client2Step2 = [ "saslContinue" =: (1 :: Int)
+                            , "conversationId" =: (at "conversationId" server1 :: Int)
+                            , "payload" =: String ""]
+        server3 <- lift $ runCommand client2Step2
+        shortcircuit (true1 "ok" server3) "server3"
+
+shortcircuit :: Monad m => Bool -> String -> ExceptT String m ()
+shortcircuit True _ = pure ()
+shortcircuit False reason = throwE (show reason)
+
+scramHI :: HashAlgorithm -> B.ByteString -> B.ByteString -> Int -> B.ByteString
+scramHI algo digest salt iters = snd $ foldl com (u1, u1) [1..(iters-1)]
     where
-    shortcircuit True f = f
+    hmacd = HMAC.hmac (hash algo) 64 digest
-    shortcircuit False _ = return False
-
-scramHI :: B.ByteString -> B.ByteString -> Int -> B.ByteString
-scramHI digest salt iters = snd $ foldl com (u1, u1) [1..(iters-1)]
-    where
-    hmacd = HMAC.hmac SHA1.hash 64 digest
     u1 = hmacd (B.concat [salt, BS.pack [0, 0, 0, 1]])
     com (u,uc) _ = let u' = hmacd u in (u', BS.pack $ BS.zipWith xor uc u')
 

mongoDB.cabal

@@ -115,6 +115,8 @@ Benchmark bench
                     , base16-bytestring
                     , binary -any
                     , bson >= 0.3 && < 0.5
+                    , conduit
+                    , conduit-extra
                     , data-default-class -any
                     , text
                     , bytestring -any
@@ -128,6 +130,7 @@ Benchmark bench
                     , random-shuffle -any
                     , monad-control >= 0.3.1
                     , lifted-base >= 0.1.0.3
+                    , transformers
                     , transformers-base >= 0.4.1
                     , hashtables >= 1.1.2.0
                     , fail

test/QuerySpec.hs

@@ -4,6 +4,7 @@
 module QuerySpec (spec) where
 import Data.String (IsString(..))
 import TestImport
+import Control.Concurrent (threadDelay)
 import Control.Exception
 import Control.Monad (forM_, when)
 import System.Environment (getEnv)
@@ -87,6 +88,21 @@ spec = around withCleanDatabase $ do
       db (count $ select ["name" =: "Yankees", "league" =: "American"] "team") `shouldReturn` 1
       _id `shouldBe` ()
 
+  describe "upsert" $ do
+    it "upserts a document twice with the same spec" $ do
+      let q = select ["name" =: "jack"] "users"
+      db $ upsert q ["color" =: "blue", "name" =: "jack"]
+      -- since there is no way to ask for a ack, we must wait for "a sufficient time"
+      -- for the write to be visible
+      threadDelay 10000
+      db (rest =<< find (select [] "users")) >>= print
+      db (count $ select ["name" =: "jack"] "users") `shouldReturn` 1
+      db $ upsert q ["color" =: "red", "name" =: "jack"]
+      threadDelay 10000
+      db (count $ select ["name" =: "jack"] "users") `shouldReturn` 1
+      Just doc <- db $ findOne (select ["name" =: "jack"] "users")
+      doc !? "color" `shouldBe` Just "red"
+
   describe "insertMany" $ do
     it "inserts documents to the collection and returns their _ids" $ do
       (_id1:_id2:_) <- db $ insertMany "team" [ ["name" =: "Yankees", "league" =: "American"]