module Data.Torrent.Piece
(
PieceIx
, PieceCount
, PieceSize
, minPieceSize
, maxPieceSize
, defaultPieceSize
, Piece (..)
, pieceSize
, isPiece
, HashArray (..)
, PieceInfo (..)
, pieceCount
, pieceLength
, pieceHashes
, pieceHash
, checkPieceLazy
, getPieceInfo
, putPieceInfo
) where
import Control.DeepSeq
import Control.Lens
import qualified Crypto.Hash.SHA1 as SHA1
import Data.Aeson (ToJSON(..), FromJSON(..), Value(..), withText)
import Data.Aeson.TH
import Data.BEncode
import Data.BEncode.Types
import Data.Bits
import Data.Bits.Extras
import Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Base64 as Base64
import Data.Char
import Data.Int
import Data.List as L
import Data.Text.Encoding as T
import Data.Typeable
import Text.PrettyPrint
import Text.PrettyPrint.Class
import Data.Torrent.Block
class Lint a where
lint :: a -> Either String a
type PieceCount = Int
optimalPieceCount :: PieceCount
optimalPieceCount = 1000
minPieceSize :: Int
minPieceSize = defaultTransferSize * 4
maxPieceSize :: Int
maxPieceSize = 4 * 1024 * 1024
toPow2 :: Int -> Int
toPow2 x = bit $ fromIntegral (leadingZeros (0 :: Int) leadingZeros x)
defaultPieceSize :: Int64 -> Int
defaultPieceSize x = max minPieceSize $ min maxPieceSize $ toPow2 pc
where
pc = fromIntegral (x `div` fromIntegral optimalPieceCount)
data Piece a = Piece
{
pieceIndex :: !PieceIx
, pieceData :: !a
} deriving (Show, Read, Eq, Typeable)
$(deriveJSON (L.map toLower . L.dropWhile isLower) ''Piece)
instance NFData (Piece a)
instance Pretty (Piece a) where
pretty Piece {..} = "Piece" <+> braces ("index" <+> "=" <+> int pieceIndex)
pieceSize :: Piece BL.ByteString -> PieceSize
pieceSize Piece {..} = fromIntegral (BL.length pieceData)
isPiece :: PieceSize -> Block BL.ByteString -> Bool
isPiece pieceLen blk @ (Block i offset _) =
offset == 0 && blockSize blk == pieceLen && i >= 0
newtype HashArray = HashArray { unHashArray :: ByteString }
deriving (Show, Read, Eq, BEncode)
instance ToJSON HashArray where
toJSON (HashArray bs) = String $ T.decodeUtf8 $ Base64.encode bs
instance FromJSON HashArray where
parseJSON = withText "HashArray" $
either fail (return . HashArray) . Base64.decode . T.encodeUtf8
data PieceInfo = PieceInfo
{ piPieceLength :: !PieceSize
, piPieceHashes :: !HashArray
} deriving (Show, Read, Eq, Typeable)
$(deriveJSON (L.map toLower . L.dropWhile isLower) ''PieceInfo)
makeLensesFor [("piPieceLength", "pieceLength")] ''PieceInfo
makeLensesFor [("piPieceHashes", "pieceHashes")] ''PieceInfo
instance NFData PieceInfo
instance Lint PieceInfo where
lint pinfo @ PieceInfo {..}
| BS.length (unHashArray piPieceHashes) `rem` hashsize == 0
, piPieceLength >= 0 = return pinfo
| otherwise = Left undefined
putPieceInfo :: PieceInfo -> BDict -> BDict
putPieceInfo PieceInfo {..} cont =
"piece length" .=! piPieceLength
.: "pieces" .=! piPieceHashes
.: cont
getPieceInfo :: Get PieceInfo
getPieceInfo = do
PieceInfo <$>! "piece length"
<*>! "pieces"
instance BEncode PieceInfo where
toBEncode = toDict . (`putPieceInfo` endDict)
fromBEncode = fromDict getPieceInfo
instance Pretty PieceInfo where
pretty PieceInfo {..} = "Piece size: " <> int piPieceLength
hashsize :: Int
hashsize = 20
slice :: Int -> Int -> ByteString -> ByteString
slice start len = BS.take len . BS.drop start
pieceHash :: PieceInfo -> PieceIx -> ByteString
pieceHash PieceInfo {..} i = slice (hashsize * i) hashsize (unHashArray piPieceHashes)
pieceCount :: PieceInfo -> PieceCount
pieceCount PieceInfo {..} = BS.length (unHashArray piPieceHashes) `quot` hashsize
isLastPiece :: PieceInfo -> PieceIx -> Bool
isLastPiece ci i = pieceCount ci == succ i
checkPieceLazy :: PieceInfo -> Piece BL.ByteString -> Bool
checkPieceLazy pinfo @ PieceInfo {..} Piece {..}
= (fromIntegral (BL.length pieceData) == piPieceLength
|| isLastPiece pinfo pieceIndex)
&& SHA1.hashlazy pieceData == pieceHash pinfo pieceIndex