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