Merge patch for Mongo 6

Merge pull request #147 from pierreMizrahi/master

CHANGELOG.md

@@ -4,6 +4,11 @@ This project adheres to [Package Versioning Policy](https://wiki.haskell.org/Pac
 
 * Get rid of `MonadFail` constraints in `Database.MongoDB.Query`
 
+## [] - 2023-??-??
+
+### Fixed
+- Collections with dot in the name
+
 ## [2.7.1.2] - 2022-10-26
 
 ### Added

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
@@ -578,6 +581,10 @@ putOpMsg cmd requestId flagBit params = do
                 putInt32 (bit $ bitOpMsg $ ExhaustAllowed)
                 putInt8 0
                 putDocument pre
+            Message{..} -> do
+                putInt32 biT
+                putInt8 0
+                putDocument $ merge [ "$db" =: mDatabase ] mParams
         Kc k -> case k of
             KillC{..} -> do
                 let n = T.splitOn "." kFullCollection
@@ -653,7 +660,11 @@ data Request =
     } | GetMore {
         gFullCollection :: FullCollection,
         gBatchSize :: Int32,
-        gCursorId :: CursorId}
+        gCursorId :: CursorId
+    } | Message {
+        mDatabase :: Text,
+        mParams :: Document
+    }
     deriving (Show, Eq)
 
 data QueryOption =
@@ -673,6 +684,7 @@ data QueryOption =
 qOpcode :: Request -> Opcode
 qOpcode Query{} = 2004
 qOpcode GetMore{} = 2005
+qOpcode Message{} = 2013
 
 opMsgOpcode :: Opcode
 opMsgOpcode = 2013
@@ -693,6 +705,10 @@ putRequest request requestId = do
             putCString gFullCollection
             putInt32 gBatchSize
             putInt64 gCursorId
+        Message{..} -> do
+            putInt32 0
+            putInt8 0
+            putDocument $ merge [ "$db" =: mDatabase ] mParams
 
 qBit :: QueryOption -> Int32
 qBit TailableCursor = bit 1

Database/MongoDB/Internal/Util.hs

@@ -92,6 +92,9 @@ bitOr = foldl (.|.) 0
 -- ^ Concat first and second together with period in between. Eg. @\"hello\" \<.\> \"world\" = \"hello.world\"@
 a <.> b = T.append a (T.cons '.' b)
 
+splitDot :: Text -> (Text, Text)
+splitDot t = let (pre, post) = T.break (== '.') t in (pre, T.drop 1 post) 
+
 true1 :: Label -> Document -> Bool
 -- ^ Is field's value a 1 or True (MongoDB use both Int and Bools for truth values). Error if field not in document or field not a Num or Bool.
 true1 k doc = case valueAt k doc of

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)
@@ -133,7 +135,7 @@ import Database.MongoDB.Internal.Protocol
   )
 import Control.Monad.Trans.Except
 import qualified Database.MongoDB.Internal.Protocol as P
-import Database.MongoDB.Internal.Util (liftIOE, loop, true1, (<.>))
+import Database.MongoDB.Internal.Util (liftIOE, loop, true1, (<.>), splitDot)
 import System.Mem.Weak (Weak)
 import Text.Read (readMaybe)
 import Prelude hiding (lookup)
@@ -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
+    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
+    shortcircuit (B.isInfixOf nonce snonce) "nonce"
     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 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 = SHA1.hash clientKey
+    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]
+
+    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)
+
     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 (true1 "ok" server2) $ do
     let serverPayload2 = B64.decodeLenient . B.pack $ at "payload" server2
     let serverData2 = parseSCRAM serverPayload2
     let serverSigComp = Map.findWithDefault "" "v" serverData2
 
-                shortcircuit (serverSig == serverSigComp) $ do
-                  let done = true1 "done" server2
-                  if done
-                    then return True
+    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 <- runCommand client2Step2
-                      shortcircuit (true1 "ok" server3) $ do
-                        return True
-    where
-    shortcircuit True f = f
-    shortcircuit False _ = return False
+        server3 <- lift $ runCommand client2Step2
+        shortcircuit (true1 "ok" server3) "server3"
 
