/BEncoding.hs

  1module BEncoding (BValue(..), encode, decode, parseFile, bdict, bdictLookup, infoHash) where
  2
  3import qualified Data.ByteString.Lazy.Char8 as B8
  4import qualified Data.ByteString.Lazy as W8
  5import qualified Data.ByteString.Char8 as SB8
  6import qualified Data.ByteString as SW8
  7import Data.Binary.Strict.Get
  8import Data.Char (isDigit, chr)
  9import Control.Monad
 10import Test.QuickCheck
 11import Data.List (intercalate)
 12import Control.DeepSeq
 13import OpenSSL.Digest (MessageDigest(SHA1))
 14import OpenSSL.Digest.ByteString.Lazy (digest)
 15import Prelude hiding (getChar, takeWhile)
 16
 17
 18data BValue = BInteger Integer
 19            | BString B8.ByteString
 20            | BList [BValue]
 21            | BDict [(BValue, BValue)]
 22            deriving (Eq, Ord)
 23instance NFData BValue where
 24    rnf (BInteger i) = rnf i
 25    rnf (BString bs) = rnf $ B8.unpack bs
 26    rnf (BList []) = ()
 27    rnf (BList (x:xs)) = rnf x `seq`
 28                       rnf (BList xs)
 29    rnf (BDict []) = ()
 30    rnf (BDict ((x, x'):xs)) = rnf x `seq`
 31                             rnf x' `seq`
 32                             rnf (BDict xs)
 33
 34encode :: BValue -> B8.ByteString
 35encode (BInteger i) = B8.singleton 'i' `B8.append`
 36                      B8.pack (show i) `B8.append`
 37                      B8.singleton 'e'
 38encode (BString s) = B8.pack (show $ B8.length s) `B8.append`
 39                     B8.singleton ':' `B8.append`
 40                     s
 41encode (BList xs) = B8.singleton 'l' `B8.append`
 42                    B8.concat (map encode xs) `B8.append`
 43                    B8.singleton 'e'
 44encode (BDict xs) = B8.singleton 'd' `B8.append`
 45                    B8.concat (map (\(k, v) ->
 46                                        encode k `B8.append` encode v
 47                                   ) xs) `B8.append`
 48                    B8.singleton 'e'
 49
 50bdict :: [(String, BValue)] -> BValue
 51bdict = BDict . map (\(s, v) ->
 52                         (BString $ B8.pack s, v)
 53                    )
 54
 55instance Show BValue where
 56    show (BInteger i) = show i
 57    show (BString s) = show $ B8.unpack s
 58    show (BList l) = show l
 59    show (BDict d) = "{ " ++ (intercalate ", " $
 60                              map (\(k, v) -> show k ++ ": " ++ show v) d
 61                             ) ++
 62                     " }"
 63
 64
 65decode :: SB8.ByteString -> Either String BValue
 66decode bs = case runGet decoder bs of
 67              (Right a, _) -> Right a
 68              (Left e, _) -> Left $ "Parse: " ++ e
 69
 70decoder :: Get BValue
 71decoder = do c1 <- getChar
 72             case c1 of
 73               'i' ->
 74                     do iS <- takeWhile (\c -> c == '-' || isDigit c)
 75                        char 'e'
 76                        let Just (i, _) = SB8.readInteger iS
 77                        return $ BInteger i
 78               d | isDigit d ->
 79                     do lS <- SB8.cons d `liftM` takeWhile isDigit
 80                        char ':'
 81                        let Just (l, _) = SB8.readInteger lS
 82                        s <- getByteString $ fromIntegral l
 83                        return $ BString $ B8.fromChunks [s]
 84               'l' ->
 85                     BList `liftM` manyTill decoder 'e'
 86               'd' ->
 87                     BDict `liftM` manyTill (do k <- decoder
 88                                                v <- decoder
 89                                                return (k, v)
 90                                            ) 'e'
 91               _ -> fail $ "unexpected type: " ++ show c1
 92    where getChar :: Get Char
 93          getChar = (chr . fromIntegral) `liftM` getWord8
 94          char :: Char -> Get ()
 95          char c = getChar >>= \c' ->
 96                   if c == c'
 97                   then return ()
 98                   else fail $ "expected " ++ show c
 99          takeWhile :: (Char -> Bool) -> Get SB8.ByteString
100          takeWhile p = do buf <- lookAhead $ remaining >>= getByteString
101                           let bufLen = SB8.length buf
102                               len = run 0
103                               run i | i >= bufLen = bufLen
104                                     | p (SB8.index buf i) = run $ i + 1
105                                     | otherwise = i
106                           getByteString len
107          manyTill :: Get a -> Char -> Get [a]
108          manyTill e c = lookAhead getChar >>= \c' ->
109                         if c == c'
110                         then getChar >> return []
111                         else do el <- e
112                                 (el:) `liftM` manyTill e c
113
114parseFile :: FilePath -> IO (Maybe BValue)
115parseFile = return . result . decode <=< SB8.readFile
116    where result (Left _) = Nothing
117          result (Right a) = Just a
118
119bdictLookup :: BValue -> String -> Maybe BValue
120bdictLookup (BDict dict) key = lookup (BString $ B8.pack key) dict
121bdictLookup _ _ = Nothing
122
123infoHash :: BValue -> IO (Maybe SW8.ByteString)
124infoHash metaInfo
125    = maybe (return Nothing) (return . Just <=< sha1) $
126      encode `liftM` (metaInfo `bdictLookup` "info")
127    where sha1 :: B8.ByteString -> IO SW8.ByteString
128          sha1 bs = SW8.pack `liftM` digest SHA1 bs
129
130{-
131instance Arbitrary Char where
132  arbitrary = chr `fmap` oneof [choose (0,127), choose (0,255)]-}
133instance Arbitrary B8.ByteString where
134    arbitrary = B8.pack `fmap`  arbitrary
135instance Arbitrary SB8.ByteString where
136    arbitrary = SB8.pack `fmap`  arbitrary
137instance Arbitrary BValue where
138    arbitrary = frequency [(10, BInteger `liftM` arbitrary),
139                           (5, resize 150 $ BString `liftM` arbitrary),
140                           (2, resize 5 $ BList `liftM` arbitrary),
141                           (2, resize 5 $ BDict `liftM` arbitrary)]
142
143propEncodeDecode val = decode (SB8.concat $ B8.toChunks $ encode val) == Right val