src/FuncTorrent/Bencode.hs

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