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