Untested announce-peer command for bittorrent dht.

3 files changed, 67 insertions, 0 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index d23aca78..6529e746 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -107,6 +107,11 @@ data DHTQuery nid ni = forall addr r tok.
     , qshowTok :: tok -> Maybe String
     }
 
+data DHTAnnouncable nid ni = forall dta r. Show r => DHTAnnouncable
+    { announceParseData :: String -> String -> IO (Either String dta)
+    , announceSendData :: dta -> Maybe ni -> IO (Maybe r)
+    }
+
 data DHTSearch nid ni = forall addr tok r. DHTSearch
     { searchThread     :: ThreadId
     , searchState      :: SearchState nid addr tok ni r
@@ -130,6 +135,7 @@ data DHT = forall nid ni. ( Show     ni
     { dhtBuckets       :: TVar (BucketList ni)
     , dhtPing          :: ni -> IO Bool
     , dhtQuery         :: Map.Map String (DHTQuery nid ni)
+    , dhtAnnouncables  :: Map.Map String (DHTAnnouncable nid ni)
     , dhtParseId       :: String -> Either String nid
     , dhtSearches      :: TVar (Map.Map (String,nid) (DHTSearch nid ni))
     , dhtFallbackNodes :: IO [ni]
@@ -440,6 +446,32 @@ clientSession s@Session{..} sock cnum h = do
                               goQuery
                               $ Map.lookup method dhtQuery
 
+        ("p", s) | Just DHT{..} <- Map.lookup netname dhts
+                    -> cmd0 $ do
+                        -- arguments: method
+                        --            data
+                        --            token
+                        --            (optional dest-ni)
+                        self <- atomically $ thisNode <$> readTVar dhtBuckets
+                        let (method,xs) = break isSpace $ dropWhile isSpace s
+                            (dtastr,ys) = break isSpace $ dropWhile isSpace xs
+                            (tokenstr,zs) = break isSpace $ dropWhile isSpace ys
+                            destination = dropWhile isSpace zs
+                            goTarget DHTAnnouncable{..} = do
+                                parseResult <- announceParseData dtastr tokenstr
+                                case parseResult of
+                                    Left e -> hPutClient h e
+                                    Right nid -> go nid >>= either (hPutClient h) (hPutClient h . show)
+                              where
+                                go | null destination = fmap (maybe (Left "Timeout.") Right)
+                                                        . flip announceSendData Nothing
+                                   | otherwise        = case readEither destination of
+                                        Right ni -> fmap (maybe (Left "Timeout.") Right)
+                                                    . flip announceSendData (Just ni)
+                                        Left  e  -> const $ return $ Left ("Bad destination: "++e)
+                        maybe (hPutClient h ("Unsupported method: "++method))
+                              goTarget
+                              $ Map.lookup method dhtAnnouncables
         ("s", s) | Just dht@DHT{..} <- Map.lookup netname dhts
                     -> cmd0 $ do
                         let (method,xs) = break isSpace s
@@ -583,6 +615,7 @@ main = do
             (bt,btR) <- Mainline.newClient swarms addr
             quitBt <- forkListener "bt" (clientNet bt)
             mainlineSearches <- atomically $ newTVar Map.empty
+            peerPort <- atomically $ newTVar 6881 -- BitTorrent client TCP port.
             let mainlineDHT bkts wantip = DHT
                     { dhtBuckets = bkts btR
                     , dhtPing    = Mainline.ping bt
@@ -604,6 +637,23 @@ main = do
                     , dhtParseId  = readEither :: String -> Either String Mainline.NodeId
                     , dhtSearches = mainlineSearches
                     , dhtFallbackNodes = Mainline.bootstrapNodes wantip
+                    , dhtAnnouncables = Map.fromList
+                                    [ ("peer", DHTAnnouncable { announceSendData = \dta -> \case
+                                                                    Just ni -> Mainline.announce bt dta ni
+                                                                    Nothing -> return Nothing
+                                                              , announceParseData = \str tokstr -> do
+                                                                    port <- atomically $ readTVar peerPort
+                                                                    let ih = read str
+                                                                        tok = read tokstr
+                                                                    return $ Right $ Mainline.mkAnnounce port ih tok
+                                                              })
+                                    , ("port", DHTAnnouncable { announceParseData = \portstr _ -> return $ readEither portstr
+                                                              , announceSendData = \dta -> \case
+                                                                    Nothing -> do atomically $ writeTVar peerPort (dta :: PortNumber)
+                                                                                  return $ Just dta
+                                                                    Just _ -> return Nothing
+                                                              })]
+
                     }
                 dhts = Map.fromList $
                         ("bt4", mainlineDHT Mainline.routing4 Want_IP4)
@@ -656,6 +706,7 @@ main = do
                     , dhtParseId  = readEither :: String -> Either String Tox.NodeId
                     , dhtSearches = toxSearches
                     , dhtFallbackNodes = return []
+                    , dhtAnnouncables = Map.empty -- TODO
                     }
                 dhts = Map.fromList $
                         ("tox4", toxDHT Tox.routing4)
diff --git a/src/Network/BitTorrent/DHT/Token.hs b/src/Network/BitTorrent/DHT/Token.hs
index 756b5a98..3da59c53 100644
--- a/src/Network/BitTorrent/DHT/Token.hs
+++ b/src/Network/BitTorrent/DHT/Token.hs
@@ -44,6 +44,7 @@ module Network.BitTorrent.DHT.Token
        , Network.BitTorrent.DHT.Token.update
        ) where
 
+import Control.Arrow
 import Control.Monad.State
 #ifdef VERSION_bencoding
 import Data.BEncode (BEncode)
@@ -75,6 +76,9 @@ newtype Token = Token BS.ByteString
 instance Show Token where
     show (Token bs) = B8.unpack $ Base16.encode bs
 
+instance Read Token where
+    readsPrec i s = pure $ (Token *** B8.unpack) $ Base16.decode (B8.pack s)
+
 -- | Meaningless token, for testing purposes only.
 instance Default Token where
   def = makeToken (0::Int) 0
diff --git a/src/Network/BitTorrent/MainlineDHT.hs b/src/Network/BitTorrent/MainlineDHT.hs
index f4ce4019..a7359bda 100644
--- a/src/Network/BitTorrent/MainlineDHT.hs
+++ b/src/Network/BitTorrent/MainlineDHT.hs
@@ -906,6 +906,15 @@ data Announce = Announce
 
   } deriving (Show, Eq, Typeable)
 
+mkAnnounce :: PortNumber -> InfoHash -> Token -> Announce
+mkAnnounce portnum info token = Announce
+    { topic = info
+    , port = portnum
+    , sessionToken = token
+    , announcedName = Nothing
+    , impliedPort = False
+    }
+
 peer_ip_key   = "ip"
 peer_id_key   = "peer id"
 peer_port_key = "port"
@@ -1078,3 +1087,6 @@ resolve want hostAndPort = do
         return $ addrAddress info
 
 
+announce :: MainlineClient -> Announce -> NodeInfo -> IO (Maybe Announced)
+announce client msg addr = do
+    mainlineSend (Method "announce_peer") id (\() -> msg) client () addr