4 files changed, 83 insertions, 27 deletions

diff --git a/DHTHandlers.hs b/DHTHandlers.hs
index 437b05f3..7ff7a3ce 100644
--- a/DHTHandlers.hs
+++ b/DHTHandlers.hs
@@ -26,6 +26,7 @@ import Data.IP
 import Data.Ord
 import Data.Maybe
 import Data.Bits
+import System.IO
 
 data TransactionId = TransactionId
  { transactionKey :: Nonce8  -- ^ Used to lookup pending query.
@@ -95,8 +96,9 @@ prefer4or6 addr iptyp = fromMaybe (ipFamily $ nodeIP addr) iptyp
 toxSpace :: R.KademliaSpace NodeId NodeInfo
 toxSpace = R.KademliaSpace
     { R.kademliaLocation = nodeId
-    , R.kademliaTestBit  = testIdBit
-    , R.kademliaXor      = xor
+    , R.kademliaTestBit  = testNodeIdBit
+    , R.kademliaXor      = xorNodeId
+    , R.kademliaSample   = sampleNodeId
     }
 
 
@@ -133,7 +135,7 @@ type Client = QR.Client String PacketKind TransactionId NodeInfo Message
 
 wrapAssym :: TransactionId -> NodeInfo -> NodeInfo -> (Nonce8 -> dta) -> Assym dta
 wrapAssym (TransactionId n8 n24) src dst dta = Assym
-    { senderKey = let NodeId pubkey = nodeId src in pubkey
+    { senderKey = id2key $ nodeId src
     , assymNonce = n24
     , assymData = dta n8
     }
@@ -158,7 +160,9 @@ unpong _               = Nothing
 
 ping :: Client -> NodeInfo -> IO Bool
 ping client addr = do
+    hPutStrLn stderr $ show addr ++ " <-- ping"
     reply <- QR.sendQuery client (serializer PingType DHTPing unpong) Ping addr
+    hPutStrLn stderr $ show addr ++ " -pong-> " ++ show reply
     maybe (return False) (\Pong -> return True) $ join reply
 
 unsendNodes :: Message -> Maybe (Assym (Nonce8,SendNodes))
@@ -170,12 +174,14 @@ unwrapNodes (SendNodes ns) = (ns,ns,())
 
 getNodes :: Client -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],()))
 getNodes client nid addr = do
+    hPutStrLn stderr $ show addr ++ " <-- getnodes " ++ show nid
     reply <- QR.sendQuery client (serializer GetNodesType DHTGetNodes unsendNodes) (GetNodes nid) addr
+    hPutStrLn stderr $ show addr ++ " -sendnodes-> " ++ show reply
     return $ fmap unwrapNodes $ join reply
 
 updateRouting :: Client -> Routing -> NodeInfo -> Message -> IO ()
 updateRouting client routing naddr msg = do
-    -- hPutStrLn stderr $ "updateRouting "++show typ
+    hPutStrLn stderr $ "updateRouting "++show (fst $ dhtMessageType $ fst $ DHTTransport.encrypt (error "updateRouting") msg naddr )
     -- TODO: check msg type
     case prefer4or6 naddr Nothing of
         Want_IP4 -> updateTable client naddr (routing4 routing) (committee4 routing) (sched4 routing)
diff --git a/DHTTransport.hs b/DHTTransport.hs
index 3d008174..013fa322 100644
--- a/DHTTransport.hs
+++ b/DHTTransport.hs
@@ -21,6 +21,7 @@ module DHTTransport
     , mapMessage
     , encrypt
     , decrypt
+    , dhtMessageType
     ) where
 
 import ToxAddress
@@ -122,7 +123,7 @@ direct bs saddr f = fanGet bs getDHT f (asymNodeInfo saddr)
 
 -- Throws an error if called with a non-internet socket.
 asymNodeInfo :: SockAddr -> Assym a -> NodeInfo
-asymNodeInfo saddr asym = either (error . mappend "asymNodeInfo: ") id $ nodeInfo (NodeId $ senderKey asym) saddr
+asymNodeInfo saddr asym = either (error . mappend "asymNodeInfo: ") id $ nodeInfo (key2id $ senderKey asym) saddr
 
 
 fanGet :: ByteString -> Get x -> (x -> a) -> (x -> b) -> Either String (a,b)
