From 577161e540d4e13d0163ca5541853b1b672fe941 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= Date: Mon, 30 Oct 2023 19:49:45 +0100 Subject: [PATCH] ws protocol: simplify code for AnnounceAction and ScrapeAction --- crates/ws/src/workers/socket.rs | 4 +- crates/ws/src/workers/swarm.rs | 8 +- crates/ws_load_test/src/utils.rs | 4 +- .../bench_deserialize_announce_request.rs | 2 +- crates/ws_protocol/src/common.rs | 100 +++--------------- crates/ws_protocol/src/lib.rs | 20 ++-- 6 files changed, 35 insertions(+), 103 deletions(-) diff --git a/crates/ws/src/workers/socket.rs b/crates/ws/src/workers/socket.rs index 419800d..a53e9f0 100644 --- a/crates/ws/src/workers/socket.rs +++ b/crates/ws/src/workers/socket.rs @@ -775,7 +775,7 @@ impl ConnectionReader { for (consumer_index, info_hashes) in info_hashes_by_worker { let in_message = InMessage::ScrapeRequest(ScrapeRequest { - action: ScrapeAction, + action: ScrapeAction::Scrape, info_hashes: Some(ScrapeRequestInfoHashes::Multiple(info_hashes)), }); @@ -877,7 +877,7 @@ impl ConnectionWriter { slab.shrink_to_fit(); OutMessage::ScrapeResponse(ScrapeResponse { - action: ScrapeAction, + action: ScrapeAction::Scrape, files: pending.stats, }) }; diff --git a/crates/ws/src/workers/swarm.rs b/crates/ws/src/workers/swarm.rs index b973f87..6cbf8dd 100644 --- a/crates/ws/src/workers/swarm.rs +++ b/crates/ws/src/workers/swarm.rs @@ -444,7 +444,7 @@ fn handle_announce_request( for (offer, offer_receiver) in offers.into_iter().zip(offer_receivers) { let offer_out_message = OfferOutMessage { - action: AnnounceAction, + action: AnnounceAction::Announce, info_hash: request.info_hash, peer_id: request.peer_id, offer: offer.offer, @@ -470,7 +470,7 @@ fn handle_announce_request( ) { if let Some(answer_receiver) = torrent_data.peers.get(&answer_receiver_id) { let answer_out_message = AnswerOutMessage { - action: AnnounceAction, + action: AnnounceAction::Announce, peer_id: request.peer_id, info_hash: request.info_hash, answer, @@ -489,7 +489,7 @@ fn handle_announce_request( } let out_message = OutMessage::AnnounceResponse(AnnounceResponse { - action: AnnounceAction, + action: AnnounceAction::Announce, info_hash: request.info_hash, complete: torrent_data.num_seeders, incomplete: torrent_data.num_leechers(), @@ -515,7 +515,7 @@ fn handle_scrape_request( let num_to_take = info_hashes.len().min(config.protocol.max_scrape_torrents); let mut out_message = ScrapeResponse { - action: ScrapeAction, + action: ScrapeAction::Scrape, files: HashMap::with_capacity(num_to_take), }; diff --git a/crates/ws_load_test/src/utils.rs b/crates/ws_load_test/src/utils.rs index 90850b5..354e180 100644 --- a/crates/ws_load_test/src/utils.rs +++ b/crates/ws_load_test/src/utils.rs @@ -58,7 +58,7 @@ fn create_announce_request( } InMessage::AnnounceRequest(AnnounceRequest { - action: AnnounceAction, + action: AnnounceAction::Announce, info_hash: state.info_hashes[info_hash_index], peer_id, bytes_left: Some(bytes_left), @@ -82,7 +82,7 @@ fn create_scrape_request(config: &Config, state: &LoadTestState, rng: &mut impl } InMessage::ScrapeRequest(ScrapeRequest { - action: ScrapeAction, + action: ScrapeAction::Scrape, info_hashes: Some(ScrapeRequestInfoHashes::Multiple(scrape_hashes)), }) } diff --git a/crates/ws_protocol/benches/bench_deserialize_announce_request.rs b/crates/ws_protocol/benches/bench_deserialize_announce_request.rs index 429c125..7b961e7 100644 --- a/crates/ws_protocol/benches/bench_deserialize_announce_request.rs +++ b/crates/ws_protocol/benches/bench_deserialize_announce_request.rs @@ -26,7 +26,7 @@ pub fn bench(c: &mut Criterion) { let offers_len = offers.len(); let request = InMessage::AnnounceRequest(AnnounceRequest { - action: AnnounceAction, + action: AnnounceAction::Announce, info_hash, peer_id, bytes_left: Some(2), diff --git a/crates/ws_protocol/src/common.rs b/crates/ws_protocol/src/common.rs index e426f70..08e00f8 100644 --- a/crates/ws_protocol/src/common.rs +++ b/crates/ws_protocol/src/common.rs @@ -30,12 +30,28 @@ pub struct OfferId( pub [u8; 20], ); +/// Serializes to and deserializes from "announce" +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum AnnounceAction { + Announce, +} + +/// Serializes to and deserializes from "scrape" +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ScrapeAction { + Scrape, +} + +/// Serializes to and deserializes from "offer" #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum RtcOfferType { Offer, } +/// Serializes to and deserializes from "answer" #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum RtcAnswerType { @@ -64,90 +80,6 @@ pub struct RtcAnswer { pub sdp: String, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct AnnounceAction; - -impl Serialize for AnnounceAction { - fn serialize(&self, serializer: S) -> Result - where - S: Serializer, - { - serializer.serialize_str("announce") - } -} - -impl<'de> Deserialize<'de> for AnnounceAction { - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - deserializer.deserialize_str(AnnounceActionVisitor) - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ScrapeAction; - -impl Serialize for ScrapeAction { - fn serialize(&self, serializer: S) -> Result - where - S: Serializer, - { - serializer.serialize_str("scrape") - } -} - -impl<'de> Deserialize<'de> for ScrapeAction { - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - deserializer.deserialize_str(ScrapeActionVisitor) - } -} - -pub struct AnnounceActionVisitor; - -impl<'de> Visitor<'de> for AnnounceActionVisitor { - type Value = AnnounceAction; - - fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { - formatter.write_str("string with value 'announce'") - } - - fn visit_str(self, v: &str) -> Result - where - E: ::serde::de::Error, - { - if v == "announce" { - Ok(AnnounceAction) - } else { - Err(E::custom("value is not 'announce'")) - } - } -} - -pub struct ScrapeActionVisitor; - -impl<'de> Visitor<'de> for ScrapeActionVisitor { - type Value = ScrapeAction; - - fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { - formatter.write_str("string with value 'scrape'") - } - - fn visit_str(self, v: &str) -> Result - where - E: ::serde::de::Error, - { - if v == "scrape" { - Ok(ScrapeAction) - } else { - Err(E::custom("value is not 'scrape'")) - } - } -} - fn serialize_20_bytes(data: &[u8; 20], serializer: S) -> Result where S: Serializer, diff --git a/crates/ws_protocol/src/lib.rs b/crates/ws_protocol/src/lib.rs index c0b38a5..8a6bed9 100644 --- a/crates/ws_protocol/src/lib.rs +++ b/crates/ws_protocol/src/lib.rs @@ -80,7 +80,7 @@ mod tests { impl Arbitrary for OfferOutMessage { fn arbitrary(g: &mut quickcheck::Gen) -> Self { Self { - action: AnnounceAction, + action: AnnounceAction::Announce, peer_id: Arbitrary::arbitrary(g), info_hash: Arbitrary::arbitrary(g), offer_id: Arbitrary::arbitrary(g), @@ -92,7 +92,7 @@ mod tests { impl Arbitrary for AnswerOutMessage { fn arbitrary(g: &mut quickcheck::Gen) -> Self { Self { - action: AnnounceAction, + action: AnnounceAction::Announce, peer_id: Arbitrary::arbitrary(g), info_hash: Arbitrary::arbitrary(g), offer_id: Arbitrary::arbitrary(g), @@ -134,7 +134,7 @@ mod tests { let numwant = offers.as_ref().map(|offers| offers.len()); Self { - action: AnnounceAction, + action: AnnounceAction::Announce, info_hash: Arbitrary::arbitrary(g), peer_id: Arbitrary::arbitrary(g), bytes_left: Arbitrary::arbitrary(g), @@ -151,7 +151,7 @@ mod tests { impl Arbitrary for AnnounceResponse { fn arbitrary(g: &mut quickcheck::Gen) -> Self { Self { - action: AnnounceAction, + action: AnnounceAction::Announce, info_hash: Arbitrary::arbitrary(g), complete: Arbitrary::arbitrary(g), incomplete: Arbitrary::arbitrary(g), @@ -163,7 +163,7 @@ mod tests { impl Arbitrary for ScrapeRequest { fn arbitrary(g: &mut quickcheck::Gen) -> Self { Self { - action: ScrapeAction, + action: ScrapeAction::Scrape, info_hashes: Arbitrary::arbitrary(g), } } @@ -194,7 +194,7 @@ mod tests { let files: Vec<(InfoHash, ScrapeStatistics)> = Arbitrary::arbitrary(g); Self { - action: ScrapeAction, + action: ScrapeAction::Scrape, files: files.into_iter().collect(), } } @@ -277,7 +277,7 @@ mod tests { ]); let expected = ScrapeRequest { - action: ScrapeAction, + action: ScrapeAction::Scrape, info_hashes: Some(info_hashes), }; @@ -300,7 +300,7 @@ mod tests { ScrapeRequestInfoHashes::Single(info_hash_from_bytes(b"aaaabbbbccccddddeeee")); let expected = ScrapeRequest { - action: ScrapeAction, + action: ScrapeAction::Scrape, info_hashes: Some(info_hashes), }; @@ -330,7 +330,7 @@ mod tests { }; let expected = ScrapeRequest { - action: ScrapeAction, + action: ScrapeAction::Scrape, info_hashes: None, }; @@ -349,7 +349,7 @@ mod tests { }; let expected = ScrapeRequest { - action: ScrapeAction, + action: ScrapeAction::Scrape, info_hashes: None, };