1 files changed, 234 insertions, 2 deletions

diff --git a/src/Network/RPC.hs b/src/Network/RPC.hs
index 7fb0e571..2333766a 100644
--- a/src/Network/RPC.hs
+++ b/src/Network/RPC.hs
@@ -1,29 +1,80 @@
 {-# LANGUAGE ConstraintKinds        #-}
-{-# LANGUAGE DeriveDataTypeable     #-}
+{-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE DeriveFunctor              #-}
+{-# LANGUAGE DeriveFoldable             #-}
+{-# LANGUAGE DeriveTraversable          #-}
+{-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE FlexibleContexts       #-}
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE MultiParamTypeClasses  #-}
 {-# LANGUAGE RankNTypes             #-}
 {-# LANGUAGE ScopedTypeVariables    #-}
 {-# LANGUAGE TypeFamilies           #-}
+{-# LANGUAGE StandaloneDeriving     #-}
 module Network.RPC where
 
+import Control.Applicative
+import qualified Text.ParserCombinators.ReadP as RP
+import Data.Digest.CRC32C
+import Data.Word
+import Data.Monoid
+import Data.Hashable
+import Data.String
 import Data.Bits
 import Data.ByteString (ByteString)
 import Data.Kind       (Constraint)
 import Data.Data
+import Data.Default
+import Data.List.Split
+import Data.Ord
+import Data.IP
 import Network.Socket
+import Text.PrettyPrint as PP hiding ((<>))
 import Text.PrettyPrint.HughesPJClass  hiding (($$), (<>))
+import Text.Read        (readMaybe)
 import Data.Serialize                  as S
 import qualified Data.ByteString.Char8 as Char8
+import qualified Data.ByteString as BS
 import Data.ByteString.Base16          as Base16
+import System.Entropy
+
+class (Eq a, Serialize a, Typeable a, Hashable a, Pretty a)
+    => Address a where
+  toSockAddr   :: a        -> SockAddr
+  fromSockAddr :: SockAddr -> Maybe a
+
+fromAddr :: (Address a, Address b) => a -> Maybe b
+fromAddr = fromSockAddr . toSockAddr
+
+-- | Note that port is zeroed.
+instance Address IPv4 where
+  toSockAddr = SockAddrInet 0 . toHostAddress
+  fromSockAddr (SockAddrInet _ h) = Just (fromHostAddress h)
+  fromSockAddr  _                 = Nothing
+
+-- | Note that port is zeroed.
+instance Address IPv6 where
+  toSockAddr h = SockAddrInet6 0 0 (toHostAddress6 h) 0
+  fromSockAddr (SockAddrInet6 _ _ h _) = Just (fromHostAddress6 h)
+  fromSockAddr  _                      = Nothing
+
+-- | Note that port is zeroed.
+instance Address IP where
+  toSockAddr (IPv4 h) = toSockAddr h
+  toSockAddr (IPv6 h) = toSockAddr h
+  fromSockAddr sa =
+        IPv4 <$> fromSockAddr sa
+    <|> IPv6 <$> fromSockAddr sa
+
+
+
 
 data MessageClass = Error | Query | Response
  deriving (Eq,Ord,Enum,Bounded,Data,Show,Read)
 
 class Envelope envelope where
     type TransactionID envelope
-    type NodeId envelope
+    data NodeId envelope
 
     envelopePayload     :: envelope a -> a
     envelopeTransaction :: envelope a -> TransactionID envelope
@@ -58,6 +109,187 @@ instance Serialize nodeid => Pretty (NodeDistance nodeid) where
   pPrint n = text $ show n
 
 
+-- | When 'get'ing an IP it must be 'isolate'd to the appropriate
+-- number of bytes since we have no other way of telling which
+-- address type we are trying to parse
+instance Serialize IP where
+    put (IPv4 ip) = put ip
+    put (IPv6 ip) = put ip
+
+    get = do
+      n <- remaining
+      case n of
+        4  -> IPv4 <$> get
+        16 -> IPv6 <$> get
+        _ -> fail (show n ++ " is the wrong number of remaining bytes to parse IP")
+
+instance Serialize IPv4 where
+    put = putWord32host    .  toHostAddress
+    get = fromHostAddress <$> getWord32host
+
+instance Serialize IPv6 where
+    put ip = put $ toHostAddress6 ip
+    get = fromHostAddress6 <$> get
+
+instance Pretty IPv4 where
+  pPrint = PP.text . show
+  {-# INLINE pPrint #-}
+
+instance Pretty IPv6 where
+  pPrint = PP.text . show
+  {-# INLINE pPrint #-}
+
+instance Pretty IP where
+  pPrint = PP.text . show
+  {-# INLINE pPrint #-}
+
+instance Hashable IPv4 where
+  hashWithSalt = hashUsing toHostAddress
+  {-# INLINE hashWithSalt #-}
+
+instance Hashable IPv6 where
+  hashWithSalt s a = hashWithSalt s (toHostAddress6 a)
+
+instance Hashable IP where
+  hashWithSalt s (IPv4 h) = hashWithSalt s h
+  hashWithSalt s (IPv6 h) = hashWithSalt s h
+
+
+
+
+
+data NodeAddr a = NodeAddr
+  { nodeHost ::                !a
+  , nodePort :: {-# UNPACK #-} !PortNumber
+  } deriving (Eq, Ord, Typeable, Functor, Foldable, Traversable)
+
+instance Show a => Show (NodeAddr a) where
+  showsPrec i NodeAddr {..}
+    = showsPrec i nodeHost <> showString ":" <> showsPrec i nodePort
+
+instance Read (NodeAddr IPv4) where
+  readsPrec i = RP.readP_to_S $ do
+    ipv4 <- RP.readS_to_P (readsPrec i)
+    _    <- RP.char ':'
+    port <- toEnum <$> RP.readS_to_P (readsPrec i)
+    return $ NodeAddr ipv4 port
+
+-- | @127.0.0.1:6882@
+instance Default (NodeAddr IPv4) where
+  def = "127.0.0.1:6882"
+
+-- | KRPC compatible encoding.
+instance Serialize a => Serialize (NodeAddr a) where
+  get = NodeAddr <$> get <*> get
+  {-# INLINE get #-}
+  put NodeAddr {..} = put nodeHost >> put nodePort
+  {-# INLINE put #-}
+
+-- | Example:
+--
+--   @nodePort \"127.0.0.1:6881\" == 6881@
+--
+instance IsString (NodeAddr IPv4) where
+  fromString str
+    | [hostAddrStr, portStr] <- splitWhen (== ':') str
+    , Just hostAddr <- readMaybe hostAddrStr
+    , Just portNum  <- toEnum <$> readMaybe portStr
+                = NodeAddr hostAddr portNum
+    | otherwise = error $ "fromString: unable to parse (NodeAddr IPv4): " ++ str
+
+instance Hashable PortNumber where
+  hashWithSalt s = hashWithSalt s . fromEnum
+  {-# INLINE hashWithSalt #-}
+
+instance Pretty PortNumber where
+  pPrint = PP.int . fromEnum
+  {-# INLINE pPrint #-}
+
+
+instance Hashable a => Hashable (NodeAddr a) where
+  hashWithSalt s NodeAddr {..} = hashWithSalt s (nodeHost, nodePort)
+  {-# INLINE hashWithSalt #-}
+
+instance Pretty ip => Pretty (NodeAddr ip) where
+  pPrint NodeAddr {..} = pPrint nodeHost <> ":" <> pPrint nodePort
+
+
+instance Serialize PortNumber where
+  get = fromIntegral <$> getWord16be
+  {-# INLINE get #-}
+  put = putWord16be . fromIntegral
+  {-# INLINE put #-}
+
+
+
+
+data NodeInfo dht addr u = NodeInfo
+  { nodeId   :: !(NodeId dht)
+  , nodeAddr :: !(NodeAddr addr)
+  , nodeAnnotation :: u
+  } deriving (Functor, Foldable, Traversable)
+
+deriving instance ( Show (NodeId dht)
+                  , Show addr
+                  , Show u ) => Show (NodeInfo dht addr u)
+
+mapAddress :: (addr -> b) -> NodeInfo dht addr u -> NodeInfo dht b u
+mapAddress f ni = ni { nodeAddr = fmap f (nodeAddr ni) }
+
+traverseAddress :: Applicative f => (addr -> f b) -> NodeInfo dht addr u -> f (NodeInfo dht b u)
+traverseAddress f ni = fmap (\addr -> ni { nodeAddr = addr }) $ traverse f (nodeAddr ni)
+
+-- Warning: Eq and Ord only look at the nodeId field.
+instance Eq (NodeId dht) => Eq (NodeInfo dht a u) where
+    a == b  = (nodeId a == nodeId b)
+
+instance Ord (NodeId dht) => Ord (NodeInfo dht a u) where
+  compare = comparing nodeId
+
+
+-- TODO WARN is the 'system' random suitable for this?
+-- | Generate random NodeID used for the entire session.
+--   Distribution of ID's should be as uniform as possible.
+--
+genNodeId :: forall dht.
+             ( Serialize (NodeId dht)
+             , FiniteBits (NodeId dht)
+             ) => IO (NodeId dht)
+genNodeId = either error id . S.decode <$> getEntropy nodeIdSize
+ where
+    nodeIdSize = finiteBitSize (undefined :: NodeId dht) `div` 8
+
+-- | Generate a random 'NodeId' within a range suitable for a bucket.  To
+-- obtain a sample for bucket number /index/ where /is_last/ indicates if this
+-- is for the current deepest bucket in our routing table:
+--
+-- > sample <- genBucketSample nid (bucketRange index is_last)
+genBucketSample :: ( FiniteBits (NodeId dht)
+                   , Serialize (NodeId dht)
+                   ) => NodeId dht -> (Int,Word8,Word8) -> IO (NodeId dht)
+genBucketSample n qmb = genBucketSample' getEntropy n qmb
+
+-- | Generalizion of 'genBucketSample' that accepts a byte generator
+-- function to use instead of the system entropy.
+genBucketSample' :: forall m dht.
+                    ( Applicative m
+                    , FiniteBits (NodeId dht)
+                    , Serialize (NodeId dht)
+                    ) =>
+    (Int -> m ByteString) -> NodeId dht -> (Int,Word8,Word8) -> m (NodeId dht)
+genBucketSample' gen self (q,m,b)
+    | q <= 0           =  either error id . S.decode <$> gen nodeIdSize
+    | q >= nodeIdSize  =  pure self
+    | otherwise        =  either error id . S.decode .  build <$> gen (nodeIdSize - q + 1)
+ where
+    nodeIdSize = finiteBitSize (undefined :: NodeId dht) `div` 8
+    build tl = BS.init hd <> BS.cons (h .|. t) (BS.tail tl)
+     where
+        hd = BS.take q $ S.encode self
+        h = xor b (complement m .&. BS.last hd)
+        t = m .&. BS.head tl
+
+
 class Envelope envelope => WireFormat raw envelope where
     type SerializableTo raw :: * -> Constraint
     type CipherContext raw envelope