3 files changed, 196 insertions, 35 deletions

diff --git a/Mainline.hs b/Mainline.hs
new file mode 100644
index 00000000..01d683d1
--- /dev/null
+++ b/Mainline.hs
@@ -0,0 +1,151 @@
+{-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE DeriveFoldable             #-}
+{-# LANGUAGE DeriveFunctor              #-}
+{-# LANGUAGE DeriveTraversable          #-}
+{-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+module Mainline where
+
+import Control.Arrow
+import Data.BEncode             as BE
+import Data.BEncode.BDict       as BE
+import Data.Bool
+import Data.ByteArray
+import Data.ByteString          (ByteString)
+import Data.ByteString          as B
+import Data.ByteString.Lazy     (toStrict)
+import Data.Data
+import Data.IP
+import qualified Data.Serialize as S
+import Data.Typeable
+import Data.Word
+import Network.Address          (Address, fromSockAddr, sockAddrPort,
+                                 toSockAddr, withPort)
+import Network.QueryResponse
+import Network.Socket
+
+newtype NodeId = NodeId ByteString
+ deriving (Eq,Ord,Show,ByteArrayAccess, BEncode)
+
+data NodeInfo = NodeInfo
+  { nodeId   :: NodeId
+  , nodeIP :: IP
+  , nodePort :: PortNumber
+  }
+
+nodeAddr :: NodeInfo -> SockAddr
+nodeAddr (NodeInfo _ ip port) = toSockAddr ip `withPort` port
+
+nodeInfo :: NodeId -> SockAddr -> Either String NodeInfo
+nodeInfo nid saddr
+    | Just ip <- fromSockAddr saddr
+    , Just port <- sockAddrPort saddr = Right $ NodeInfo nid ip port
+    | otherwise                       = Left "Address family not supported."
+
+-- | Types of RPC errors.
+data ErrorCode
+    -- | Some error doesn't fit in any other category.
+  = GenericError
+
+    -- | Occurs when server fail to process procedure call.
+  | ServerError
+
+    -- | Malformed packet, invalid arguments or bad token.
+  | ProtocolError
+
+    -- | Occurs when client trying to call method server don't know.
+  | MethodUnknown
+    deriving (Show, Read, Eq, Ord, Bounded, Typeable, Data)
+
+-- | According to the table:
+-- <http://bittorrent.org/beps/bep_0005.html#errors>
+instance Enum ErrorCode where
+  fromEnum GenericError  = 201
+  fromEnum ServerError   = 202
+  fromEnum ProtocolError = 203
+  fromEnum MethodUnknown = 204
+  {-# INLINE fromEnum #-}
+  toEnum 201 = GenericError
+  toEnum 202 = ServerError
+  toEnum 203 = ProtocolError
+  toEnum 204 = MethodUnknown
+  toEnum _   = GenericError
+  {-# INLINE toEnum #-}
+
+instance BEncode ErrorCode where
+  toBEncode = toBEncode . fromEnum
+  {-# INLINE toBEncode #-}
+  fromBEncode b = toEnum <$> fromBEncode b
+  {-# INLINE fromBEncode #-}
+
+data Error = Error
+    { errorCode    :: !ErrorCode  -- ^ The type of error.
+    , errorMessage :: !ByteString -- ^ Human-readable text message.
+    } deriving ( Show, Eq, Ord, Typeable, Data, Read )
+
+newtype TransactionId = TransactionId Word16
+ deriving (Eq, Ord, Show, BEncode)
+
+newtype Method = Method ByteString
+ deriving (Eq, Ord, Show, BEncode)
+
+data Message a = Q { msgOrigin   :: NodeId
+                   , msgID       :: TransactionId
+                   , qryPayload  :: a
+                   , qryMethod   :: Method
+                   , qryReadOnly :: Bool }
+
+               | R { msgOrigin      :: NodeId
+                   , msgID          :: TransactionId
+                   , rspPayload     :: Either Error a
+                   , rspReflectedIP :: Maybe SockAddr }
+
+instance BE.BEncode (Message BValue) where
+    toBEncode = encodeMessage
+    fromBEncode = error "fromBEncode"
+
+encodeMessage (Q origin tid a meth ro)
+    = case a of
+        BDict args -> encodeQuery tid meth (BDict $ genericArgs origin ro `union` args)
+        _          -> encodeQuery tid meth a -- XXX: Not really a valid query.
+encodeMessage (R origin tid v ip)
+    = case v of
+        Right vals -> encodeResponse tid vals (BString . encodeAddr <$> ip)
+        Left  err  -> encodeError tid err
+
+encodeAddr :: SockAddr -> ByteString
+encodeAddr (SockAddrInet port addr)
+    = S.runPut (S.putWord32host addr >> S.put (fromIntegral port :: Word16))
+encodeAddr (SockAddrInet6 port _ addr _)
+    = S.runPut (S.put addr >> S.put (fromIntegral port :: Word16))
+encodeAddr _ = B.empty
+
+genericArgs nodeid ro =
+       "id" .=! nodeid
+    .: "ro" .=? bool Nothing (Just (1 :: Int)) ro
+    .: endDict
+
+encodeError     tid (Error ecode emsg)       = encodeAny tid "e" (ecode,emsg) id
+encodeResponse  tid rvals              rip   = encodeAny tid "r" rvals        ("ip" .=? rip   .:)
+encodeQuery     tid qmeth              qargs = encodeAny tid "q" qmeth        ("a"  .=! qargs .:)
+
+encodeAny tid key val aux = toDict $
+    aux $  key .=! val
+        .: "t" .=! tid
+        .: "y" .=! key
+        .: endDict
+
+parsePacket :: ByteString -> SockAddr -> Either String (Message BValue, NodeInfo)
+parsePacket bs addr = do pkt <- BE.decode bs
+                         ni <- nodeInfo (msgOrigin pkt) addr
+                         return (pkt, ni)
+
+encodePacket :: Message BValue -> NodeInfo -> (ByteString, SockAddr)
+encodePacket msg ni = ( toStrict $ BE.encode msg
+                      , nodeAddr ni )
+
+newClient addr = do
+    udp <- udpTransport addr
+    let net = layerTransport parsePacket encodePacket udp
+    return net
+
diff --git a/src/Network/Address.hs b/src/Network/Address.hs
index 0fbd191a..d34b0431 100644
--- a/src/Network/Address.hs
+++ b/src/Network/Address.hs
@@ -83,6 +83,7 @@ module Network.Address
          -- * Utils
        , libUserAgent
        , sockAddrPort
+       , withPort
        , getBindAddress
        ) where
 
@@ -147,6 +148,12 @@ sockAddrPort (SockAddrInet6 p _ _ _) = Just p
 sockAddrPort _                       = Nothing
 {-# INLINE sockAddrPort #-}
 
+withPort :: SockAddr -> PortNumber -> SockAddr
+withPort (SockAddrInet  _ ip           ) port = SockAddrInet port ip
+withPort (SockAddrInet6 _ flow ip scope) port = SockAddrInet6 port flow ip scope
+withPort addr                            _    = addr
+{-# INLINE withPort #-}
+
 instance Address a => Address (NodeAddr a) where
   toSockAddr NodeAddr {..} = setPort nodePort $ toSockAddr nodeHost
   fromSockAddr sa = NodeAddr <$> fromSockAddr sa <*> sockAddrPort sa
diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 5b8bcda4..58db3a71 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -39,9 +39,9 @@ import Data.Maybe
 --  > r <- sendQuery client method q
 --  > -- Quit client.
 --  > quitServer
-forkListener :: Client err tbl meth tid addr x -> IO (IO ())
+forkListener :: Client err tbl meth tid addr x ctx -> IO (IO ())
 forkListener client = do
-    thread_id <- fork $ do
+    thread_id <- forkIO $ do
         myThreadId >>= flip labelThread "listener"
         fix $ handleMessage client
     return $ do
@@ -52,12 +52,12 @@ forkListener client = do
 -- out if 'forkListener' was never invoked to spawn a thread receive and
 -- dispatch the response.
 sendQuery ::
-    forall err a b tbl x meth tid addr.
-        Client err tbl meth tid addr x  -- ^ A query/response implementation.
-        -> Method addr x meth a b       -- ^ Information for marshalling the query.
-        -> a                            -- ^ The outbound query.
-        -> addr                         -- ^ Destination address of query.
-        -> IO (Maybe b)                 -- ^ The response, or 'Nothing' if it timed out.
+    forall err a b tbl x ctx meth tid addr.
+        Client err tbl meth tid addr x ctx -- ^ A query/response implementation.
+        -> MethodSerializer addr x ctx meth a b      -- ^ Information for marshalling the query.
+        -> a                               -- ^ The outbound query.
+        -> addr                            -- ^ Destination address of query.
+        -> IO (Maybe b)                    -- ^ The response, or 'Nothing' if it timed out.
 sendQuery (Client net d err pending whoami) meth q addr = do
     mvar <- newEmptyMVar
     tid <- atomically $ do
@@ -65,8 +65,8 @@ sendQuery (Client net d err pending whoami) meth q addr = do
         let (tid, tbl') = dispatchRegister (tableMethods d) mvar tbl
         writeTVar pending tbl'
         return tid
-    self <- whoami
-    sendMessage net addr (wrapQuery meth self addr q)
+    (self,ctx) <- whoami
+    sendMessage net addr (wrapQuery meth ctx self addr q)
     mres <- timeout (methodTimeout meth) $ takeMVar mvar
     case mres of
         Just x -> return $ Just $ unwrapResponse meth x
@@ -78,20 +78,21 @@ sendQuery (Client net d err pending whoami) meth q addr = do
 -- * Implementing a query\/response 'Client'.
 
 -- | All inputs required to implement a query\/response client.
-data Client err tbl meth tid addr x = Client
+data Client err tbl meth tid addr x ctx = Client
     { -- | The 'Transport' used to dispatch and receive packets.
       clientNet :: Transport err addr x
       -- | Methods for handling inbound packets.
-    , clientDispatcher :: DispatchMethods tbl err meth tid addr x
+    , clientDispatcher :: DispatchMethods tbl err meth tid addr x ctx
       -- | Methods for reporting various conditions.
     , clientErrorReporter :: ErrorReporter addr x meth tid err
       -- | State necessary for routing inbound responses and assigning unique
       -- /tid/ values for outgoing queries.
     , clientPending :: TVar tbl
-      -- | An action yielding this client\'s own address.  It is invoked once on
-      -- each outbound and inbound packet.  It is valid for this to always
-      -- return the same value.
-    , clientMyAddress :: IO addr
+      -- | An action yielding this client\'s own address along with some
+      -- context neccessary for serializing outgoing packets.  It is invoked
+      -- once on each outbound and inbound packet.  It is valid for this to
+      -- always return the same value.
+    , clientContext :: IO (addr,ctx)
     }
 
 -- | An incomming message can be classified into three cases.
@@ -101,12 +102,12 @@ data MessageClass err meth tid
     | IsUnknown err  -- ^ None of the above.
 
 -- | Handler for an inbound query of type _x_ from an address of type _addr_.
-data MethodHandler err addr x = forall a b. MethodHandler
+data MethodHandler err addr x ctx = forall a b. MethodHandler
     { -- | Parse the query into a more specific type for this method.
       methodParse :: x -> Either err a
-      -- | Serialize the response type for transmission.  Origin and destination
-      -- addresses for the packet are supplied in case they are required.
-    , methodSerialize :: addr -> addr -> b -> x
+      -- | Serialize the response for transmission, given a context /ctx/ and the origin
+      -- and destination addresses.
+    , methodSerialize :: ctx -> addr -> addr -> b -> x
       -- | Fully typed action to perform upon the query.  The remote origin
       -- address of the query is provided to the handler.
     , methodAction :: addr -> a -> IO b
@@ -115,18 +116,19 @@ data MethodHandler err addr x = forall a b. MethodHandler
 -- | Attempt to invoke a 'MethodHandler' upon a given inbound query.  If the
 -- parse is successful, the returned IO action will construct our reply.
 -- Otherwise, a parse err is returned.
-dispatchQuery :: MethodHandler err addr x -- ^ Handler to invoke.
-                 -> addr                  -- ^ Our own address, to which the query was sent.
-                 -> x                     -- ^ The query packet.
-                 -> addr                  -- ^ The origin address of the query.
+dispatchQuery :: MethodHandler err addr x ctx -- ^ Handler to invoke.
+                 -> ctx                       -- ^ Arbitrary context used during serialization.
+                 -> addr                      -- ^ Our own address, to which the query was sent.
+                 -> x                         -- ^ The query packet.
+                 -> addr                      -- ^ The origin address of the query.
                  -> Either err (IO x)
-dispatchQuery (MethodHandler unwrapQ wrapR f) self x addr =
-    fmap (\a -> wrapR self addr <$> f addr a) $ unwrapQ x
+dispatchQuery (MethodHandler unwrapQ wrapR f) ctx self x addr =
+    fmap (\a -> wrapR ctx self addr <$> f addr a) $ unwrapQ x
 
 -- | These four parameters are required to implement an ougoing query.  A
--- peer-to-peer algorithm will define a 'Method' for every 'MethodHandler' that
+-- peer-to-peer algorithm will define a 'MethodSerializer' for every 'MethodHandler' that
 -- might be returned by 'lookupHandler'.
-data Method addr x meth a b = Method
+data MethodSerializer addr x ctx meth a b = MethodSerializer
     { -- | Seconds to wait for a response.
       methodTimeout :: Int
       -- | A method identifier used for error reporting.  This needn't be the
@@ -134,8 +136,9 @@ data Method addr x meth a b = Method
     , method :: meth
       -- | Serialize the outgoing query /a/ into a transmitable packet /x/.
       -- The /addr/ arguments are, respectively, our own origin address and the
-      -- destination of the request.
-    , wrapQuery :: addr -> addr -> a -> x
+      -- destination of the request.  The /ctx/ argument is useful for attaching
+      -- auxillary notations on all outgoing packets.
+    , wrapQuery :: ctx -> addr -> addr -> a -> x
       -- | Parse an inbound packet /x/ into a response /b/ for this query.
     , unwrapResponse :: x -> b
     }
@@ -227,11 +230,11 @@ transactionTableMethods insert delete lookup generate = TableMethods
     }
 
 -- | A set of methods neccessary for dispatching incomming packets.
-data DispatchMethods tbl err meth tid addr x = DispatchMethods
+data DispatchMethods tbl err meth tid addr x ctx = DispatchMethods
     { -- | Clasify an inbound packet as a query or response.
       classifyInbound :: x -> MessageClass err meth tid
       -- | Lookup the handler for a inbound query.
-    , lookupHandler :: meth -> Maybe (MethodHandler err addr x)
+    , lookupHandler :: meth -> Maybe (MethodHandler err addr x ctx)
       -- | Methods for handling incomming responses.
     , tableMethods :: TableMethods tbl tid x
     }
@@ -264,7 +267,7 @@ data ErrorReporter addr x meth tid err = ErrorReporter
 -- 'fix' in a forked thread that loops until 'awaitMessage' returns 'Nothing'
 -- or throws an exception.
 handleMessage ::
-    Client err tbl meth tid addr x
+    Client err tbl meth tid addr x ctx
     -> IO ()
     -> IO ()
 handleMessage (Client net d err pending whoami) again = do
@@ -276,10 +279,10 @@ handleMessage (Client net d err pending whoami) again = do
             IsQuery meth -> case lookupHandler d meth of
                                 Nothing -> reportMissingHandler err meth addr plain
                                 Just m  -> do
-                                  self <- whoami
+                                  (self,ctx) <- whoami
                                   either (reportParseError err)
                                          (>>= sendMessage net addr)
-                                         (dispatchQuery m self plain addr)
+                                         (dispatchQuery m ctx self plain addr)
             IsResponse tid -> do
                 action <- atomically $ do
                     ts0 <- readTVar pending