3 files changed, 226 insertions, 0 deletions

diff --git a/src/Network/BitTorrent/PeerWire/Block.hs b/src/Network/BitTorrent/PeerWire/Block.hs
new file mode 100644
index 00000000..33e3dead
--- /dev/null
+++ b/src/Network/BitTorrent/PeerWire/Block.hs
@@ -0,0 +1,56 @@
+module Network.BitTorrent.PeerWire.Block
+       ( BlockIx(..), Block(..)
+       , defaultBlockSize
+       , blockRange, ixRange, pieceIx
+       , isPiece
+       ) where
+
+import           Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+import Data.Int
+
+
+data BlockIx = BlockIx {
+    ixPiece  :: {-# UNPACK #-} !Int -- ^ Zero-based piece index.
+  , ixOffset :: {-# UNPACK #-} !Int -- ^ Zero-based byte offset within the piece.
+  , ixLength :: {-# UNPACK #-} !Int -- ^ Block size starting from offset.
+  } deriving (Show, Eq)
+
+data Block = Block {
+    blkPiece  :: Int         -- ^ Zero-based piece index.
+  , blkOffset :: Int         -- ^ Zero-based byte offset within the piece.
+  , blkData   :: ByteString  -- ^ Payload.
+  } deriving (Show, Eq)
+
+
+-- | Widely used semi-official block size.
+defaultBlockSize :: Int
+defaultBlockSize = 16 * 1024
+
+
+isPiece :: Int -> Block -> Bool
+isPiece pieceSize (Block i offset bs) =
+  offset == 0 && B.length bs == pieceSize && i >= 0
+{-# INLINE isPiece #-}
+
+pieceIx :: Int -> Int -> BlockIx
+pieceIx i = BlockIx i 0
+{-# INLINE pieceIx #-}
+
+blockRange :: (Num a, Integral a) => Int -> Block -> (a, a)
+blockRange pieceSize blk = (offset, offset + len)
+  where
+    offset = fromIntegral pieceSize * fromIntegral (blkPiece blk)
+           + fromIntegral (blkOffset blk)
+    len    = fromIntegral (B.length (blkData blk))
+{-# INLINE blockRange #-}
+{-# SPECIALIZE blockRange :: Int -> Block -> (Int64, Int64) #-}
+
+ixRange :: (Num a, Integral a) => Int -> BlockIx -> (a, a)
+ixRange pieceSize ix = (offset, offset + len)
+  where
+    offset = fromIntegral  pieceSize * fromIntegral (ixPiece ix)
+           + fromIntegral (ixOffset ix)
+    len    = fromIntegral (ixLength ix)
+{-# INLINE ixRange #-}
+{-# SPECIALIZE ixRange :: Int -> BlockIx -> (Int64, Int64) #-}
diff --git a/src/Network/BitTorrent/PeerWire/Handshake.hs b/src/Network/BitTorrent/PeerWire/Handshake.hs
new file mode 100644
index 00000000..6ce37887
--- /dev/null
+++ b/src/Network/BitTorrent/PeerWire/Handshake.hs
@@ -0,0 +1,79 @@
+-- |
+--   Copyright   :  (c) Sam T. 2013
+--   License     :  MIT
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+{-# LANGUAGE OverloadedStrings #-}
+module Network.BitTorrent.PeerWire.Handshake
+       ( Handshake
+       , handshakeMaxSize
+       , defaultBTProtocol, defaultReserved, defaultHandshake
+       , handshake
+       ) where
+
+import Control.Applicative
+import Data.Word
+import           Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+import Data.Serialize as S
+import Data.Torrent.InfoHash
+import Network
+import Network.Socket.ByteString
+
+import Network.BitTorrent.PeerID
+
+
+-- | In order to establish the connection between peers we should send 'Handshake'
+--   message. The 'Handshake' is a required message and must be the first message
+--   transmitted by the peer to the another peer.
+data Handshake = Handshake {
+    hsProtocol    :: ByteString  -- ^ Identifier of the protocol.
+  , hsReserved    :: Word64      -- ^ Reserved bytes, rarely used.
+  , hsInfoHash    :: InfoHash    -- ^ Hash from the metainfo file.
+    -- This /should be/ same hash that is transmitted in tracker requests.
+  , hsPeerID      :: PeerID      -- ^ Peer id of the initiator.
+    -- This is /usually the same peer id that is transmitted in tracker requests.
+  } deriving (Show, Eq)
+
+instance Serialize Handshake where
+  put hs = do
+    putWord8 (fromIntegral (B.length (hsProtocol hs)))
+    putByteString (hsProtocol hs)
+    putWord64be   (hsReserved hs)
+    put (hsInfoHash hs)
+    put (hsPeerID hs)
+
+  get = do
+    len  <- getWord8
+    Handshake <$> getBytes (fromIntegral len)
+              <*> getWord64be
+              <*> get
+              <*> get
+
+-- | Maximum size of handshake message in bytes.
+handshakeMaxSize :: Int
+handshakeMaxSize = 1 + 256 + 8 + 20 + 20
+
+-- | Default protocol string "BitTorrent protocol" as is.
+defaultBTProtocol :: ByteString
+defaultBTProtocol = "BitTorrent protocol"
+
+-- | Default reserved word is 0.
+defaultReserved :: Word64
+defaultReserved = 0
+
+-- | Length of info hash and peer id is unchecked, so it /should/ be equal 20.
+defaultHandshake :: InfoHash -> PeerID -> Handshake
+defaultHandshake = Handshake defaultBTProtocol defaultReserved
+
+
+-- TODO check if hash the same
+-- | Handshaking with a peer specified by the second argument.
+--
+handshake :: Socket -> Handshake -> IO (Either String Handshake)
+handshake sock hs = do
+  sendAll sock (S.encode hs)
+  r <- recv sock handshakeMaxSize
+  return (S.decode r)
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
new file mode 100644
index 00000000..1bcb2ee5
--- /dev/null
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -0,0 +1,91 @@
+module Network.BitTorrent.PeerWire.Message
+       ( Message(..)
+       ) where
+
+import Control.Applicative
+import           Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+
+import Data.Serialize
+
+import Network.BitTorrent.PeerWire.Block
+
+
+-- TODO comment message constructors
+data Message = KeepAlive
+             | Choke
+             | Unchoke
+             | Interested
+             | NotInterested
+             | Have     Int
+             | Bitfield ByteString
+             | Request  BlockIx
+             | Piece    Block
+             | Cancel   BlockIx
+             | Port     Int
+               deriving (Show, Eq)
+
+instance Serialize BlockIx where
+  {-# SPECIALIZE instance Serialize BlockIx #-}
+  get = BlockIx <$> getInt <*> getInt <*> getInt
+  {-# INLINE get #-}
+
+  put ix = do putInt (ixPiece ix)
+              putInt (ixOffset ix)
+              putInt (ixLength ix)
+  {-# INLINE put #-}
+
+instance Serialize Message where
+  get = do
+    len <- getInt
+    lookAhead $ ensure len
+    if len == 0 then return KeepAlive -- FIX check if BS is empty instead of reading len
+      else do
+        mid <- getWord8
+        case mid of
+          0 -> return Choke
+          1 -> return Unchoke
+          2 -> return Interested
+          3 -> return NotInterested
+          4 -> Have     <$> getInt
+          5 -> Bitfield <$> getBytes (pred len)
+          6 -> Request  <$> get
+          7 -> Piece    <$> getBlock (len - 9)
+          8 -> Cancel   <$> get
+          9 -> (Port . fromIntegral) <$> getWord16be
+          _ -> fail $ "unknown message ID: " ++ show mid
+
+    where
+      getBlock :: Int -> Get Block
+      getBlock len = Block <$> getInt <*> getInt <*> getBytes len
+      {-# INLINE getBlock #-}
+
+  put KeepAlive     = putInt 0
+  put Choke         = putInt 1  >> putWord8 0
+  put Unchoke       = putInt 1  >> putWord8 1
+  put Interested    = putInt 1  >> putWord8 2
+  put NotInterested = putInt 1  >> putWord8 3
+  put (Have i)      = putInt 5  >> putWord8 4 >> putInt i
+  put (Bitfield b)  = putInt l  >> putWord8 5 >> putByteString b
+    where l = succ (B.length b)
+          {-# INLINE l #-}
+  put (Request blk) = putInt 13 >> putWord8 6 >> put blk
+  put (Piece   blk) = putInt l  >> putWord8 7 >> putBlock
+    where l = 9 + B.length (blkData blk)
+          {-# INLINE l #-}
+          putBlock = do putInt (blkPiece blk)
+                        putInt (blkOffset  blk)
+                        putByteString (blkData blk)
+          {-# INLINE putBlock #-}
+
+  put (Cancel  blk) = putInt 13 >> putWord8 8 >> put blk
+  put (Port    p  ) = putInt 3  >> putWord8 9 >> putWord16be (fromIntegral p)
+
+
+getInt :: Get Int
+getInt = fromIntegral <$> getWord32be
+{-# INLINE getInt #-}
+
+putInt :: Putter Int
+putInt = putWord32be . fromIntegral
+{-# INLINE putInt #-}