Make bs2objid safe, add unsafe version under bs2objid'

ozataman · Mar 29, 2011 · c893ab5 · c893ab5
1 parent ce9bf84
commit c893ab5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/src/Snap/Extension/DB/MongoDB.hs b/src/Snap/Extension/DB/MongoDB.hs
@@ -31,6 +31,7 @@ module Snap.Extension.DB.MongoDB
     -- * Utility Functions
   , getObjId
   , bs2objid
+  , bs2objid'
   , objid2bs
   , lp
 

diff --git a/src/Snap/Extension/DB/MongoDB/Utils.hs b/src/Snap/Extension/DB/MongoDB/Utils.hs
@@ -9,6 +9,7 @@ import           Data.ByteString.Internal (c2w, w2c)
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Char8 as B8
 import           Data.ByteString (ByteString)
+import           Data.Maybe (fromJust)
 import qualified Data.Text as T
 import qualified Data.CompactString.Internal as CSI
 import qualified Data.CompactString.UTF8 as CS
@@ -37,11 +38,20 @@ objid2bs (Oid a b) = B8.pack . showHex a . showChar '-' . showHex b $ ""
 
 ------------------------------------------------------------------------------
 -- | Convert 'ByteString' into 'ObjectId'
-bs2objid :: ByteString -> ObjectId
+bs2objid :: ByteString -> Maybe ObjectId
-bs2objid bs = Oid a b
+bs2objid bs = do
-  where (a',b') = break (== '-') . B8.unpack $ bs
+  case B8.split '-' bs of
-        a = fst . head . readHex $ a'
+    (a':b':_) -> do
-        b = fst . head . readHex $ drop 1 b'
+      a <- fmap fst . headMay . readHex . B8.unpack $ a'
+      b <- fmap fst . headMay . readHex . B8.unpack $ b'
+      return $ Oid a b
+    _ -> Nothing
+
+------------------------------------------------------------------------------
+-- | Like 'bs2objid', but may blow with an error if the 'ByteString' can't be
+-- converted to an 'ObjectId'
+bs2objid' :: ByteString -> ObjectId
+bs2objid' = fromJust . bs2objid
 
 
 bs2cs :: ByteString -> UString