Fix request and response checking in MoreToCome processing

Addresses mongodb-haskell#153
darycabrera · Nov 19, 2024 · 5825c90 · 5825c90
1 parent c083ead
commit 5825c90
Showing 1 changed file with 23 additions and 24 deletions.
diff --git a/Database/MongoDB/Internal/Protocol.hs b/Database/MongoDB/Internal/Protocol.hs
@@ -287,9 +287,11 @@ callOpMsg pipe request flagBit params = do
         (rt, r) ->
           case r of
               ReplyOpMsg{..} ->
-                if flagBits == [MoreToCome]
-                  then yieldResponses .| foldlC mergeResponses p
-                  else return $ (rt, check reqId p)
+                if rt /= reqId then
+                  error $ "expected response id (" ++ show rt ++ ") to match request id (" ++ show reqId ++ ")"
+                else
+                  if flagBits == [MoreToCome] then yieldResponses .| foldlC mergeResponses p
+                  else return (rt, r)
               _ -> error "Impossible" -- see comment above
     yieldResponses = repeatWhileMC
           (do
@@ -298,27 +300,24 @@ callOpMsg pipe request flagBit params = do
              readMVar var >>= either throwIO return :: IO Response
           )
           checkFlagBit
-    mergeResponses p@(rt,rep) p' =
-      case (p, p') of
-          ((_, r), (_, r')) ->
-            case (r, r') of
-                (ReplyOpMsg _ sec _, ReplyOpMsg _ sec' _) -> do
-                    let (section, section') = (head sec, head sec')
-                        (cur, cur') = (maybe Nothing cast $ look "cursor" section,
-                                      maybe Nothing cast $ look "cursor" section')
-                    case (cur, cur') of
-                      (Just doc, Just doc') -> do
-                        let (docs, docs') =
-                              ( fromJust $ cast $ valueAt "nextBatch" doc :: [Document]
-                              , fromJust $ cast $ valueAt "nextBatch" doc' :: [Document])
-                            id' = fromJust $ cast $ valueAt "id" doc' :: Int32
-                        (rt, check id' (rt, rep{ sections = docs' ++ docs })) -- todo: avoid (++)
-                        -- Since we use this to process moreToCome messages, we
-                        -- know that there will be a nextBatch key in the document
-                      _ ->  error "Impossible"
-                _ -> error "Impossible" -- see comment above
-    check requestId (responseTo, reply) = if requestId == responseTo then reply else
-        error $ "expected response id (" ++ show responseTo ++ ") to match request id (" ++ show requestId ++ ")"
+    mergeResponses (rt, rep) (_, rep') =
+      case (rep, rep') of
+        (ReplyOpMsg _ sec _, ReplyOpMsg _ sec' _) -> do
+            let (section, section') = (head sec, head sec')
+                (cur, cur') = ( cast =<< look "cursor" section
+                              , cast =<< look "cursor" section'
+                              )
+            case (cur, cur') of
+              (Just doc, Just doc') -> do
+                let (docs, docs') =
+                      ( fromJust $ cast $ valueAt "nextBatch" doc :: [Document]
+                      , fromJust $ cast $ valueAt "nextBatch" doc' :: [Document]
+                      )
+                (rt, rep{ sections = docs' ++ docs }) -- todo: avoid (++)
+                -- Since we use this to process moreToCome messages, we
+                -- know that there will be a nextBatch key in the document
+              _ ->  error "Impossible"
+        _ -> error "Impossible" -- see comment above
 
 -- * Message