diff --git a/OnionTransport.hs b/OnionTransport.hs
index e05e2ca0..6901038d 100644
--- a/OnionTransport.hs
+++ b/OnionTransport.hs
@@ -108,7 +108,7 @@ instance Serialize (OnionMessage Encrypted) where
 
 onionToOwner :: Assym a -> ReturnPath N3 -> SockAddr -> Either String OnionToOwner
 onionToOwner assym ret3 saddr = do
-    ni <- nodeInfo (NodeId $ senderKey assym) saddr
+    ni <- nodeInfo (key2id $ senderKey assym) saddr
     return $ OnionToOwner ni ret3
 -- data CookieAddress = WithoutCookie NodeInfo | CookieAddress Cookie SockAddr
 
diff --git a/ToxAddress.hs b/ToxAddress.hs
index c95c221b..6a724d0f 100644
--- a/ToxAddress.hs
+++ b/ToxAddress.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE ApplicativeDo              #-}
 {-# LANGUAGE BangPatterns               #-}
 {-# LANGUAGE CPP                        #-}
 {-# LANGUAGE DataKinds                  #-}
@@ -13,7 +14,18 @@
 {-# LANGUAGE ScopedTypeVariables        #-}
 {-# LANGUAGE TupleSections              #-}
 {-# LANGUAGE TypeApplications           #-}
-module ToxAddress (NodeInfo(..),NodeId(..),nodeInfo,nodeAddr,zeroID,key2id,id2key,getIP) where
+module ToxAddress
+    ( NodeInfo(..)
+    , NodeId
+    , nodeInfo
+    , nodeAddr
+    , zeroID
+    , key2id
+    , id2key
+    , getIP
+    , xorNodeId
+    , testNodeIdBit
+    , sampleNodeId) where
 
 import Control.Applicative
 import Control.Monad
@@ -43,6 +55,9 @@ import qualified Text.ParserCombinators.ReadP as RP
 import Text.Read
 import Data.Bits
 import ToxCrypto
+import Foreign.Ptr
+import Data.Function
+import System.Endian
 
 -- | perform io for hashes that do allocation and ffi.
 -- unsafeDupablePerformIO is used when possible as the
@@ -56,22 +71,34 @@ unsafeDoIO = unsafeDupablePerformIO
 unsafeDoIO = unsafePerformIO
 #endif
 
-unpackPublicKey :: PublicKey -> [Word64]
+unpackPublicKey :: ByteArrayAccess bs => bs -> [Word64]
 unpackPublicKey bs = loop 0
   where loop i
-            | i == 4    = []
+            | i == (BA.length bs `div` 8) = []
             | otherwise =
-                let !v = unsafeDoIO $ BA.withByteArray bs (\p -> peekElemOff p i)
+                let !v = unsafeDoIO $ BA.withByteArray bs (\p -> fromBE64 <$> peekElemOff p i)
                  in v : loop (i+1)
 
-newtype NodeId = NodeId PublicKey
- deriving (Eq,ByteArrayAccess) -- (Eq,Ord,ByteArrayAccess, Bits, Hashable)
+packPublicKey :: BA.ByteArray bs => [Word64] -> bs
+packPublicKey ws = BA.allocAndFreeze (8 * length ws) $
+    flip fix ws $ \loop ys ptr -> case ys of
+        [] -> return ()
+        x:xs -> do poke ptr (toBE64 x)
+                   loop xs (plusPtr ptr 8)
+
+newtype NodeId = NodeId [Word64]
+ deriving (Eq,Ord) -- ByteArrayAccess) -- (Eq,Ord,ByteArrayAccess, Bits, Hashable)
+
 
 key2id :: PublicKey -> NodeId
-key2id = NodeId
+key2id = NodeId . unpackPublicKey
+
+bs2id :: ByteString -> NodeId
+bs2id bs = NodeId . unpackPublicKey $ throwCryptoError . publicKey $ bs
 
 id2key :: NodeId -> PublicKey
-id2key (NodeId key) = key
+id2key (NodeId key) = throwCryptoError . publicKey $ (packPublicKey key :: BA.Bytes)
+
 
 {-
 id2key :: NodeId -> PublicKey
@@ -87,33 +114,55 @@ key2id pk = case S.decode (BA.convert pk) of
 
 -}
 
+{-
 instance Ord NodeId where
     compare (NodeId a) (NodeId b) = compare (unpackPublicKey a) (unpackPublicKey b)
+-}
 
 zeroID :: NodeId
-zeroID = NodeId $ throwCryptoError $ publicKey $ B.replicate 32 0
+zeroID = NodeId $ replicate 4 0 -- throwCryptoError $ publicKey $ B.replicate 32 0
 
 instance Read NodeId where
     readsPrec _ str
         | (bs, xs) <- Base16.decode $ C8.pack str
         , CryptoPassed pub <- publicKey bs -- B.length bs == 32
-                    = [ (NodeId pub, drop 64 str) ]
+                    = [ (key2id pub, drop 64 str) ]
         | otherwise = []
 
 instance Show NodeId where
-    show (NodeId bs) = C8.unpack $ Base16.encode $ BA.convert bs
+    show nid = C8.unpack $ Base16.encode $ BA.convert $ id2key nid
 
 instance S.Serialize NodeId where
-    get = NodeId <$> getPublicKey
-    put (NodeId bs) = putPublicKey bs
-
-instance Bits NodeId where -- TODO
+    get = key2id <$> getPublicKey
+    put nid = putPublicKey $ id2key nid
 
 instance Hashable NodeId where
-    hashWithSalt salt (NodeId key) = hashWithSalt salt (BA.convert key :: ByteString)
-
-instance FiniteBits NodeId where
-    finiteBitSize _ = 256
+    hashWithSalt salt (NodeId key) = salt `xor` fromIntegral (byteSwap64 $ head key)
+
+-- instance FiniteBits NodeId where finiteBitSize _ = 256
+
+testNodeIdBit :: NodeId -> Word -> Bool
+testNodeIdBit (NodeId ws) i
+  | fromIntegral i < 256 -- 256 bits
+  , (q, r) <- quotRem (fromIntegral i) 64
+  = testBit (ws !! q) (63 - r)
+  |     otherwise      = False
+
+xorNodeId :: NodeId -> NodeId -> NodeId
+xorNodeId (NodeId xs) (NodeId ys) = NodeId $ zipWith xor xs ys
+
+sampleNodeId :: Applicative m => (Int -> m ByteString) -> NodeId -> (Int,Word8,Word8) -> m NodeId
+sampleNodeId gen (NodeId self) (q,m,b)
+    | q <= 0    = bs2id <$> gen 32
+    | q >= 32   = pure (NodeId self)
+    | let (qw,r) = (q+7) `divMod` 8 -- How many Word64 to prepend?
+          bw = shiftL (fromIntegral b) (8*(7-r))
+          mw = bw - 1 :: Word64
+          (hd, t0 : _) = splitAt (qw-1) self
+          h = xor bw (complement mw .&. t0)
+    = flip fmap (gen $ 8 * (4 - (qw-1)) ) $ \bs ->
+        let (w:ws) = unpackPublicKey bs
+        in NodeId $ hd ++ (h .|. (w .&. mw)) : ws
 
 data NodeInfo = NodeInfo
   { nodeId   :: NodeId
@@ -156,7 +205,7 @@ instance FromJSON NodeInfo where
               <|> maybe empty (return . IPv4) (ip4str >>= readMaybe)
         let (bs,_) = Base16.decode (C8.pack nidstr)
         guard (B.length bs == 32)
-        return $ NodeInfo (NodeId $ throwCryptoError . publicKey $ bs) ip (fromIntegral (portnum :: Word16))
+        return $ NodeInfo (bs2id bs) ip (fromIntegral (portnum :: Word16))
 
 getIP :: Word8 -> S.Get IP
 getIP 0x02 = IPv4 <$> S.get
@@ -199,7 +248,7 @@ instance Read NodeInfo where
                       RP.char '@' RP.+++ RP.satisfy isSpace
                       addrstr <- parseAddr
                       nid <- case Base16.decode $ C8.pack hexhash of
-                              (bs,_) | B.length bs==32 -> return (NodeId $ throwCryptoError . publicKey $ bs)
+                              (bs,_) | B.length bs==32 -> return (bs2id bs)
                               _                        -> fail "Bad node id."
                       return (nid,addrstr)
     (nid,addrstr) <- ( nodeidAt RP.+++ ( (zeroID,) <$> parseAddr) )