2 files changed, 72 insertions, 14 deletions

diff --git a/Kademlia.hs b/Kademlia.hs
index fdc1fdd2..4920467a 100644
--- a/Kademlia.hs
+++ b/Kademlia.hs
@@ -270,10 +270,11 @@ refreshBucket :: forall nid ni addr.
                  Search nid addr ni ni -> TVar (BucketList ni) -> nid -> Int -> IO ()
 refreshBucket sch var nid n = do
     tbl <- atomically (readTVar var)
-    sample <- if n+1 == R.defaultBucketCount -- Is this the last bucket?
-                then return nid              -- Yes? Search our own id.
-                else genBucketSample nid     -- No? Generate a random id.
-                        (bucketRange n (n + 1 < bktCount tbl))
+    let count = bktCount tbl
+    sample <- if n+1 >= count            -- Is this the last bucket?
+                then return nid          -- Yes? Search our own id.
+                else genBucketSample nid -- No? Generate a random id.
+                        (bucketRange n (n + 1 < count))
     resultCounter <- atomically $ newTVar Set.empty
     let fullcount = R.defaultBucketSize
     let checkBucketFull :: ni -> STM Bool
@@ -315,7 +316,7 @@ bootstrap ::
 bootstrap sch var ping ns ns0 = do
     -- First, ping the given nodes so that they are added to
     -- our routing table.
-    withTaskGroup 20 $ \g -> do
+    withTaskGroup 50 $ \g -> do
         got_response <- or <$> mapConcurrently g ping ns
         -- We resort to the hardcoded fallback nodes only when we got no
         -- responses.  This is to lesson the burden on well-known boostrap
@@ -323,13 +324,20 @@ bootstrap sch var ping ns ns0 = do
         when (not got_response) $ do
             _ <- mapConcurrently g ping ns0
             return ()
+
+    mvar <- newMVar ()
     -- Now run searches until all the buckets are full.  On a small network,
     -- this may never quit.
     --
     -- TODO: For small networks, we should give up on filling a nearby bucket
     -- at some point and move on to one farther away.
     fix $ \again -> do
+        takeMVar mvar
         tbl <- atomically $ readTVar var
+        fork $ do
+            -- Force a delay in case the search returns too quickly
+            threadDelay (60 * 1000000)
+            putMVar mvar ()
         let shp    = reverse $ zip (R.shape tbl) [0 .. ]
             unfull = filter ( (< R.defaultBucketSize) . fst ) shp
         case dropWhile ((> R.defaultBucketCount - 1) . snd) unfull of
diff --git a/Mainline.hs b/Mainline.hs
index 911f0dbc..b24f847d 100644
--- a/Mainline.hs
+++ b/Mainline.hs
@@ -73,6 +73,9 @@ import Control.Concurrent.Lifted
 import GHC.Conc                  (labelThread)
 #endif
 import Control.Exception (SomeException(..),handle)
+import Data.Aeson (FromJSON,ToJSON,(.=))
+import qualified Data.Aeson as JSON
+import Text.Read
 
 newtype NodeId = NodeId ByteString
  deriving (Eq,Ord,ByteArrayAccess, BEncode, Bits, Hashable)
@@ -97,6 +100,36 @@ data NodeInfo = NodeInfo
   }
  deriving (Eq,Ord)
 
+instance ToJSON NodeInfo where
+    toJSON (NodeInfo nid (IPv4 ip) port)
+        = JSON.object [ "node-id" .= show nid
+                      , "ipv4" .= show ip
+                      , "port" .= (fromIntegral port :: Int)
+                      ]
+    toJSON (NodeInfo nid (IPv6 ip6) port)
+        | Just ip <- un4map ip6
+          = JSON.object [ "node-id" .= show nid
+                        , "ipv4" .= show ip
+                        , "port" .= (fromIntegral port :: Int)
+                        ]
+        | otherwise
+          = JSON.object [ "node-id" .= show nid
+                        , "ipv6" .= show ip6
+                        , "port" .= (fromIntegral port :: Int)
+                        ]
+instance FromJSON NodeInfo where
+    parseJSON (JSON.Object v) = do
+        nidstr <- v JSON..: "node-id"
+        ip6str <- v JSON..:? "ipv6"
+        ip4str <- v JSON..:? "ipv4"
+        portnum <- v JSON..: "port"
+        ip <- maybe empty (return . IPv6) (ip6str >>= readMaybe)
+              <|> maybe empty (return . IPv4) (ip4str >>= readMaybe)
+        let (bs,_) = Base16.decode (Char8.pack nidstr)
+        guard (B.length bs == 20)
+        return $ NodeInfo (NodeId bs) ip (fromIntegral (portnum :: Word16))
+
+
 -- The Hashable instance depends only on the IP address and port number.  It is
 -- used to compute the announce token.
 instance Hashable NodeInfo where
