1 files changed, 70 insertions, 21 deletions

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) )