X-Git-Url: https://git.rkrishnan.org/?a=blobdiff_plain;f=src%2FFuncTorrent%2FBencode.hs;h=b55b3fe15509e5e4745a5451fd9509d842eb7c05;hb=9bb062d46ca3124fac06fd2a4c636b603e60f1cf;hp=2b641a7f4b51f57b8ef4f831c6156a20f79c054c;hpb=677f43e145aa2bb62f580be8b2fd72eccf212c1f;p=functorrent.git

diff --git a/src/FuncTorrent/Bencode.hs b/src/FuncTorrent/Bencode.hs
index 2b641a7..b55b3fe 100644
--- a/src/FuncTorrent/Bencode.hs
+++ b/src/FuncTorrent/Bencode.hs
@@ -1,24 +1,41 @@
+{-
+ - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
+ -
+ - This file is part of FuncTorrent.
+ -
+ - FuncTorrent is free software; you can redistribute it and/or modify
+ - it under the terms of the GNU General Public License as published by
+ - the Free Software Foundation; either version 3 of the License, or
+ - (at your option) any later version.
+ -
+ - FuncTorrent is distributed in the hope that it will be useful,
+ - but WITHOUT ANY WARRANTY; without even the implied warranty of
+ - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ - GNU General Public License for more details.
+ -
+ - You should have received a copy of the GNU General Public License
+ - along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
+ -}
+
 {-# LANGUAGE OverloadedStrings #-}
 module FuncTorrent.Bencode
     (BVal(..)
-    , bValToBList
-    , bValToBstr
-    , bValToInfoDict
     , bValToInteger
     , bstrToString
     , decode
+    , decodeWithLeftOvers
     , encode
     ) where
 
 import Prelude hiding (length, concat)
 
-import Control.Applicative ((<*))
 import Data.ByteString (ByteString, length, concat)
 import Data.ByteString.Char8 (unpack, pack)
-import Data.Functor ((<$>))
-import Data.Map.Strict (Map, fromList, keys, (!))
+import Data.Char (isLetter, isAscii)
+import Data.Map.Strict (Map, fromList, toList)
 import Text.ParserCombinators.Parsec
 import qualified Text.Parsec.ByteString as ParsecBS
+import Test.QuickCheck
 
 data BVal = Bint Integer
           | Bstr ByteString
@@ -26,25 +43,38 @@ data BVal = Bint Integer
           | Bdict (Map String BVal)
             deriving (Ord, Eq, Show)
 
+genNonEmptyString :: Gen String
+genNonEmptyString = arbitrary `suchThat` (\s -> ((s /= "") &&
+                                                 (all isAscii s) &&
+                                                 (all isLetter s)))
+
+instance Arbitrary ByteString where
+  arbitrary = pack <$> arbitrary
+
+instance Arbitrary BVal where
+  arbitrary = sized bval
+              where
+                bval :: Int -> Gen BVal
+                bval 0 = oneof [ Bint <$> arbitrary
+                               , Bstr <$> arbitrary]
+                bval n = oneof [ Bint <$> arbitrary
+                               , Bstr <$> arbitrary
+                               , Blist <$> vectorOf (n `div` 2) (bval (n `div` 4))
+                               , do keys <- vectorOf (n `div` 2) genNonEmptyString
+                                    vals <- vectorOf (n `div` 2) (bval (n `div` 4))
+                                    return $ Bdict $ fromList $ zip keys vals ]
+
 -- getters
 bValToInteger :: BVal -> Maybe Integer
 bValToInteger (Bint x) = Just x
 bValToInteger _        = Nothing
 
-bValToBstr :: BVal  -> Maybe ByteString
-bValToBstr (Bstr bs) = Just bs
-bValToBstr _         = Nothing
-
-bValToBList :: BVal    -> Maybe [BVal]
-bValToBList (Blist lst) = Just lst
-bValToBList _           = Nothing
-
-bValToInfoDict :: BVal     -> Maybe (Map String BVal)
-bValToInfoDict (Bdict dict) = Just dict
-bValToInfoDict _            = Nothing
+bValToBytestr :: BVal  -> Maybe ByteString
+bValToBytestr (Bstr bs) = Just bs
+bValToBytestr _         = Nothing
 
 bstrToString :: BVal -> Maybe String
-bstrToString bval     = unpack <$> bValToBstr bval
+bstrToString bval     = unpack <$> bValToBytestr bval
 
 -- $setup
 -- >>> import Data.Either
@@ -59,8 +89,7 @@ bstrToString bval     = unpack <$> bValToBstr bval
 -- Right ""
 --
 bencStr :: ParsecBS.Parser ByteString
-bencStr = do _ <- spaces
-             ds <- many1 digit <* char ':'
+bencStr = do ds <- many1 digit <* char ':'
              s <- count (read ds) anyChar
              return (pack s)
 
@@ -83,8 +112,7 @@ bencStr = do _ <- spaces
 -- >>> isLeft $ parse bencInt "Bint" (pack "i002e")
 -- True
 bencInt :: ParsecBS.Parser Integer
-bencInt = do _ <- spaces
-             ds <- between (char 'i') (char 'e') numbers
+bencInt = do ds <- between (char 'i') (char 'e') numbers
              return (read ds)
                where numbers = do d' <- char '-' <|> digit
                                   ds' <- many digit
@@ -106,8 +134,7 @@ bencInt = do _ <- spaces
 -- >>> parse bencList "Blist" (pack "l4:spam4:eggsli42eee")
 -- Right [Bstr "spam",Bstr "eggs",Blist [Bint 42]]
 bencList :: ParsecBS.Parser [BVal]
-bencList = do _ <- spaces
-              between (char 'l') (char 'e') (many bencVal)
+bencList = between (char 'l') (char 'e') (many bencVal)
 
 -- | parse dict
 --
@@ -121,9 +148,14 @@ bencList = do _ <- spaces
 -- Right (fromList [("publisher",Bstr "bob"),("publisher-webpage",Bstr "www.example.com"),("publisher.location",Bstr "home")])
 bencDict :: ParsecBS.Parser (Map String BVal)
 bencDict = between (char 'd') (char 'e') $ fromList <$> many kvpair
-  where kvpair = do k <- bencStr
+  where kvpair = do k <- bdictKey
                     v <- bencVal
-                    return (unpack k, v)
+                    return (k, v)
+        bdictKey = do
+          ds <- many1 digit <* char ':'
+          s <- count (read ds) anyChar
+          return s
+
 
 bencVal :: ParsecBS.Parser BVal
 bencVal = Bstr <$> bencStr <|>
@@ -134,6 +166,10 @@ bencVal = Bstr <$> bencStr <|>
 decode :: ByteString -> Either ParseError BVal
 decode = parse bencVal "BVal"
 
+decodeWithLeftOvers :: ByteString -> Either ParseError (BVal, ByteString)
+decodeWithLeftOvers = parse ((,) <$> bencVal <*> (fmap pack leftOvers)) "BVal with LeftOvers"
+  where leftOvers = manyTill anyToken eof
+
 -- Encode BVal into a bencoded ByteString. Inverse of decode
 
 -- TODO: Use builders and lazy byte string to get O(1) concatenation over O(n)
@@ -159,5 +195,9 @@ decode = parse bencVal "BVal"
 encode :: BVal -> ByteString
 encode (Bstr bs) = pack $ show (length bs) ++ ":" ++ unpack bs
 encode (Bint i) = pack $ "i" ++ show i ++ "e"
-encode (Blist xs) = pack $ "l" ++ unpack (concat $ map encode xs) ++ "e"
-encode (Bdict d) = concat [concat ["d", encode . Bstr . pack $ k , encode (d ! k) , "e"] | k <- keys d]
+encode (Blist xs) = concat ["l", concat $ map encode xs, "e"]
+encode (Bdict d) = concat ["d", concat kvlist, "e"]
+    where
+      kvlist :: [ByteString]
+      kvlist = [encPair kv | kv <- toList d]
+      encPair (k, v) = concat [encode (Bstr (pack k)), encode v]