-scramHI :: B.ByteString -> B.ByteString -> Int -> B.ByteString
-scramHI digest salt iters = snd $ foldl com (u1, u1) [1..(iters-1)]
+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
-    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')
 
@@ -1272,9 +1305,9 @@ find q@Query{selection, batchSize} = do
         qr <- queryRequestOpMsg False q
         let newQr =
               case fst qr of
-                Req qry ->
-                  let coll = last $ T.splitOn "." (qFullCollection qry)
-                  in (Req $ qry {qSelector = merge (qSelector qry) [ "find" =: coll ]}, snd qr)
+                Req P.Query{..} ->
+                  let coll = last $ T.splitOn "." qFullCollection
+                  in (Req $ P.Query {qSelector = merge qSelector [ "find" =: coll ], ..}, snd qr)
                 -- queryRequestOpMsg only returns Cmd types constructed via Req
                 _ -> error "impossible"
         dBatch <- liftIO $ requestOpMsg pipe newQr []
@@ -1312,6 +1345,9 @@ findCommand q@Query{..} = do
         | predicate a = Just (f a)
         | otherwise   = Nothing
 
+isHandshake :: Document -> Bool
+isHandshake = (== ["isMaster" =: (1 :: Int32)])
+
 findOne :: (MonadIO m) => Query -> Action m (Maybe Document)
 -- ^ Fetch first document satisfying query or @Nothing@ if none satisfy it
 findOne q = do
@@ -1321,8 +1357,7 @@ findOne q = do
             rq <- liftIO $ request pipe [] qr
             Batch _ _ docs <- liftDB $ fulfill rq
             return (listToMaybe docs)
-        isHandshake = (== ["isMaster" =: (1 :: Int32)])  $ selector $ selection q :: Bool
-    if isHandshake
+    if isHandshake (selector $ selection q)
       then legacyQuery
       else do
         let sd = P.serverData pipe
@@ -1332,14 +1367,14 @@ findOne q = do
             qr <- queryRequestOpMsg False q {limit = 1}
             let newQr =
                   case fst qr of
-                    Req qry ->
-                      let coll = last $ T.splitOn "." (qFullCollection qry)
+                    Req P.Query{..} ->
+                      let coll = last $ T.splitOn "." qFullCollection
                           -- We have to understand whether findOne is called as
                           -- command directly. This is necessary since findOne is used via
                           -- runCommand as a vehicle to execute any type of commands and notices.
-                          labels = catMaybes $ map (\f -> look f $ qSelector qry) (noticeCommands ++ adminCommands) :: [Value]
+                          labels = catMaybes $ map (\f -> look f qSelector) (noticeCommands ++ adminCommands) :: [Value]
                       in if null labels
-                           then (Req $ qry {qSelector = merge (qSelector qry) [ "find" =: coll ]}, snd qr)
+                           then (Req P.Query {qSelector = merge qSelector [ "find" =: coll ], ..}, snd qr)
                            else qr
                     _ -> error "impossible"
             rq <- liftIO $ requestOpMsg pipe newQr []
@@ -1526,7 +1561,7 @@ requestOpMsg pipe (Req r, remainingLimit) params = do
   promise <- liftIOE ConnectionFailure $ P.callOpMsg pipe r Nothing params
   let protectedPromise = liftIOE ConnectionFailure promise
   return $ fromReply remainingLimit =<< protectedPromise
-requestOpMsg _ (Nc _, _) _ = error "requestOpMsg: Only messages of type Query are supported"
+requestOpMsg _ _ _ = error "requestOpMsg: Only messages of type Query are supported"
 
 fromReply :: Maybe Limit -> Reply -> DelayedBatch
 -- ^ Convert Reply to Batch or Failure
@@ -1844,9 +1879,29 @@ type Command = Document
 -- ^ A command is a special query or action against the database. See <http://www.mongodb.org/display/DOCS/Commands> for details.
 
 runCommand :: (MonadIO m) => Command -> Action m Document