@@ -145,11 +178,14 @@ putNodeInfo6 (NodeInfo (NodeId nid) (IPv6 ip) port)
 putNodeInfo6 _ = return ()
 
 
--- TODO: We should use a SocketAddrInet6 address for a dual-stack listen
--- socket.  Therefore, the behavior of this method should depend on the bind
--- address for outbound packets.
+-- | TODO: This should depend on the bind address to support IPv4-only. For
+-- now, in order to support dual-stack listen, we're going to assume IPv6 is
+-- wanted and map IPv4 addresses accordingly.
 nodeAddr :: NodeInfo -> SockAddr
-nodeAddr (NodeInfo _ ip port) = setPort port $ toSockAddr ip
+nodeAddr (NodeInfo _ ip port) =
+    case ip of
+        IPv4 ip4 -> setPort port $ toSockAddr (ipv4ToIPv6 ip4)
+        IPv6 ip6 -> setPort port $ toSockAddr ip6
 
 nodeInfo :: NodeId -> SockAddr -> Either String NodeInfo
 nodeInfo nid saddr
@@ -503,6 +539,16 @@ newClient addr = do
 
 -- | Modifies a purely random 'NodeId' to one that is related to a given
 -- routable address in accordance with BEP 42.
+--
+-- Test vectors from the spec:
+--
+-- IP           rand  example node ID
+-- ============ ===== ==========================================
+-- 124.31.75.21   1   5fbfbf f10c5d6a4ec8a88e4c6ab4c28b95eee4 01
+-- 21.75.31.124  86   5a3ce9 c14e7a08645677bbd1cfe7d8f956d532 56
+-- 65.23.51.170  22   a5d432 20bc8f112a3d426c84764f8c2a1150e6 16
+-- 84.124.73.14  65   1b0321 dd1bb1fe518101ceef99462b947a01ff 41
+-- 43.213.53.83  90   e56f6c bf5b7c4be0237986d5243b87aa6d5130 5a
 bep42 :: SockAddr -> NodeId -> Maybe NodeId
 bep42 addr (NodeId r)
  | Just ip <- fmap S.encode (fromSockAddr addr :: Maybe IPv4)
@@ -513,10 +559,9 @@ bep42 addr (NodeId r)
  where
     ip4mask = "\x03\x0f\x3f\xff" :: ByteString
     ip6mask = "\x01\x03\x07\x0f\x1f\x3f\x7f\xff" :: ByteString
-    nbhood_select = (B.last r :: Word8) .&. 7
-    nodeIdSize = 20
-    retr n = pure $ B.drop (nodeIdSize - n) $ S.encode r
-    crc = (`B.append` B.replicate 16 0) . S.encode . crc32c . B.pack
+    nbhood_select = B.last r .&. 7
+    retr n = pure $ B.drop (B.length r - n) r
+    crc = S.encode . crc32c . B.pack
     applyMask ip = case B.zipWith (.&.) msk ip of
                     (b:bs) -> (b .|. shiftL nbhood_select 5) : bs
                     bs     -> bs
@@ -524,6 +569,7 @@ bep42 addr (NodeId r)
                   | otherwise         =  ip6mask
 
 
+
 defaultHandler :: ByteString -> Handler
 defaultHandler meth = MethodHandler decodePayload errorPayload returnError
  where
@@ -540,7 +586,11 @@ mainlineKademlia client committee var sched
                         io2 <- touchBucket mainlineSpace (15*60) var sched tr
                         return $ do
                             io1 >> io2
-                            hPutStrLn stderr ("Buckets: "++show tr)
+                            hPutStrLn stderr $ unwords
+                                [ "Buckets: "
+                                , show (transitionedTo tr)
+                                , show (transitioningNode tr)
+                                ]
                    }