diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2015, fro_ozen + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of fro_ozen nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/kademlia.cabal b/kademlia.cabal new file mode 100644 --- /dev/null +++ b/kademlia.cabal @@ -0,0 +1,52 @@ +name: kademlia +version: 1.0.0.0 +homepage: https://github.com/froozen/kademlia +bug-reports: https://github.com/froozen/kademlia/issues +synopsis: An implementation of the Kademlia DHT Protocol +description: + . + A haskell implementation of the Kademlia distributed hashtable, an efficient + way to store and lookup values distributed over a P2P network. + . + The implementation is based on the paper + /Kademlia: A Peer-to-peer Information System Based on the XOR Metric/: + ) + by Petar Maymounkov and David Mazières. + . + This library aims to be very simple and pleasant to use, with the downside of + deciding some of the implementation details, like timeout intervals and + k-bucket size, for the user. + +license: BSD3 +license-file: LICENSE +author: fro_ozen +maintainer: fro_ozen +category: Network + +build-type: Simple +cabal-version: >=1.10 + +source-repository head + type: git + location: https://github.com/froozen/kademlia.git + +library + exposed-modules: Network.Kademlia + + other-modules: Network.Kademlia.Networking, Network.Kademlia.Types, + Network.Kademlia.Protocol, Network.Kademlia.Instance, + Network.Kademlia.Protocol.Parsing, Network.Kademlia.Tree, + Network.Kademlia.ReplyQueue, + Network.Kademlia.Implementation + + build-depends: base >= 4 && < 5, + network >=2.6 && <2.7, + mtl >=2.1.3.1, + bytestring >=0.10 && <0.11, + transformers >=0.3, + containers >=0.5.5.1, + stm >=2.4.3, + transformers-compat >=0.3.3 + + hs-source-dirs: src + default-language: Haskell2010 diff --git a/src/Network/Kademlia.hs b/src/Network/Kademlia.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia.hs @@ -0,0 +1,145 @@ +{-| +Module : Network.Kademlia +Description : Implementation of the Kademlia DHT +License: BSD3 +Maintainer: fro_ozen@gmx.de +Stability: experimental +Portability: GHC + +A haskell implementation of the Kademlia distributed hashtable, an efficient +way to store and lookup values distributed over a P2P network. + +The implementation is based on the paper by Petar Maymounkov and David Mazières:
+/Kademlia: A Peer-to-peer Information System Based on the XOR Metric/: +() + +This library aims to be very simple and pleasant to use, with the downside of +deciding some of the implementation details, like timeout intervals and +k-bucket size, for the user. + += How to use it + +To get started with this library, first import it. The import has to be +qualified, as the module uses the same function names as some other modules. + +> import qualified Network.Kademlia as K + +Next, you need to decide on the types you want to use as the values to be stored +in the DHT and the keys to acces them by. As soon as you've decided on them, you +have to make them instances of the "Serialize" typeclass, so they can be sent over +the network. + +> import qualified Data.ByteString as B +> import qualified Data.ByteString.Char8 as C +> import Control.Arrow (first) +> +> -- The type this example will use as value +> type Person = data { +> age :: Int +> , name :: String +> } +> deriving (Show) +> +> instance K.Serialize Person where +> toBS = C.pack . show +> fromBS bs = +> case (reads :: ReadS Person) . C.unpack $ bs of +> [] -> Left "Failed to parse Person." +> (result, rest):_ -> Right (result, C.pack rest) +> +> -- The type this example will use as key for the lookups +> newtype KademliaID = KademliaID B.ByteString +> +> instance K.Serialize KademliaID where +> toBS (KademliaID bs) +> | B.length bs >= 5 = B.take 5 bs +> | otherwise = error "KademliaID to short!" +> +> fromBS bs +> | B.length bs >= 5 = Right . first KademliaID . B.splitAt 5 $ bs +> | otherwise = Left "ByteString too short!" +> + +As you could see in the example above, for the algorithm to work, you have to make +sure the serialized keys are of a fixed length. There is no such constraint for +the values. + +Now you're ready to dive in and use the DHT: + +> main = do +> -- Create the first instance, which will serve as the first node of the +> -- network +> firstInstance <- K.create 12345 . KademliaID . C.pack $ "hello" +> +> -- Create the second instance and make it join the network +> secondInstance <- K.create 12346 . KademliaID . C.pack $ "uAleu" +> K.joinNetwork secondInstance ("localhost", 12345, "hello") +> +> -- Store an example value in the network +> let exampleValue = Person 25 "Alan Turing" +> K.store secondInstance (KademliaID . C.pack $ "raxqT") exampleValue +> +> -- Look up the value +> result <- K.lookup firstInstance . KademliaID . C.pack $ "raxqT" +> print result +> +> -- Close the instances +> K.close firstInstance +> K.close secondInstance + +As promised, the usage of the actual DHT is rather easy. There are a few things +to note, though: + + * To join an existing network, you need to know the hostname, listening port + and id of a node that is already part of that network + * When you don't need access to the DHT anymore, make sure to close the instances. + This closes opened sockets and kills the threads running in the background + +Another thing to note is, that you are responsible for assigning ids to nodes +and keys to values, as well as making sure these are unique. The Kademlia paper +doesn't propose any measures for this and, as this library is just a +implementation of the system proposed in it, this library doesn't implement +anything to handle this. + +-} + +module Network.Kademlia + ( KademliaInstance + , create + , close + , I.lookup + , I.store + , Network.Kademlia.joinNetwork + , Serialize(..) + ) where + +import Network.Kademlia.Networking +import Network.Kademlia.Instance +import qualified Network.Kademlia.Tree as T +import Network.Kademlia.Types +import Network.Kademlia.ReplyQueue +import Network.Kademlia.Implementation as I +import Prelude hiding (lookup) +import Control.Monad (void, forM_) +import Control.Concurrent.Chan +import Control.Concurrent.STM + +-- | Create a new KademliaInstance corresponding to a given Id on a given port +create :: (Serialize i, Ord i, Serialize a, Eq a, Eq i) => + Int -> i -> IO (KademliaInstance i a) +create port id = do + h <- openOn (show port) id + inst <- newInstance id h + start inst + return inst + +-- | Make a KademliaInstance join the network the supplied Node is a part of +joinNetwork :: (Serialize i, Ord i, Eq i, Serialize a) => KademliaInstance i a + -> (String, Int, i) -> IO () +joinNetwork inst (host, port, i) = let peer = Peer host . fromIntegral $ port + node = Node peer i + in I.joinNetwork inst node + +-- | Stop a KademliaInstance by closing it +close :: KademliaInstance i a -> IO () +close = closeK . handle diff --git a/src/Network/Kademlia/Implementation.hs b/src/Network/Kademlia/Implementation.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia/Implementation.hs @@ -0,0 +1,310 @@ +{-| +Module : Network.Kademlia.Implementation +Description : The details of the lookup algorithm + +"Network.Kademlia.Implementation" contains the actual implementations of the +different Kademlia Network Algorithms. +-} + +module Network.Kademlia.Implementation + ( lookup + , store + , joinNetwork + ) where + +import Network.Kademlia.Networking +import Network.Kademlia.Instance +import qualified Network.Kademlia.Tree as T +import Network.Kademlia.Types +import Network.Kademlia.ReplyQueue +import Prelude hiding (lookup) +import Control.Monad (forM_, unless, when) +import Control.Monad.Trans.State hiding (state) +import Control.Concurrent.Chan +import Control.Concurrent.STM +import Control.Monad.IO.Class (liftIO) +import Data.List (delete, find, (\\)) +import Data.Maybe (isJust, fromJust) + + +-- Lookup the value corresponding to a key in the DHT +lookup :: (Serialize i, Serialize a, Eq i, Ord i) => KademliaInstance i a -> i + -> IO (Maybe a) +lookup inst id = runLookup go inst id + where go = startLookup sendS cancel checkSignal + + -- Return Nothing on lookup failure + cancel = return Nothing + + -- When receiving a RETURN_VALUE command, finish the lookup, then + -- cache the value in the closest peer that didn't return it and + -- finally return the value + checkSignal (Signal origin (RETURN_VALUE _ value)) = do + -- Abuse the known list for saving the peers that are *known* to + -- store the value + modify $ \s -> s { known = [origin] } + + -- Finish the lookup, recording which nodes returned the value + finish + + -- Store the value in the closest peer that didn't return the + -- value + known <- gets known + polled <- gets polled + let rest = polled \\ known + unless (null rest) $ do + let cachePeer = peer . head . sortByDistanceTo rest $ id + liftIO . send (handle inst) cachePeer . STORE id $ value + + -- Return the value + return . Just $ value + + -- When receiving a RETURN_NODES command, throw the nodes into the + -- lookup loop and continue the lookup + checkSignal (Signal _ (RETURN_NODES _ nodes)) = + continueLookup nodes sendS continue cancel + + -- Continuing always means waiting for the next signal + continue = waitForReply cancel checkSignal + + -- Send a FIND_VALUE command, looking for the supplied id + sendS = sendSignal (FIND_VALUE id) + + -- As long as there still are pending requests, wait for the next one + finish = do + pending <- gets pending + unless (null pending) $ waitForReply (return ()) finishCheck + + -- Record the nodes which return the value + finishCheck (Signal origin (RETURN_VALUE _ _)) = do + known <- gets known + modify $ \s -> s { known = origin:known } + finish + finishCheck _ = finish + +-- Store assign a value to a key and store it in the DHT +store :: (Serialize i, Serialize a, Eq i, Ord i) => + KademliaInstance i a -> i -> a -> IO () +store inst key val = runLookup go inst key + where go = startLookup sendS end checkSignal + + -- Always add the nodes into the loop and continue the lookup + checkSignal (Signal _ (RETURN_NODES _ nodes)) = + continueLookup nodes sendS continue end + + -- Continuing always means waiting for the next signal + continue = waitForReply end checkSignal + + -- Send a FIND_NODE command, looking for the node corresponding to the + -- key + sendS = sendSignal (FIND_NODE key) + + -- Run the lookup as long as possible, to make sure the nodes closest + -- to the key were polled. + end = do + polled <- gets polled + + unless (null polled) $ do + let h = handle inst + -- Select the peer closest to the key + storePeer = peer . head . sortByDistanceTo polled $ key + -- Send it a STORE command + liftIO . send h storePeer . STORE key $ val + +-- | Make a KademliaInstance join the network a supplied Node is in +joinNetwork :: (Serialize i, Serialize a, Eq i, Ord i) => KademliaInstance i a + -> Node i -> IO () +joinNetwork inst node = ownId >>= runLookup go inst + where go = do + -- Poll the supplied node + sendS node + -- Run a normal lookup from thereon out + waitForReply cancel checkSignal + + -- Do nothing upon failure or when the join operation has terminated + cancel = return () + + -- Retrieve your own id + ownId = + fmap T.extractId . atomically . readTVar . sTree . state $ inst + + -- Always add the nodes into the loop and continue the lookup + checkSignal (Signal _ (RETURN_NODES _ nodes)) = + continueLookup nodes sendS continue cancel + + -- Continuing always means waiting for the next signal + continue = waitForReply cancel checkSignal + + -- Send a FIND_NODE command, looking up your own id + sendS node = liftIO ownId >>= flip sendSignal node . FIND_NODE + +-- | The state of a lookup +data LookupState i a = LookupState { + inst :: KademliaInstance i a + , targetId :: i + , replyChan :: Chan (Reply i a) + , known :: [Node i] + , pending :: [Node i] + , polled :: [Node i] + } + +-- | MonadTransformer context of a lookup +type LookupM i a = StateT (LookupState i a) IO + +-- Run a LookupM, returning its result +runLookup :: LookupM i a b -> KademliaInstance i a -> i ->IO b +runLookup lookup inst id = do + chan <- newChan + let state = LookupState inst id chan [] [] [] + + evalStateT lookup state + +-- The initial phase of the normal kademlia lookup operation +startLookup :: (Serialize i, Serialize a, Eq i, Ord i) => (Node i -> LookupM i a ()) + -> LookupM i a b -> (Signal i a -> LookupM i a b) -> LookupM i a b +startLookup sendSignal cancel onSignal = do + inst <- gets inst + tree <- liftIO . atomically . readTVar . sTree . state $ inst + chan <- gets replyChan + id <- gets targetId + + -- Find the three nodes closest to the supplied id + case T.findClosest tree id 3 of + [] -> cancel + closest -> do + -- Send a signal to each of the Nodes + forM_ closest sendSignal + + -- Add them to the list of known nodes. At this point, it will + -- be empty, therfore just overwrite it. + modify $ \s -> s { known = closest } + + -- Start the recursive lookup + waitForReply cancel onSignal + +-- Wait for the next reply and handle it appropriately +waitForReply :: (Serialize i, Serialize a, Ord i) => LookupM i a b + -> (Signal i a -> LookupM i a b) -> LookupM i a b +waitForReply cancel onSignal = do + chan <- gets replyChan + sPending <- gets pending + known <- gets known + inst <- gets inst + polled <- gets polled + + result <- liftIO . readChan $ chan + case result of + -- If there was a reply + Answer sig@(Signal node _) -> do + -- Insert the node into the tree, as it might be a new one or it + -- would have to be refreshed + liftIO . insertNode inst $ node + + -- Remove the node from the list of nodes with pending replies + modify $ \s -> s { pending = delete node sPending } + + -- Call the signal handler + onSignal sig + + -- On timeout + Timeout registration -> do + let id = replyOrigin registration + + -- Find the node corresponding to the id + -- + -- ReplyQueue guarantees us, that it will be in polled, therefore + -- we can use fromJust + let node = fromJust . find (\n -> nodeId n == id) $ polled + + -- Remove every trace of the node's existance + liftIO . deleteNode inst $ id + modify $ \s -> s { + pending = delete node sPending + , known = delete node known + , polled = delete node polled + } + + -- Continue, if there still are pending responses + updatedPending <- gets pending + if not . null $ updatedPending + then waitForReply cancel onSignal + else cancel + + Closed -> cancel + +-- Decide wether, and which node to poll and react appropriately. +-- +-- This is the meat of kademlia lookups +continueLookup :: (Serialize i, Serialize a, Eq i) => [Node i] + -> (Node i -> LookupM i a ()) -> LookupM i a b -> LookupM i a b + -> LookupM i a b +continueLookup nodes sendSignal continue end = do + known <- gets known + id <- gets targetId + pending <- gets pending + polled <- gets polled + + -- Pick the k closest known nodes, that haven't been polled yet + let newKnown = take 7 . filter (`notElem` polled) $ nodes ++ known + + -- If there the k closest nodes haven't been polled yet + closestPolled <- closestPolled newKnown + if (not . null $ newKnown) && not closestPolled + then do + -- Send signal to the closest node, that hasn't + -- been polled yet + let next = head . sortByDistanceTo newKnown $ id + sendSignal next + + -- Update known + modify $ \s -> s { known = newKnown } + + -- Continue the lookup + continue + + -- If there are still pending replies + else if not . null $ pending + -- Wait for the pending replies to finish + then continue + -- Stop recursive lookup + else end + + where closestPolled known = do + polled <- gets polled + closest <- closest known + + return . all (`elem` polled) $ closest + + closest known = do + id <- gets targetId + polled <- gets polled + + -- Return the 7 closest nodes, the lookup had contact with + return . take 7 . sortByDistanceTo (known ++ polled) $ id + +-- Send a signal to a node +sendSignal :: (Serialize i, Serialize a, Eq i) => Command i a + -> Node i -> LookupM i a () +sendSignal cmd node = do + h <- fmap handle . gets $ inst + chan <- gets replyChan + polled <- gets polled + pending <- gets pending + + -- Send the signal + liftIO . send h (peer node) $ cmd + + -- Expect an appropriate reply to the command + liftIO . expect h regs $ chan + + -- Mark the node as polled and pending + modify $ \s -> s { + polled = node:polled + , pending = node:pending + } + + -- Determine the appropriate ReplyRegistrations to the command + where regs = case cmd of + (FIND_NODE id) -> RR [R_RETURN_NODES id] (nodeId node) + (FIND_VALUE id) -> + RR [R_RETURN_NODES id, R_RETURN_VALUE id] (nodeId node) diff --git a/src/Network/Kademlia/Instance.hs b/src/Network/Kademlia/Instance.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia/Instance.hs @@ -0,0 +1,175 @@ +{-| +Module : Network.Kademlia.Instance +Description : Implementation of the KademliaInstance type + +"Network.Kademlia.Instance" implements the KademliaInstance type, as well +as all the things that need to happen in the background to get a working +Kademlia instance. +-} + +module Network.Kademlia.Instance + ( KademliaInstance(..) + , KademliaState(..) + , start + , newInstance + , insertNode + , deleteNode + , lookupNode + ) where + +import Control.Concurrent +import Control.Concurrent.Chan +import Control.Concurrent.STM +import Control.Monad (void, forever, when, join, forM_, forever) +import Control.Monad.Trans +import Control.Monad.Trans.State +import Control.Monad.Trans.Reader +import Control.Monad.IO.Class (liftIO) +import System.IO.Error (catchIOError) +import qualified Data.Map as M +import Data.Maybe (catMaybes) +import Data.Function (on) + +import Network.Kademlia.Networking +import qualified Network.Kademlia.Tree as T +import Network.Kademlia.Types +import Network.Kademlia.ReplyQueue + +-- | The handle of a running Kademlia Node +data KademliaInstance i a = KI { + handle :: KademliaHandle i a + , state :: KademliaState i a + } + +-- | Representation of the data the KademliaProcess carries +data KademliaState i a = KS { + sTree :: TVar (T.NodeTree i) + , values :: TVar (M.Map i a) + } + +-- | Create a new KademliaInstance from an Id and a KademliaHandle +newInstance :: (Serialize i) => + i -> KademliaHandle i a -> IO (KademliaInstance i a) +newInstance id handle = do + tree <- atomically . newTVar . T.create $ id + values <- atomically . newTVar $ M.empty + return . KI handle . KS tree $ values + +insertNode :: (Serialize i, Ord i) => KademliaInstance i a -> Node i -> IO () +insertNode (KI _ (KS sTree _)) node = atomically $ do + tree <- readTVar sTree + writeTVar sTree . T.insert tree $ node + +deleteNode :: (Serialize i, Ord i) => KademliaInstance i a -> i -> IO () +deleteNode (KI _ (KS sTree _)) id = atomically $ do + tree <- readTVar sTree + writeTVar sTree . T.delete tree $ id + +lookupNode :: (Serialize i, Ord i) => KademliaInstance i a -> i -> IO (Maybe (Node i)) +lookupNode (KI _ (KS sTree _)) id = atomically $ do + tree <- readTVar sTree + return . T.lookup tree $ id + +insertValue :: (Ord i) => i -> a -> KademliaInstance i a -> IO () +insertValue key value (KI _ (KS _ values)) = atomically $ do + vals <- readTVar values + writeTVar values $ M.insert key value vals + +lookupValue :: (Ord i) => i -> KademliaInstance i a -> IO (Maybe a) +lookupValue key (KI _ (KS _ values)) = atomically $ do + vals <- readTVar values + return . M.lookup key $ vals + +-- | Start the background process for a KademliaInstance +start :: (Serialize i, Ord i, Serialize a, Eq i, Eq a) => + KademliaInstance i a -> IO () +start inst = do + chan <- newChan + startRecvProcess (handle inst) chan + pingId <- forkIO . pingProcess inst $ chan + spreadId <- forkIO . spreadValueProcess $ inst + void . forkIO $ backgroundProcess inst chan [pingId, spreadId] + +-- | The actual process running in the background +backgroundProcess :: (Serialize i, Ord i, Serialize a, Eq i, Eq a) => + KademliaInstance i a -> Chan (Reply i a) -> [ThreadId] -> IO () +backgroundProcess inst chan threadIds = do + reply <- liftIO . readChan $ chan + + case reply of + Answer sig -> do + let node = source sig + + -- Handle the signal + handleCommand (command sig) (peer node) inst + + -- Insert the node into the tree, if it's allready known, it will + -- be refreshed + insertNode inst node + + backgroundProcess inst chan threadIds + + -- Delete timed out nodes + Timeout registration -> deleteNode inst . replyOrigin $ registration + + -- Kill pingProcess and stop on Closed + Closed -> mapM_ killThread threadIds + +-- | Ping all known nodes every five minutes to make sure they are still present +pingProcess :: (Serialize i, Serialize a, Eq i) => KademliaInstance i a + -> Chan (Reply i a) -> IO () +pingProcess (KI h (KS sTree _)) chan = forever $ do + threadDelay fiveMinutes + + tree <- atomically . readTVar $ sTree + forM_ (allNodes tree) $ \node -> do + -- Send PING and expect a PONG + send h (peer node) PING + expect h (RR [R_PONG] (nodeId node)) $ chan + + where fiveMinutes = 300000000 + allNodes = join . catMaybes . map snd + +-- | Store all values stored in the node in the 7 closest known nodes every day +spreadValueProcess :: (Serialize i, Serialize a, Eq i) => KademliaInstance i a + -> IO () +spreadValueProcess (KI h (KS sTree sValues)) = forever $ do + threadDelay day + + values <- atomically . readTVar $ sValues + tree <- atomically . readTVar $ sTree + + mapMWithKey (sendRequests tree) $ values + + where day = 24 * 60 * 60 * 1000000 + sendRequests tree key val = do + let closest = T.findClosest tree key 7 + forM_ closest $ \node -> send h (peer node) (STORE key val) + + mapMWithKey :: (k -> v -> IO a) -> M.Map k v -> IO [a] + mapMWithKey f m = sequence . map snd . M.toList . M.mapWithKey f $ m + +-- | Handles the differendt Kademlia Commands appropriately +handleCommand :: (Serialize i, Eq i, Ord i, Serialize a) => + Command i a -> Peer -> KademliaInstance i a -> IO () +-- Simply answer a PING with a PONG +handleCommand PING peer inst = send (handle inst) peer PONG +-- Return a KBucket with the closest Nodes +handleCommand (FIND_NODE id) peer inst = returnNodes peer id inst +-- Insert the value into the values Map +handleCommand (STORE key value) _ inst = insertValue key value inst +-- Return the value, if known, or the closest other known Nodes +handleCommand (FIND_VALUE key) peer inst = do + result <- lookupValue key inst + case result of + Just value -> liftIO $ send (handle inst) peer $ RETURN_VALUE key value + Nothing -> returnNodes peer key inst +handleCommand _ _ _ = return () + +-- | Return a KBucket with the closest Nodes to a supplied Id +returnNodes :: (Serialize i, Eq i, Ord i, Serialize a) => + Peer -> i -> KademliaInstance i a -> IO () +returnNodes peer id (KI h (KS sTree _)) = do + tree <- atomically . readTVar $ sTree + let nodes = T.findClosest tree id 7 + liftIO $ send h peer (RETURN_NODES id nodes) diff --git a/src/Network/Kademlia/Networking.hs b/src/Network/Kademlia/Networking.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia/Networking.hs @@ -0,0 +1,137 @@ +{-| +Module : Network.Kademlia.Networking +Description : All of the UDP network code + +Network.Kademlia.Networking implements all the UDP network functionality. +-} + +module Network.Kademlia.Networking + ( openOn + , startRecvProcess + , send + , expect + , closeK + , KademliaHandle + ) where + +-- Just to make sure I'll only use the ByteString functions +import Network.Socket hiding (send, sendTo, recv, recvFrom, Closed) +import qualified Network.Socket.ByteString as S +import Data.ByteString +import Control.Monad (forever, unless) +import Control.Exception (finally) +import Control.Concurrent +import Control.Concurrent.STM +import Control.Concurrent.Chan +import Control.Concurrent.MVar +import System.IO.Error (ioError, userError) + +import Network.Kademlia.Types +import Network.Kademlia.Protocol +import Network.Kademlia.ReplyQueue + +-- | A handle to a UDP socket running the Kademlia connection +data KademliaHandle i a = KH { + kSock :: Socket + , sendThread :: ThreadId + , sendChan :: Chan (Command i a, Peer) + , replyQueue :: ReplyQueue i a + , recvThread :: MVar ThreadId + } + +-- | Open a Kademlia connection on specified port and return a corresponding +-- KademliaHandle +openOn :: (Serialize i, Serialize a) => String -> i -> IO (KademliaHandle i a) +openOn port id = withSocketsDo $ do + -- Get addr to bind to + (serveraddr:_) <- getAddrInfo + (Just (defaultHints {addrFlags = [AI_PASSIVE]})) + Nothing (Just port) + + -- Create socket and bind to it + sock <- socket (addrFamily serveraddr) Datagram defaultProtocol + bindSocket sock (addrAddress serveraddr) + + chan <- newChan + tId <- forkIO . sendProcess sock id $ chan + rq <- emptyReplyQueue + mvar <- newEmptyMVar + + -- Return the handle + return $ KH sock tId chan rq mvar + +sendProcess :: (Serialize i, Serialize a) => Socket -> i + -> Chan (Command i a, Peer) -> IO () +sendProcess sock id chan = (withSocketsDo . forever $ do + (cmd, Peer host port) <- readChan chan + + -- Get Peer's address + (peeraddr:_) <- getAddrInfo Nothing (Just host) + (Just . show . fromIntegral $ port) + + -- Send the signal + let sig = serialize id cmd + S.sendTo sock sig (addrAddress peeraddr)) + -- Close socket on exception (ThreadKilled) + `finally` sClose sock + +-- | Dispatch the receiving process +-- +-- Receive a signal and first try to dispatch it via the ReplyQueue. If that +-- fails, send it to the supplied default channel instead. +-- +-- This throws an exception if called a second time. +startRecvProcess :: (Serialize i, Serialize a, Eq i, Eq a) => KademliaHandle i a + -> Chan (Reply i a) -> IO () +startRecvProcess kh defaultChan = do + tId <- forkIO $ (withSocketsDo . forever $ do + -- Read from socket + (received, addr) <- S.recvFrom (kSock kh) 1500 + -- Try to create peer + peer <- toPeer addr + case peer of + Nothing -> return () + Just p -> + -- Try parsing the signal + case parse p received of + Left _ -> return () + Right sig -> do + -- Try to dispatch the signal + success <- dispatch sig $ replyQueue kh + + unless success $ + -- Send it to the default channel + writeChan defaultChan $ Answer sig) + + -- Send Closed reply to all handlers + `finally` do + flush . replyQueue $ kh + writeChan defaultChan Closed + + success <- tryPutMVar (recvThread kh) tId + unless success . ioError . userError $ "Receiving process already running" + +-- | Send a Signal to a Peer over the connection corresponding to the +-- KademliaHandle +send :: (Serialize i, Serialize a) => KademliaHandle i a -> Peer -> Command i a + -> IO () +send kh peer cmd = writeChan (sendChan kh) (cmd, peer) + +-- | Register a handler channel for a Reply +expect :: (Serialize i, Serialize a, Eq i) => KademliaHandle i a + -> ReplyRegistration i -> Chan (Reply i a) -> IO () +expect kh reg = register reg . replyQueue $ kh + +-- | Close the connection corresponding to a KademliaHandle +closeK :: KademliaHandle i a -> IO () +closeK kh = do + -- Kill sendThread + killThread . sendThread $ kh + + -- Kill recvThread + empty <- isEmptyMVar . recvThread $ kh + unless empty $ do + tId <- takeMVar . recvThread $ kh + killThread tId + + yield diff --git a/src/Network/Kademlia/Protocol.hs b/src/Network/Kademlia/Protocol.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia/Protocol.hs @@ -0,0 +1,57 @@ +{-| +Module : Network.Kademlia.Protocol +Description : Implementation of the actual protocol + +Network.Kademlia.Protocol implements the parsing and serialisation of +ByteStrings into 'Protocol'-Values. +-} + +module Network.Kademlia.Protocol + ( serialize + , parse + ) where + +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as L +import qualified Data.ByteString.Char8 as C +import Data.ByteString.Builder (toLazyByteString, word16BE) +import Data.Word (Word8) +import Data.List (foldl') + +import Network.Kademlia.Types +import Network.Kademlia.Protocol.Parsing + +-- | Retrieve the assigned protocolId +commandId :: Command i a -> Word8 +commandId PING = 0 +commandId PONG = 1 +commandId (STORE _ _) = 2 +commandId (FIND_NODE _) = 3 +commandId (RETURN_NODES _ _) = 4 +commandId (FIND_VALUE _) = 5 +commandId (RETURN_VALUE _ _) = 6 + +-- | Turn the command arguments into a ByteString +commandArgs :: (Serialize i, Serialize a) => Command i a -> B.ByteString +commandArgs PING = B.empty +commandArgs PONG = B.empty +commandArgs (STORE k v) = toBS k `B.append` toBS v +commandArgs (FIND_NODE id) = toBS id +commandArgs (FIND_VALUE k) = toBS k +commandArgs (RETURN_VALUE id v) = toBS id `B.append` toBS v +commandArgs (RETURN_NODES id kb) = toBS id `B.append` + foldl' B.append B.empty (fmap nodeToArg kb) + +nodeToArg :: (Serialize i) => Node i -> B.ByteString +nodeToArg node = id `B.append` C.pack (host ++ " ") `B.append` port + where id = toBS . nodeId $ node + host = peerHost . peer $ node + port = toBinary . fromIntegral . peerPort . peer $ node + -- Converts a Word16 into a two character ByteString + toBinary = B.concat . L.toChunks . toLazyByteString . word16BE + +-- | Turn a command into a sendable ByteString +serialize :: (Serialize i, Serialize a) => i -> Command i a -> B.ByteString +serialize id command = cId `B.cons` toBS id `B.append` args + where cId = commandId command + args = commandArgs command diff --git a/src/Network/Kademlia/Protocol/Parsing.hs b/src/Network/Kademlia/Protocol/Parsing.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia/Protocol/Parsing.hs @@ -0,0 +1,131 @@ +{-| +Module : Network.Kademlia.Protocol.Parsing +Description : Implementation of the protocol parsing + +Network.Kademlia.Protocol.Parsing implements the actual protocol parsing. + +It made sense to split it off Network.Kademlia.Protocol as it made both cleaner +and more readable. +-} + +module Network.Kademlia.Protocol.Parsing where + +import qualified Data.ByteString as B +import qualified Data.ByteString.Char8 as C +import Control.Monad (liftM, liftM2) +import Control.Monad.Trans (lift) +import Control.Monad.State +import Control.Monad.Trans.Except +import Text.Read (readMaybe) +import Data.Word (Word8, Word16) +import Data.Bits (shiftL) + +import Network.Kademlia.Types + +type Parse = ExceptT String (State B.ByteString) + +-- | Parse a signal from a ByteString +-- +-- (This needs to be supplied a Peer, to be able to create a complete Signal) +parse :: (Serialize i, Serialize a) => Peer -> B.ByteString -> Either String (Signal i a) +parse peer = evalState (runExceptT $ parseSignal peer) + +-- | Parses the parsable parts of a signal +parseSignal :: (Serialize i, Serialize a) => Peer -> Parse (Signal i a) +parseSignal peer = do + cId <- parseCommandId + id <- parseSerialize + cmd <- parseCommand cId + let node = Node peer id + return $ Signal node cmd + +-- | Parses a Serialize +parseSerialize :: (Serialize a) => Parse a +parseSerialize = do + bs <- lift get + case fromBS bs of + Left err -> throwE err + Right (id, rest) -> do + lift . put $ rest + return id + +-- | Parses a CommandId +parseCommandId :: Parse Int +parseCommandId = do + bs <- lift get + case B.uncons bs of + Nothing -> throwE "uncons returned Nothing" + Just (id, rest) -> do + lift . put $ rest + return $ fromIntegral id + +-- | Splits after a certain character +parseSplit :: Char -> Parse B.ByteString +parseSplit c = do + bs <- lift get + if B.null bs + then throwE "ByteString empty" + else do + let (result, rest) = C.span (/=c) bs + lift . put $ rest + return result + +-- | Skips one character +skipCharacter :: Parse () +skipCharacter = do + bs <- lift get + if B.null bs + then throwE "ByteString empty" + else lift . put $ B.drop 1 bs + +-- | Parses an Int +parseInt :: Parse Int +parseInt = do + bs <- lift get + case C.readInt bs of + Nothing -> throwE "Failed to parse an Int" + Just (n, rest) -> do + lift . put $ rest + return n + +-- | Parses two Word8s from a ByteString into one Word16 +parseWord16 :: Parse Word16 +parseWord16 = do + bs <- lift get + if B.length bs < 2 + then throwE "ByteString to short" + else do + let (words, rest) = B.splitAt 2 bs + lift . put $ rest + return . joinWords . B.unpack $ words + where + joinWords [a, b] = (toWord16 a `shiftL` 8) + toWord16 b + + toWord16 :: Word8 -> Word16 + toWord16 = fromIntegral + +-- | Parses a Node's info +parseNode :: (Serialize i) => Parse (Node i) +parseNode = do + id <- parseSerialize + host <- parseSplit ' ' + skipCharacter + port <- parseWord16 + let peer = Peer (C.unpack host) (fromIntegral port) + return $ Node peer id + +-- | Parses a trailing k-bucket +parseKBucket :: (Serialize i) => Parse (KBucket i) +parseKBucket = liftM2 (:) parseNode parseKBucket + `catchE` \_ -> return [] + +-- | Parses the rest of a command corresponding to an id +parseCommand :: (Serialize i, Serialize a) => Int -> Parse (Command i a) +parseCommand 0 = return PING +parseCommand 1 = return PONG +parseCommand 2 = liftM2 STORE parseSerialize parseSerialize +parseCommand 3 = FIND_NODE `liftM` parseSerialize +parseCommand 4 = liftM2 RETURN_NODES parseSerialize parseKBucket +parseCommand 5 = FIND_VALUE `liftM` parseSerialize +parseCommand 6 = liftM2 RETURN_VALUE parseSerialize parseSerialize +parseCommand _ = throwE "Invalid id" diff --git a/src/Network/Kademlia/ReplyQueue.hs b/src/Network/Kademlia/ReplyQueue.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia/ReplyQueue.hs @@ -0,0 +1,142 @@ +{-| +Module : Network.Kademlia.ReplyQueue +Description : A queue allowing to register handlers for expected replies + +Network.Kademlia.ReplyQueue implements a Queue designed for registering +handlers for expected replies. + +The handlers are represented by unbound channels from Control.Concurrency.Chan. +-} + +module Network.Kademlia.ReplyQueue + ( ReplyType(..) + , ReplyRegistration(..) + , Reply(..) + , ReplyQueue + , emptyReplyQueue + , register + , dispatch + , flush + ) where + +import Control.Concurrent +import Control.Concurrent.STM +import Control.Concurrent.Chan +import Control.Monad (liftM, forM_) +import Control.Monad.Trans.Maybe +import Data.List (find, delete) + +import Network.Kademlia.Types + +-- | The different types a replied signal could possibly have. +-- +-- Note that these are only those Command types, which are replies to some +-- sort of request. Therefore, most Command types aren't contained in this +-- type. +data ReplyType i = R_PONG + | R_RETURN_VALUE i + | R_RETURN_NODES i + deriving (Eq) + +-- | The representation of registered replies +data ReplyRegistration i = RR { + replyTypes :: [ReplyType i] + , replyOrigin :: i + } deriving (Eq) + +-- | Convert a Signal into its ReplyRegistration representation +toRegistration :: Signal i a -> Maybe (ReplyRegistration i) +toRegistration sig = case rType . command $ sig of + Nothing -> Nothing + Just rt -> Just (RR [rt] origin) + where origin = nodeId . source $ sig + + rType :: Command i a -> Maybe (ReplyType i) + rType PONG = Just R_PONG + rType (RETURN_VALUE id _) = Just (R_RETURN_VALUE id) + rType (RETURN_NODES id _) = Just (R_RETURN_NODES id) + rType _ = Nothing + +-- | Compare wether two ReplyRegistrations match +matchRegistrations :: (Eq i) => ReplyRegistration i -> ReplyRegistration i -> Bool +matchRegistrations (RR rtsA idA) (RR rtsB idB) = + idA == idB && (all (`elem` rtsA) rtsB || all (`elem` rtsB) rtsA) + +-- | The actual type of a replay +data Reply i a = Answer (Signal i a) + | Timeout (ReplyRegistration i) + | Closed + deriving (Eq) + +-- | The actual type representing a ReplyQueue +newtype ReplyQueue i a = RQ (TVar [(ReplyRegistration i, Chan (Reply i a), ThreadId)]) + +-- | Create an empty ReplyQueue +emptyReplyQueue :: IO (ReplyQueue i a) +emptyReplyQueue = atomically . liftM RQ $ newTVar [] + +-- | Register a channel as handler for a reply +register :: (Eq i) => ReplyRegistration i -> ReplyQueue i a -> Chan (Reply i a) + -> IO () +register reg (RQ rq) chan = do + tId <- timeoutThread chan reg (RQ rq) + atomically $ do + queue <- readTVar $ rq + writeTVar rq $ queue ++ [(reg, chan, tId)] + +timeoutThread :: (Eq i) => Chan (Reply i a) -> ReplyRegistration i + -> ReplyQueue i a -> IO ThreadId +timeoutThread chan reg (RQ rq) = forkIO $ do + -- Wait 5 seconds + threadDelay 5000000 + + -- Remove the ReplyRegistration from the ReplyQueue + myTId <- myThreadId + atomically $ do + queue <- readTVar $ rq + case find (\(_, _, tId) -> tId == myTId) queue of + Just rqElem -> writeTVar rq $ delete rqElem queue + _ -> return () + + -- Send Timeout signal + writeChan chan . Timeout $ reg + +-- | Try to send a received Signal over the registered handler channel and +-- return wether it succeeded +dispatch :: (Eq i) => Signal i a -> ReplyQueue i a -> IO Bool +dispatch sig (RQ rq) = do + result <- atomically $ do + queue <- readTVar $ rq + case toRegistration sig of + Just regA -> case find (matches regA) queue of + Just reg -> do + -- Remove registration from queue + writeTVar rq $ delete reg queue + return . Just $ reg + + Nothing -> return Nothing + Nothing -> return Nothing + + case result of + Just (_, chan, tId) -> do + -- Kill the timeout thread + killThread tId + + -- Send the signal + writeChan chan $ Answer sig + return True + _ -> return False + + where matches regA (regB, _, _) = matchRegistrations regA regB + +-- | Send Closed signal to all handlers and empty ReplyQueue +flush :: ReplyQueue i a -> IO () +flush (RQ rq) = do + queue <- atomically $ do + queue <- readTVar $ rq + writeTVar rq $ [] + return queue + + forM_ queue $ \(_, chan, tId) -> do + killThread tId + writeChan chan Closed diff --git a/src/Network/Kademlia/Tree.hs b/src/Network/Kademlia/Tree.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia/Tree.hs @@ -0,0 +1,202 @@ +{-| +Module : Network.Kademlia.Tree +Description : Implementation of the Node Storage Tree + +Network.Kademlia.Tree implements the Node Storage Tree used to store +and look up the known nodes. + +This module is designed to be used as a qualified import. +-} + +module Network.Kademlia.Tree + ( NodeTree + , create + , insert + , lookup + , delete + , refresh + , findClosest + , extractId + ) where + +import Network.Kademlia.Types +import qualified Data.List as L (delete, find) +import Prelude hiding (lookup, split) +import Control.Monad (liftM) +import Control.Arrow (first, second) +import Data.Function (on) + +-- | Type used for building the Node Storage Tree +type NodeTree i = [(Bool, Maybe (KBucket i))] + +-- | Structure used for easier modification of the NodeTree +type Zipper i = (NodeTree i, NodeTree i) + +-- | Move the Zipper along an Id +seek :: (Serialize i) => NodeTree i -> i -> Zipper i +seek tree id = go tree $ toByteStruct id + where go [] _ = ([], []) + go (pair@(bit, bucket):rest) (b:bs) + | ends rest = ([], pair:rest) + | bit == b = first (pair:) $ go rest bs + | otherwise = ([], pair:rest) + +-- | Cheks wether a NodeTree ends +ends :: NodeTree i -> Bool +ends ((_, Just _):_) = False +ends ((_, Nothing):_) = True + +-- | Apply a function to the KBucket a Node with a given Id would be in +applyTo :: (Serialize i, Eq i) => + (KBucket i -> a) -- ^ Function to apply at matched position + -> a -- ^ Default value + -> NodeTree i -- ^ NodeTree to apply to + -> i -- ^ Position to apply at + -> a +applyTo f end tree id = case seek tree id of + (_, []) -> end + (_, (_, Nothing):_) -> end + (_, (_, Just bucket):_) -> f bucket + +-- | Modify a NodeTree at the position a Node with a given Id would have +modifyTreeAt :: (Serialize i, Eq i) => + ((Bool, Maybe (KBucket i)) -> (Bool, Maybe (KBucket i))) + -- ^ Function to apply to corresponding TreeNode + -> NodeTree i -- ^ NodeTree to modify + -> i -- ^ Position to modify at + -> NodeTree i +modifyTreeAt f tree id = case seek tree id of + (beg, []) -> beg + (beg, pair:end) -> beg ++ f pair : end + +-- | Modify the KBucket a node of a given Id would be in +modifyKBucket :: (Serialize i, Eq i) => + (KBucket i -> KBucket i) -- ^ Modification funciton + -> NodeTree i -- ^ Node tree to modify + -> i -- ^ Postition to modify at + -> NodeTree i +modifyKBucket f = modifyTreeAt (second . fmap $ f) + +-- | Create a NodeTree corresponding to the Owner-Node's Id +create :: (Serialize i) => i -> NodeTree i +create id = zip (toByteStruct id) (repeat Nothing) + +-- | Insert a node into a NodeTree +insert :: (Serialize i, Eq i, Ord i) => NodeTree i -> Node i -> NodeTree i +insert tree node = case seek tree . nodeId $ node of + -- The tree is empty, create first KBucket + (_, (b, Nothing):xs) -> (b, Just [node]):xs + + -- Normal case + (beg, (b, Just bucket):xs) + -- At least refresh the Node, as it has been active + | node `elem` bucket -> refresh tree . nodeId $ node + -- The last bucket may always be split + | full bucket && ends xs -> let new = split tree . extractId $ tree + in insert new node + -- If the bucket is full and can't be split, the Node isn't inserted + | full bucket -> tree + -- Just insert the Node + | otherwise -> beg ++ (b, Just $ node:bucket):xs + + where full b = length b >= 7 + +-- Extract original Id from NodeTree +extractId :: (Serialize i) => NodeTree i -> i +extractId tree = fromByteStruct bs + where bs = foldr (\x id -> fst x:id) [] tree + +-- | Split the last bucket +-- +-- This function does some quite unsafe pattern matching for the sake of not +-- ending up even longer than it already is. It is only used internally and +-- all the assumptions made by those patterns are provable, so it's ok. +split :: (Serialize i, Ord i) => NodeTree i -> i -> NodeTree i +split tree id = let (begin, (b, Just bucket):xs) = seek tree id + (this, next) = doSplit bucket + in begin ++ (b, Just this) : injectBucket next xs + + where doSplit [] = ([], []) + doSplit (node:ns) = + -- More matching bytes than the index means that a node can be + -- moved to a later bucket. + if countMatching (toByteStruct . nodeId $ node) + (toByteStruct id) > index + then second (node:) $ doSplit ns + else first (node:) $ doSplit ns + + index = let (beg, _) = seek tree id in length beg + countMatching [] [] = 0 + countMatching (a:as) (b:bs) + | a == b = 1 + countMatching as bs + | otherwise = 0 + + injectBucket bucket ((b, _):xs) = (b, Just bucket):xs + +-- | Lookup a node within a NodeTree +lookup :: (Serialize i, Eq i) => NodeTree i -> i -> Maybe (Node i) +lookup tree id = applyTo f Nothing tree id + where f = L.find $ idMatches id + +-- | Delete a Node corresponding to a supplied Id from a NodeTree +delete :: (Serialize i, Eq i) => NodeTree i -> i -> NodeTree i +delete tree id = modifyKBucket f tree id + where f = filter $ not . idMatches id + +-- | Refresh the node corresponding to a supplied Id by placing it at the first +-- index of it's KBucket +refresh :: (Serialize i, Eq i) => NodeTree i -> i -> NodeTree i +refresh tree id = modifyKBucket f tree id + where f bucket = case L.find (idMatches id) bucket of + Just node -> node : L.delete node bucket + _ -> bucket + +-- | Find the k closest Nodes to a given Id +-- +-- Uset to implemenet RETURN_NODES +findClosest :: (Serialize i, Eq i) => NodeTree i -> i -> Int -> KBucket i +findClosest tree id n = case seek tree id of + -- The tree is empty + (_, (_, Nothing):xs) -> [] + + -- Normal case + (beg, (_, Just bk):xs) + -- The bucket contains enough Nodes on its own + | length bk == n -> bk + -- We need to retrieve Nodes from other buckets as well + | otherwise -> let + missing = n - length bk + in if ends xs + -- If it's the last one, take nodes from higher up in + -- the hierarchy + then let higher = next missing $ reverse beg + in take n . flip sortByDistanceTo id $ bk ++ higher + -- Else retrieve the missing amount of Nodes by calling + -- findClosest with an Id whose first differing bit doesn't + -- differ. + -- (Sounds complicated, but the tests prove that it actually + -- works this way) + else let treeId = extractId tree + newId = id `alignedTo` treeId + other = findClosest tree newId missing + in bk ++ other + where -- Pick the n closest Nodes from the tree + next _ [] = [] + next n ((_, Nothing):xs) = next n xs + next n ((_, Just bk):xs) + | length bk == n = bk + | length bk < n = bk ++ next (n - length bk) xs + -- Take the n closest Nodes + | otherwise = take n . sortByDistanceTo bk $ id + + -- Change the first differing bit of idA to match idB + idA `alignedTo` idB = fromByteStruct . alignF idA $ idB + alignF = align `on` toByteStruct + align [] [] = [] + align (a:as) (b:bs) + | a == b = a : align as bs + | otherwise = b : as + +-- | Helper function used for KBucket manipulation +idMatches :: (Eq i) => i -> Node i -> Bool +idMatches id node = id == nodeId node diff --git a/src/Network/Kademlia/Types.hs b/src/Network/Kademlia/Types.hs new file mode 100644 --- /dev/null +++ b/src/Network/Kademlia/Types.hs @@ -0,0 +1,109 @@ +{-| +Module : Network.Kademlia.Types +Description : Definitions of a few types + +Network.Kademlia.Types defines a few types that are used throughout the +Network.Kademlia codebase. +-} + +module Network.Kademlia.Types + ( Peer(..) + , toPeer + , Node(..) + , KBucket + , sortByDistanceTo + , Serialize(..) + , Signal(..) + , Command(..) + , ByteStruct(..) + , toByteStruct + , fromByteStruct + , distance + ) where + +import Network.Socket (SockAddr(..), PortNumber, inet_ntoa, inet_addr) +import qualified Data.ByteString as B (ByteString, foldr, pack) +import Data.Bits (testBit, setBit, zeroBits) +import Data.List (sortBy) +import Data.Function (on) + +-- | Representation of an UDP peer +data Peer = Peer { + peerHost :: String + , peerPort :: PortNumber + } deriving (Eq, Ord, Show) + +-- | Representation of a Kademlia Node, containing a Peer and an Id +data Node i = Node { + peer :: Peer + , nodeId :: i + } deriving (Eq, Ord, Show) + +-- | Aliases to make the code more readable by using the same names as the +-- papers +type KBucket i = [Node i] + +-- | Sort a bucket by the closeness of its nodes to a give Id +sortByDistanceTo :: (Serialize i) => KBucket i -> i -> KBucket i +sortByDistanceTo bucket id = unpack . sort . pack $ bucket + where pack bk = zip bk $ map f bk + f = distance id . nodeId + sort = sortBy (compare `on` snd) + unpack = map fst + +-- | A structure serializable into and parsable from a ByteString +class Serialize a where + fromBS :: B.ByteString -> Either String (a, B.ByteString) + toBS :: a -> B.ByteString + +-- | A Structure made up of bits, represented as a list of Bools +type ByteStruct = [Bool] + +-- | Converts a Serialize into a ByteStruct +toByteStruct :: (Serialize a) => a -> ByteStruct +toByteStruct s = B.foldr (\w bits -> convert w ++ bits) [] $ toBS s + where convert w = foldr (\i bits -> testBit w i : bits) [] [0..7] + +-- | Convert a ByteStruct back to its ByteString form +fromByteStruct :: (Serialize a) => ByteStruct -> a +fromByteStruct bs = case fromBS s of + (Right (converted, _)) -> converted + (Left err) -> error $ "Failed to convert from ByteStruct: " ++ err + where s = B.pack . foldr (\i ws -> createWord i : ws) [] $ indexes + indexes = [0..(length bs `div` 8) -1] + createWord i = let pos = i * 8 + in foldr changeBit zeroBits [pos..pos+7] + + changeBit i w = if bs !! i + then setBit w (i `mod` 8) + else w + +-- Calculate the distance between two Ids, as specified in the Kademlia paper +distance :: (Serialize i) => i -> i -> ByteStruct +distance idA idB = let bsA = toByteStruct idA + bsB = toByteStruct idB + in zipWith xor bsA bsB + where xor a b = not (a && b) && (a || b) + +-- | Try to convert a SockAddr to a Peer +toPeer :: SockAddr -> IO (Maybe Peer) +toPeer (SockAddrInet port host) = do + hostname <- inet_ntoa host + return $ Just $ Peer hostname port +toPeer _ = return Nothing + +-- | Representation of a protocl signal +data Signal i v = Signal { + source :: Node i + , command :: Command i v + } deriving (Show, Eq) + +-- | Representations of the different protocol commands +data Command i a = PING + | PONG + | STORE i a + | FIND_NODE i + | RETURN_NODES i (KBucket i) + | FIND_VALUE i + | RETURN_VALUE i a + deriving (Eq, Show)