From b2a3c034b0932100b5a7d081247804988bb02655 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Sun, 3 Apr 2022 19:17:57 +0200
Subject: [PATCH] http_protocol: allow announce/scrape url suffix; use in load
 test

---
 aquatic_http_load_test/src/config.rs  |  3 +++
 aquatic_http_load_test/src/network.rs |  2 +-
 aquatic_http_protocol/src/request.rs  | 18 +++++++++++-------
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/aquatic_http_load_test/src/config.rs b/aquatic_http_load_test/src/config.rs
index 6f88e79..a581e89 100644
--- a/aquatic_http_load_test/src/config.rs
+++ b/aquatic_http_load_test/src/config.rs
@@ -18,6 +18,8 @@ pub struct Config {
     /// opened as quickly as possible, which is useful when the tracker
     /// does not keep connections alive.
     pub connection_creation_interval_ms: u64,
+    /// Announce/scrape url suffix. Use `/my_token/` to get `/announce/my_token/`
+    pub url_suffix: String,
     pub duration: usize,
     pub torrents: TorrentConfig,
     #[cfg(feature = "cpu-pinning")]
@@ -56,6 +58,7 @@ impl Default for Config {
             num_workers: 1,
             num_connections: 128,
             connection_creation_interval_ms: 10,
+            url_suffix: "".into(),
             duration: 0,
             torrents: TorrentConfig::default(),
             #[cfg(feature = "cpu-pinning")]
diff --git a/aquatic_http_load_test/src/network.rs b/aquatic_http_load_test/src/network.rs
index eb0a8b4..a8fc57a 100644
--- a/aquatic_http_load_test/src/network.rs
+++ b/aquatic_http_load_test/src/network.rs
@@ -135,7 +135,7 @@ impl Connection {
                 let request =
                     create_random_request(&self.config, &self.load_test_state, &mut self.rng);
 
-                request.write(&mut self.tls.writer())?;
+                request.write(&mut self.tls.writer(), self.config.url_suffix.as_bytes())?;
                 self.queued_responses += 1;
 
                 self.send_new_request = false;
diff --git a/aquatic_http_protocol/src/request.rs b/aquatic_http_protocol/src/request.rs
index 7f29cf7..1934575 100644
--- a/aquatic_http_protocol/src/request.rs
+++ b/aquatic_http_protocol/src/request.rs
@@ -22,8 +22,10 @@ pub struct AnnounceRequest {
 }
 
 impl AnnounceRequest {
-    fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<()> {
-        output.write_all(b"GET /announce?info_hash=")?;
+    fn write<W: Write>(&self, output: &mut W, url_suffix: &[u8]) -> ::std::io::Result<()> {
+        output.write_all(b"GET /announce")?;
+        output.write_all(url_suffix)?;
+        output.write_all(b"?info_hash=")?;
         urlencode_20_bytes(self.info_hash.0, output)?;
 
         output.write_all(b"&peer_id=")?;
@@ -167,8 +169,10 @@ pub struct ScrapeRequest {
 }
 
 impl ScrapeRequest {
-    fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<()> {
+    fn write<W: Write>(&self, output: &mut W, url_suffix: &[u8]) -> ::std::io::Result<()> {
         output.write_all(b"GET /scrape?")?;
+        output.write_all(url_suffix)?;
+        output.write_all(b"?")?;
 
         let mut first = true;
 
@@ -307,10 +311,10 @@ impl Request {
         }
     }
 
-    pub fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<()> {
+    pub fn write<W: Write>(&self, output: &mut W, url_suffix: &[u8]) -> ::std::io::Result<()> {
         match self {
-            Self::Announce(r) => r.write(output),
-            Self::Scrape(r) => r.write(output),
+            Self::Announce(r) => r.write(output, url_suffix),
+            Self::Scrape(r) => r.write(output, url_suffix),
         }
     }
 }
@@ -436,7 +440,7 @@ mod tests {
 
             let mut bytes = Vec::new();
 
-            request.write(&mut bytes).unwrap();
+            request.write(&mut bytes, &[]).unwrap();
 
             let parsed_request = Request::from_bytes(&bytes[..]).unwrap();