--- ^ Run command against the database and return its result
-runCommand c = fromMaybe err <$> findOne (query c "$cmd") where
-    err = error $ "Nothing returned for command: " ++ show c
+runCommand params = do
+    pipe <- asks mongoPipe
+    if isHandshake params || maxWireVersion (P.serverData pipe) < 17
+      then runCommandLegacy pipe params
+      else runCommand' pipe params
+
+runCommandLegacy :: MonadIO m => Pipe -> Selector -> ReaderT MongoContext m Document
+runCommandLegacy pipe params = do
+    qr <- queryRequest False (query params "$cmd") {limit = 1}
+    rq <- liftIO $ request pipe [] qr
+    Batch _ _ docs <- liftDB $ fulfill rq
+    case docs of
+      [doc] -> pure doc
+      _ -> error $ "Nothing returned for command: " <> show params
+
+runCommand' :: MonadIO m => Pipe -> Selector -> ReaderT MongoContext m Document
+runCommand' pipe params = do
+    ctx <- ask
+    rq <- liftIO $ requestOpMsg pipe ( Req (P.Message (mongoDatabase ctx) params), Just 1) []
+    Batch _ _ docs <- liftDB $ fulfill rq
+    case docs of
+      [doc] -> pure doc
+      _ -> error $ "Nothing returned for command: " <> show params
 
 runCommand1 :: (MonadIO m) => Text -> Action m Document
 -- ^ @runCommand1 foo = runCommand [foo =: 1]@

mongoDB.cabal

@@ -60,15 +60,17 @@ Library
                     , fail
                     , dns
                     , http-types
-
-  if flag(_old-network)
-     -- "Network.BSD" is only available in network < 2.9
-     build-depends: network < 2.9
-  else
-     -- "Network.BSD" has been moved into its own package `network-bsd`
-     build-depends: network >= 3.0
+                    , network >= 3.0
                     , network-bsd >= 2.7 && < 2.9
 
+  --if flag(_old-network)
+     ---- "Network.BSD" is only available in network < 2.9
+     --build-depends: network < 2.9
+  --else
+     ---- "Network.BSD" has been moved into its own package `network-bsd`
+     --build-depends: network >= 3.0
+                    --, network-bsd >= 2.7 && < 2.9
+
   Exposed-modules:  Database.MongoDB
                     Database.MongoDB.Admin
                     Database.MongoDB.Connection
@@ -115,6 +117,8 @@ Benchmark bench
                     , base16-bytestring
                     , binary -any
                     , bson >= 0.3 && < 0.5
+                    , conduit
+                    , conduit-extra
                     , data-default-class -any
                     , text
                     , bytestring -any
@@ -128,6 +132,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
@@ -135,14 +140,16 @@ Benchmark bench
                     , http-types
                     , criterion
                     , tls >= 1.3.0
-
-  if flag(_old-network)
-     -- "Network.BSD" is only available in network < 2.9
-     build-depends: network < 2.9
-  else
-     -- "Network.BSD" has been moved into its own package `network-bsd`
-     build-depends: network >= 3.0
+                    , network >= 3.0
                     , network-bsd >= 2.7 && < 2.9
 
+  --if flag(_old-network)
+     ---- "Network.BSD" is only available in network < 2.9
+     --build-depends: network < 2.9
+  --else
+     ---- "Network.BSD" has been moved into its own package `network-bsd`
+     --build-depends: network >= 3.0
+                    --, network-bsd >= 2.7 && < 2.9
+
   default-language: Haskell2010
   default-extensions: OverloadedStrings

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)
@@ -75,6 +76,14 @@ spec = around withCleanDatabase $ do
       db thisDatabase `shouldReturn` testDBName
       db (useDb anotherDBName thisDatabase) `shouldReturn` anotherDBName
 
+  describe "collectionWithDot" $ do
+    it "uses a collection with dots in the name" $ do
+      let coll = "collection.with.dot"
+      _id <- db $ insert coll ["name" =: "jack", "color" =: "blue"]
+      Just doc <- db $ findOne (select ["name" =: "jack"] coll)
+      doc !? "color" `shouldBe` (Just "blue")
+
+
   describe "insert" $ do
     it "inserts a document to the collection and returns its _id" $ do
       _id <- db $ insert "team" ["name" =: "Yankees", "league" =: "American"]
@@ -87,6 +96,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"]