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
-    let hmac = HMAC.hmac SHA1.hash 64
+authSCRAMWith algo un pw = toAuthResult $ do
+    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)]
-    server1 <- runCommand client1
+    let 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
-        let serverPayload1 = B64.decodeLenient . B.pack . at "payload" $ server1
-        let serverData1 = parseSCRAM serverPayload1
-        let iterations = read . B.unpack $ Map.findWithDefault "1" "i" serverData1
-        let salt = B64.decodeLenient $ Map.findWithDefault "" "s" serverData1
-        let snonce = Map.findWithDefault "" "r" serverData1
+    shortcircuit (true1 "ok" server1) (show server1)
+    let serverPayload1 = B64.decodeLenient . B.pack . at "payload" $ server1
+    let serverData1 = parseSCRAM serverPayload1
+    let iterations = read . B.unpack $ Map.findWithDefault "1" "i" serverData1
+    let salt = B64.decodeLenient $ Map.findWithDefault "" "s" serverData1
+    let snonce = Map.findWithDefault "" "r" serverData1
 
-        shortcircuit (B.isInfixOf nonce snonce) $ do
-            let withoutProof = B.concat [B.pack "c=biws,r=", snonce]
-            let digestS = B.pack $ T.unpack un ++ ":mongo:" ++ T.unpack pw
-            let digest = B16.encode $ MD5.hash digestS
-            let saltedPass = scramHI digest salt iterations
-            let clientKey = hmac saltedPass (B.pack "Client Key")
-            let storedKey = SHA1.hash clientKey
-            let authMsg = B.concat [firstBare, B.pack ",", serverPayload1, B.pack ",", withoutProof]
-            let clientSig = hmac storedKey authMsg
-            let pval = B64.encode . BS.pack $ BS.zipWith xor clientKey clientSig
-            let clientFinal = B.concat [withoutProof, B.pack ",p=", pval]
-            let serverKey = hmac saltedPass (B.pack "Server Key")
-            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 (B.isInfixOf nonce snonce) "nonce"
+    let withoutProof = B.concat [B.pack "c=biws,r=", snonce]
+    let digest = case algo of
+          SHA1 -> B16.encode $ MD5.hash $ B.pack $ T.unpack un ++ ":mongo:" ++ T.unpack pw
+          SHA256 -> B.pack $ T.unpack $ saslprep pw
+    let saltedPass = scramHI algo digest salt iterations
+    let clientKey = hmac saltedPass (B.pack "Client Key")
+    let storedKey = hash algo clientKey
+    let authMsg = B.concat [firstBare, B.pack ",", serverPayload1, B.pack ",", withoutProof]
+    let clientSig = hmac storedKey authMsg
+    let pval = B64.encode . BS.pack $ BS.zipWith xor clientKey clientSig
+    let clientFinal = B.concat [withoutProof, B.pack ",p=", pval]
 
-            shortcircuit (true1 "ok" server2) $ do
-                let serverPayload2 = B64.decodeLenient . B.pack $ at "payload" server2
-                let serverData2 = parseSCRAM serverPayload2
-                let serverSigComp = Map.findWithDefault "" "v" serverData2
+    let client2 =
+            [ "saslContinue" =: (1 :: Int)
+            , "conversationId" =: (at "conversationId" server1 :: Int)
+            , "payload" =: B.unpack (B64.encode clientFinal)
+            ]
+    server2 <- lift $ runCommand client2
+    shortcircuit (true1 "ok" server2) (show server2)
 
-                shortcircuit (serverSig == serverSigComp) $ do
-                  let done = true1 "done" server2
-                  if done
-                    then return True
-                    else do
-                      let client2Step2 = [ "saslContinue" =: (1 :: Int)
-                                         , "conversationId" =: (at "conversationId" server1 :: Int)
-                                         , "payload" =: String ""]
-                      server3 <- runCommand client2Step2
-                      shortcircuit (true1 "ok" server3) $ do
-                        return True
+    let serverKey = hmac saltedPass (B.pack "Server Key")
+    let serverSig = B64.encode $ hmac serverKey authMsg
+    let serverPayload2 = B64.decodeLenient . B.pack $ at "payload" server2
+    let serverData2 = parseSCRAM serverPayload2
+    let serverSigComp = Map.findWithDefault "" "v" serverData2
+
+    shortcircuit (serverSig == serverSigComp) "server signature does not match"
+    if true1 "done" server2
+      then return ()
+      else do
+        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
-    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
+    hmacd = HMAC.hmac (hash algo) 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"]