Check uri scheme before making RPC

2 files changed, 19 insertions, 1 deletions

diff --git a/src/Network/BitTorrent/Tracker/RPC/UDP.hs b/src/Network/BitTorrent/Tracker/RPC/UDP.hs
index 160397dd..0cb8e95e 100644
--- a/src/Network/BitTorrent/Tracker/RPC/UDP.hs
+++ b/src/Network/BitTorrent/Tracker/RPC/UDP.hs
@@ -212,6 +212,9 @@ data RpcException
     -- | Unable to lookup hostname;
   | HostLookupFailed
 
+    -- | Expecting 'udp:', but some other scheme provided.
+  | UnrecognizedScheme String
+
     -- | Tracker exists but not responding for specific number of seconds.
   | TimeoutExpired Int
 
@@ -250,7 +253,9 @@ resolveURI _       = throwIO HostUnknown
 
 -- TODO caching?
 getTrackerAddr :: Manager -> URI -> IO SockAddr
-getTrackerAddr _ = resolveURI
+getTrackerAddr _ uri
+  | uriScheme uri == "udp:" = resolveURI uri
+  |        otherwise        = throwIO (UnrecognizedScheme (uriScheme uri))
 
 {-----------------------------------------------------------------------
   Tokens
diff --git a/tests/Network/BitTorrent/Tracker/RPC/UDPSpec.hs b/tests/Network/BitTorrent/Tracker/RPC/UDPSpec.hs
index 598b2866..8a1ffc01 100644
--- a/tests/Network/BitTorrent/Tracker/RPC/UDPSpec.hs
+++ b/tests/Network/BitTorrent/Tracker/RPC/UDPSpec.hs
@@ -97,6 +97,19 @@ spec = parallel $ do
       scrape mgr (trackerURI badTracker) [def] `shouldThrow` isSomeException
 
   describe "RPC" $ do
+    describe "announce" $ do
+      it "must fail on bad scheme" $ do
+        withManager rpcOpts $ \ mgr -> do
+          q <- arbitrarySample
+          announce mgr "magnet://a.com" q
+            `shouldThrow` (== UnrecognizedScheme "magnet:")
+
+    describe "scrape" $ do
+      it "must fail on bad scheme" $ do
+        withManager rpcOpts $ \ mgr -> do
+          scrape mgr "magnet://a.com" []
+            `shouldThrow` (== UnrecognizedScheme "magnet:")
+
     forM_ (L.filter isUdpTracker trackers) $ \ TrackerEntry {..} ->
       context trackerName $ do