src/FuncTorrent/Bencode.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Bencode
   3     (BVal(..)
   4     , bValToBList
   5     , bValToBytestr
   6     , bValToInfoDict
   7     , bValToInteger
   8     , bstrToString
   9     , decode
  10     , encode
  11     ) where
  12
  13 import Prelude hiding (length, concat)
  14
  15 import Data.ByteString (ByteString, length, concat)
  16 import Data.ByteString.Char8 (unpack, pack)
  17 import Data.Map.Strict (Map, fromList, toList)
  18 import Text.ParserCombinators.Parsec
  19 import qualified Text.Parsec.ByteString as ParsecBS
  20 import Test.QuickCheck
  21
  22 data BVal = Bint Integer
  23           | Bstr ByteString
  24           | Blist [BVal]
  25           | Bdict (Map String BVal)
  26             deriving (Ord, Eq, Show)
  27
  28 instance Arbitrary ByteString where
  29   arbitrary = pack <$> arbitrary
  30
  31 instance Arbitrary BVal where
  32   arbitrary = sized bval
  33               where
  34                 bval :: Int -> Gen BVal
  35                 bval 0 = oneof [ Bint <$> arbitrary
  36                                , Bstr <$> arbitrary]
  37                 bval n = oneof [ Bint <$> arbitrary
  38                                , Bstr <$> arbitrary
  39                                , Blist <$> vectorOf n (bval (n `div` 4))
  40                                , do keys <- vectorOf n arbitrary
  41                                     vals <- vectorOf n (bval (n `div` 4))
  42                                     return $ Bdict $ fromList $ zip keys vals ]
  43
  44 -- getters
  45 bValToInteger :: BVal -> Maybe Integer
  46 bValToInteger (Bint x) = Just x
  47 bValToInteger _        = Nothing
  48
  49 bValToBytestr :: BVal  -> Maybe ByteString
  50 bValToBytestr (Bstr bs) = Just bs
  51 bValToBytestr _         = Nothing
  52
  53 bValToBList :: BVal    -> Maybe [BVal]
  54 bValToBList (Blist lst) = Just lst
  55 bValToBList _           = Nothing
  56
  57 bValToInfoDict :: BVal     -> Maybe (Map String BVal)
  58 bValToInfoDict (Bdict dict) = Just dict
  59 bValToInfoDict _            = Nothing
  60
  61 bstrToString :: BVal -> Maybe String
  62 bstrToString bval     = unpack <$> bValToBytestr bval
  63
  64 -- $setup
  65 -- >>> import Data.Either
  66
  67 -- | parse strings
  68 --
  69 -- >>> parse bencStr "Bstr" (pack "4:spam")
  70 -- Right "spam"
  71 -- >>> parse bencStr "Bstr" (pack "0:")
  72 -- Right ""
  73 -- >>> parse bencStr "Bstr" (pack "0:hello")
  74 -- Right ""
  75 --
  76 bencStr :: ParsecBS.Parser ByteString
  77 bencStr = do ds <- many1 digit <* char ':'
  78              s <- count (read ds) anyChar
  79              return (pack s)
  80
  81 -- | parse integers
  82 --
  83 -- >>> parse bencInt "Bint" (pack "i42e")
  84 -- Right 42
  85 -- >>> parse bencInt "Bint" (pack "i123e")
  86 -- Right 123
  87 -- >>> parse bencInt "Bint" (pack "i1e")
  88 -- Right 1
  89 -- >>> parse bencInt "Bint" (pack "i0e")
  90 -- Right 0
  91 -- >>> parse bencInt "Bint" (pack "i-1e")
  92 -- Right (-1)
  93 -- >>> isLeft $ parse bencInt "Bint" (pack "i01e")
  94 -- True
  95 -- >>> isLeft $ parse bencInt "Bint" (pack "i00e")
  96 -- True
  97 -- >>> isLeft $ parse bencInt "Bint" (pack "i002e")
  98 -- True
  99 bencInt :: ParsecBS.Parser Integer
 100 bencInt = do ds <- between (char 'i') (char 'e') numbers
 101              return (read ds)
 102                where numbers = do d' <- char '-' <|> digit
 103                                   ds' <- many digit
 104                                   parseNumber d' ds'
 105                      parseNumber '0' []  = return "0"
 106                      parseNumber '0' _ = unexpected "numbers cannot be left-padded with zeros"
 107                      parseNumber '-' []  = unexpected "sign without any digits"
 108                      parseNumber '-' (d'':_) | d'' == '0' = unexpected "numbers cannot be left-padded with zeros"
 109                      parseNumber d'' ds'' = return (d'':ds'')
 110
 111 -- | parse lists
 112 --
 113 -- >>> parse bencList "Blist" (pack "le")
 114 -- Right []
 115 -- >>> parse bencList "Blist" (pack "l4:spam4:eggse")
 116 -- Right [Bstr "spam",Bstr "eggs"]
 117 -- >>> parse bencList "Blist" (pack "l4:spami42ee")
 118 -- Right [Bstr "spam",Bint 42]
 119 -- >>> parse bencList "Blist" (pack "l4:spam4:eggsli42eee")
 120 -- Right [Bstr "spam",Bstr "eggs",Blist [Bint 42]]
 121 bencList :: ParsecBS.Parser [BVal]
 122 bencList = between (char 'l') (char 'e') (many bencVal)
 123
 124 -- | parse dict
 125 --
 126 -- >>> parse bencDict "Bdict" (pack "de")
 127 -- Right (fromList [])
 128 -- >>> parse bencDict "Bdict" (pack "d3:cow3:moo4:spam4:eggse")
 129 -- Right (fromList [("cow",Bstr "moo"),("spam",Bstr "eggs")])
 130 -- >>> parse bencDict "Bdict" (pack "d4:spaml1:a1:bee")
 131 -- Right (fromList [("spam",Blist [Bstr "a",Bstr "b"])])
 132 -- >>> parse bencDict "Bdict" (pack "d9:publisher3:bob17:publisher-webpage15:www.example.com18:publisher.location4:homee")
 133 -- Right (fromList [("publisher",Bstr "bob"),("publisher-webpage",Bstr "www.example.com"),("publisher.location",Bstr "home")])
 134 bencDict :: ParsecBS.Parser (Map String BVal)
 135 bencDict = between (char 'd') (char 'e') $ fromList <$> many kvpair
 136   where kvpair = do k <- bencStr
 137                     v <- bencVal
 138                     return (unpack k, v)
 139
 140 bencVal :: ParsecBS.Parser BVal
 141 bencVal = Bstr <$> bencStr <|>
 142           Bint <$> bencInt <|>
 143           Blist <$> bencList <|>
 144           Bdict <$> bencDict
 145
 146 decode :: ByteString -> Either ParseError BVal
 147 decode = parse bencVal "BVal"
 148
 149 -- Encode BVal into a bencoded ByteString. Inverse of decode
 150
 151 -- TODO: Use builders and lazy byte string to get O(1) concatenation over O(n)
 152 -- provided by lists.
 153
 154 -- TODO: encode . decode pair might be a good candidate for Quickcheck.
 155 -- | encode bencoded-values
 156 --
 157 -- >>> encode (Bstr (pack ""))
 158 -- "0:"
 159 -- >>> encode (Bstr (pack "spam"))
 160 -- "4:spam"
 161 -- >>> encode (Bint 0)
 162 -- "i0e"
 163 -- >>> encode (Bint 42)
 164 -- "i42e"
 165 -- >>> encode (Blist [(Bstr (pack "spam")), (Bstr (pack "eggs"))])
 166 -- "l4:spam4:eggse"
 167 -- >>> encode (Blist [])
 168 -- "le"
 169 -- >>> encode (Bdict (fromList [("spam", Bstr $ pack "eggs")]))
 170 -- "d4:spam4:eggse"
 171 encode :: BVal -> ByteString
 172 encode (Bstr bs) = pack $ show (length bs) ++ ":" ++ unpack bs
 173 encode (Bint i) = pack $ "i" ++ show i ++ "e"
 174 encode (Blist xs) = concat ["l", concat $ map encode xs, "e"]
 175 encode (Bdict d) = concat ["d", concat kvlist, "e"]
 176     where
 177       kvlist :: [ByteString]
 178       kvlist = [encPair kv | kv <- toList d]
 179       encPair (k, v) = concat [encode . Bstr . pack $ k, encode v]