chore: initial sanitized public snapshot

2026-07-18 16:01:12 +08:00
commit 7c6b6a3746
321 changed files with 76896 additions and 0 deletions
@@ -0,0 +1,488 @@
+use super::codec::{
+    dht_dict_get_bytes, dht_dict_get_dict, dht_dict_get_int, dht_dict_get_list, dht_encode_dict,
+    dht_parse_value,
+};
+use super::compact::{decode_compact_dht_nodes, encode_compact_dht_nodes};
+use std::collections::BTreeMap;
+
+use super::{
+    DhtAnnouncePeerQueryModel, DhtBencodeValue, DhtCompactNodeModel, DhtErrorModel,
+    DhtFindNodeQueryModel, DhtFindNodeResponseModel, DhtGetPeersQueryModel,
+    DhtGetPeersResponseModel, DhtMessageBody, DhtMessageModel, DhtPingQueryModel,
+    DhtPingResponseModel, DhtQueryModel, DhtResponseModel,
+};
+
+impl DhtMessageModel {
+    #[must_use]
+    /// Builds a DHT `ping` query.
+    pub fn ping_query(transaction_id: impl Into<Vec<u8>>, node_id: impl Into<Vec<u8>>) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Query(DhtQueryModel::Ping(DhtPingQueryModel {
+                node_id: node_id.into(),
+            })),
+        }
+    }
+
+    #[must_use]
+    /// Builds a DHT `find_node` query.
+    pub fn find_node_query(
+        transaction_id: impl Into<Vec<u8>>,
+        node_id: impl Into<Vec<u8>>,
+        target: impl Into<Vec<u8>>,
+    ) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Query(DhtQueryModel::FindNode(DhtFindNodeQueryModel {
+                node_id: node_id.into(),
+                target: target.into(),
+            })),
+        }
+    }
+
+    #[must_use]
+    /// Builds a DHT `get_peers` query.
+    pub fn get_peers_query(
+        transaction_id: impl Into<Vec<u8>>,
+        node_id: impl Into<Vec<u8>>,
+        info_hash: impl Into<Vec<u8>>,
+    ) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Query(DhtQueryModel::GetPeers(DhtGetPeersQueryModel {
+                node_id: node_id.into(),
+                info_hash: info_hash.into(),
+            })),
+        }
+    }
+
+    #[must_use]
+    /// Builds a DHT `announce_peer` query.
+    pub fn announce_peer_query(
+        transaction_id: impl Into<Vec<u8>>,
+        node_id: impl Into<Vec<u8>>,
+        info_hash: impl Into<Vec<u8>>,
+        port: u16,
+        token: impl Into<Vec<u8>>,
+        implied_port: bool,
+    ) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Query(DhtQueryModel::AnnouncePeer(DhtAnnouncePeerQueryModel {
+                node_id: node_id.into(),
+                info_hash: info_hash.into(),
+                port,
+                token: token.into(),
+                implied_port,
+            })),
+        }
+    }
+
+    #[must_use]
+    /// Builds a DHT `ping` response.
+    pub fn ping_response(transaction_id: impl Into<Vec<u8>>, node_id: impl Into<Vec<u8>>) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Response(DhtResponseModel::Ping(DhtPingResponseModel {
+                node_id: node_id.into(),
+            })),
+        }
+    }
+
+    #[must_use]
+    /// Builds a DHT `find_node` response.
+    pub fn find_node_response(
+        transaction_id: impl Into<Vec<u8>>,
+        node_id: impl Into<Vec<u8>>,
+        nodes: Vec<DhtCompactNodeModel>,
+    ) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Response(DhtResponseModel::FindNode(DhtFindNodeResponseModel {
+                node_id: node_id.into(),
+                nodes,
+            })),
+        }
+    }
+
+    #[must_use]
+    /// Builds a DHT `get_peers` response.
+    pub fn get_peers_response(
+        transaction_id: impl Into<Vec<u8>>,
+        node_id: impl Into<Vec<u8>>,
+        token: Option<Vec<u8>>,
+        nodes: Option<Vec<u8>>,
+        values: Vec<Vec<u8>>,
+    ) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Response(DhtResponseModel::GetPeers(DhtGetPeersResponseModel {
+                node_id: node_id.into(),
+                token,
+                nodes,
+                values,
+            })),
+        }
+    }
+
+    #[must_use]
+    /// Builds a DHT `announce_peer` response.
+    pub fn announce_peer_response(
+        transaction_id: impl Into<Vec<u8>>,
+        node_id: impl Into<Vec<u8>>,
+    ) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Response(DhtResponseModel::Ping(DhtPingResponseModel {
+                node_id: node_id.into(),
+            })),
+        }
+    }
+
+    #[must_use]
+    /// Builds a DHT error response.
+    pub fn error_response(
+        transaction_id: impl Into<Vec<u8>>,
+        code: i64,
+        message: impl Into<String>,
+    ) -> Self {
+        Self {
+            transaction_id: transaction_id.into(),
+            body: DhtMessageBody::Error(DhtErrorModel {
+                code,
+                message: message.into(),
+            }),
+        }
+    }
+
+    #[must_use]
+    /// Returns the raw DHT transaction id bytes.
+    pub fn transaction_id(&self) -> &[u8] {
+        &self.transaction_id
+    }
+
+    #[must_use]
+    /// Returns the DHT query method name when the message body is a query.
+    pub fn method(&self) -> Option<&'static str> {
+        match &self.body {
+            DhtMessageBody::Query(DhtQueryModel::Ping(_)) => Some("ping"),
+            DhtMessageBody::Query(DhtQueryModel::FindNode(_)) => Some("find_node"),
+            DhtMessageBody::Query(DhtQueryModel::GetPeers(_)) => Some("get_peers"),
+            DhtMessageBody::Query(DhtQueryModel::AnnouncePeer(_)) => Some("announce_peer"),
+            _ => None,
+        }
+    }
+
+    #[must_use]
+    /// Returns whether this message body is a DHT query.
+    pub fn is_query(&self) -> bool {
+        matches!(self.body, DhtMessageBody::Query(_))
+    }
+
+    #[must_use]
+    /// Serializes the DHT message as a bencoded root dictionary.
+    pub fn to_bencode_bytes(&self) -> Vec<u8> {
+        dht_encode_dict(&self.as_bencode_root())
+    }
+
+    /// Parses a DHT message from a bencoded payload.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error when the payload is not a supported DHT message dictionary.
+    pub fn from_bencode_bytes(input: &[u8]) -> Result<Self, String> {
+        let (value, next) = dht_parse_value(input, 0)?;
+        if next != input.len() {
+            return Err("trailing bytes after dht message".to_owned());
+        }
+        let DhtBencodeValue::Dict(root) = value else {
+            return Err("dht message must be a bencoded dictionary".to_owned());
+        };
+        let transaction_id = dht_dict_get_bytes(&root, b"t")
+            .ok_or_else(|| "missing dht transaction id".to_owned())?;
+        let message_type =
+            dht_dict_get_bytes(&root, b"y").ok_or_else(|| "missing dht message type".to_owned())?;
+
+        match message_type.as_slice() {
+            b"q" => parse_dht_query_message(transaction_id, &root),
+            b"r" => parse_dht_response_message(transaction_id, &root),
+            b"e" => parse_dht_error_message(transaction_id, &root),
+            _ => Err("unsupported dht message type".to_owned()),
+        }
+    }
+
+    #[expect(
+        clippy::too_many_lines,
+        reason = "torrent roundtrip test keeps the end-to-end fixture in one place for auditability"
+    )]
+    /// Rebuilds the DHT message as the bencode root dictionary used on the wire.
+    fn as_bencode_root(&self) -> BTreeMap<Vec<u8>, DhtBencodeValue> {
+        let mut root = BTreeMap::new();
+        root.insert(
+            b"t".to_vec(),
+            DhtBencodeValue::Bytes(self.transaction_id.clone()),
+        );
+
+        match &self.body {
+            DhtMessageBody::Query(query) => {
+                root.insert(b"y".to_vec(), DhtBencodeValue::Bytes(b"q".to_vec()));
+                match query {
+                    DhtQueryModel::Ping(query) => {
+                        root.insert(b"q".to_vec(), DhtBencodeValue::Bytes(b"ping".to_vec()));
+                        let mut args = BTreeMap::new();
+                        args.insert(
+                            b"id".to_vec(),
+                            DhtBencodeValue::Bytes(query.node_id.clone()),
+                        );
+                        root.insert(b"a".to_vec(), DhtBencodeValue::Dict(args));
+                    }
+                    DhtQueryModel::FindNode(query) => {
+                        root.insert(b"q".to_vec(), DhtBencodeValue::Bytes(b"find_node".to_vec()));
+                        let mut args = BTreeMap::new();
+                        args.insert(
+                            b"id".to_vec(),
+                            DhtBencodeValue::Bytes(query.node_id.clone()),
+                        );
+                        args.insert(
+                            b"target".to_vec(),
+                            DhtBencodeValue::Bytes(query.target.clone()),
+                        );
+                        root.insert(b"a".to_vec(), DhtBencodeValue::Dict(args));
+                    }
+                    DhtQueryModel::GetPeers(query) => {
+                        root.insert(b"q".to_vec(), DhtBencodeValue::Bytes(b"get_peers".to_vec()));
+                        let mut args = BTreeMap::new();
+                        args.insert(
+                            b"id".to_vec(),
+                            DhtBencodeValue::Bytes(query.node_id.clone()),
+                        );
+                        args.insert(
+                            b"info_hash".to_vec(),
+                            DhtBencodeValue::Bytes(query.info_hash.clone()),
+                        );
+                        root.insert(b"a".to_vec(), DhtBencodeValue::Dict(args));
+                    }
+                    DhtQueryModel::AnnouncePeer(query) => {
+                        root.insert(
+                            b"q".to_vec(),
+                            DhtBencodeValue::Bytes(b"announce_peer".to_vec()),
+                        );
+                        let mut args = BTreeMap::new();
+                        args.insert(
+                            b"id".to_vec(),
+                            DhtBencodeValue::Bytes(query.node_id.clone()),
+                        );
+                        args.insert(
+                            b"info_hash".to_vec(),
+                            DhtBencodeValue::Bytes(query.info_hash.clone()),
+                        );
+                        args.insert(
+                            b"port".to_vec(),
+                            DhtBencodeValue::Int(i64::from(query.port)),
+                        );
+                        args.insert(
+                            b"token".to_vec(),
+                            DhtBencodeValue::Bytes(query.token.clone()),
+                        );
+                        if query.implied_port {
+                            args.insert(b"implied_port".to_vec(), DhtBencodeValue::Int(1));
+                        }
+                        root.insert(b"a".to_vec(), DhtBencodeValue::Dict(args));
+                    }
+                }
+            }
+            DhtMessageBody::Response(response) => {
+                root.insert(b"y".to_vec(), DhtBencodeValue::Bytes(b"r".to_vec()));
+                let mut payload = BTreeMap::new();
+                match response {
+                    DhtResponseModel::Ping(response) => {
+                        payload.insert(
+                            b"id".to_vec(),
+                            DhtBencodeValue::Bytes(response.node_id.clone()),
+                        );
+                    }
+                    DhtResponseModel::FindNode(response) => {
+                        payload.insert(
+                            b"id".to_vec(),
+                            DhtBencodeValue::Bytes(response.node_id.clone()),
+                        );
+                        if !response.nodes.is_empty() {
+                            payload.insert(
+                                b"nodes".to_vec(),
+                                DhtBencodeValue::Bytes(encode_compact_dht_nodes(&response.nodes)),
+                            );
+                        }
+                    }
+                    DhtResponseModel::GetPeers(response) => {
+                        payload.insert(
+                            b"id".to_vec(),
+                            DhtBencodeValue::Bytes(response.node_id.clone()),
+                        );
+                        if let Some(token) = &response.token {
+                            payload
+                                .insert(b"token".to_vec(), DhtBencodeValue::Bytes(token.clone()));
+                        }
+                        if let Some(nodes) = &response.nodes {
+                            payload
+                                .insert(b"nodes".to_vec(), DhtBencodeValue::Bytes(nodes.clone()));
+                        }
+                        if !response.values.is_empty() {
+                            payload.insert(
+                                b"values".to_vec(),
+                                DhtBencodeValue::List(
+                                    response
+                                        .values
+                                        .iter()
+                                        .cloned()
+                                        .map(DhtBencodeValue::Bytes)
+                                        .collect(),
+                                ),
+                            );
+                        }
+                    }
+                }
+                root.insert(b"r".to_vec(), DhtBencodeValue::Dict(payload));
+            }
+            DhtMessageBody::Error(error) => {
+                root.insert(b"y".to_vec(), DhtBencodeValue::Bytes(b"e".to_vec()));
+                root.insert(
+                    b"e".to_vec(),
+                    DhtBencodeValue::List(vec![
+                        DhtBencodeValue::Int(error.code),
+                        DhtBencodeValue::Bytes(error.message.as_bytes().to_vec()),
+                    ]),
+                );
+            }
+        }
+        root
+    }
+}
+
+/// Parses a DHT query message body from the decoded root dictionary.
+fn parse_dht_query_message(
+    transaction_id: Vec<u8>,
+    root: &BTreeMap<Vec<u8>, DhtBencodeValue>,
+) -> Result<DhtMessageModel, String> {
+    let method =
+        dht_dict_get_bytes(root, b"q").ok_or_else(|| "missing dht query method".to_owned())?;
+    let arguments =
+        dht_dict_get_dict(root, b"a").ok_or_else(|| "missing dht query arguments".to_owned())?;
+    let body = match method.as_slice() {
+        b"ping" => {
+            let node_id = dht_dict_get_bytes(arguments, b"id")
+                .ok_or_else(|| "missing dht ping id".to_owned())?;
+            DhtMessageBody::Query(DhtQueryModel::Ping(DhtPingQueryModel { node_id }))
+        }
+        b"find_node" => {
+            let node_id = dht_dict_get_bytes(arguments, b"id")
+                .ok_or_else(|| "missing dht find_node id".to_owned())?;
+            let target = dht_dict_get_bytes(arguments, b"target")
+                .ok_or_else(|| "missing dht find_node target".to_owned())?;
+            DhtMessageBody::Query(DhtQueryModel::FindNode(DhtFindNodeQueryModel {
+                node_id,
+                target,
+            }))
+        }
+        b"get_peers" => {
+            let node_id = dht_dict_get_bytes(arguments, b"id")
+                .ok_or_else(|| "missing dht get_peers id".to_owned())?;
+            let info_hash = dht_dict_get_bytes(arguments, b"info_hash")
+                .ok_or_else(|| "missing dht get_peers info_hash".to_owned())?;
+            DhtMessageBody::Query(DhtQueryModel::GetPeers(DhtGetPeersQueryModel {
+                node_id,
+                info_hash,
+            }))
+        }
+        b"announce_peer" => {
+            let node_id = dht_dict_get_bytes(arguments, b"id")
+                .ok_or_else(|| "missing dht announce_peer id".to_owned())?;
+            let info_hash = dht_dict_get_bytes(arguments, b"info_hash")
+                .ok_or_else(|| "missing dht announce_peer info_hash".to_owned())?;
+            let port = dht_dict_get_int(arguments, b"port")
+                .ok_or_else(|| "missing dht announce_peer port".to_owned())?;
+            let token = dht_dict_get_bytes(arguments, b"token")
+                .ok_or_else(|| "missing dht announce_peer token".to_owned())?;
+            let implied_port =
+                dht_dict_get_int(arguments, b"implied_port").is_some_and(|value| value != 0);
+            DhtMessageBody::Query(DhtQueryModel::AnnouncePeer(DhtAnnouncePeerQueryModel {
+                node_id,
+                info_hash,
+                port: u16::try_from(port)
+                    .map_err(|_| "dht announce_peer port out of range".to_owned())?,
+                token,
+                implied_port,
+            }))
+        }
+        _ => return Err("unsupported dht query method".to_owned()),
+    };
+    Ok(DhtMessageModel {
+        transaction_id,
+        body,
+    })
+}
+
+/// Parses a DHT response message body from the decoded root dictionary.
+fn parse_dht_response_message(
+    transaction_id: Vec<u8>,
+    root: &BTreeMap<Vec<u8>, DhtBencodeValue>,
+) -> Result<DhtMessageModel, String> {
+    let payload =
+        dht_dict_get_dict(root, b"r").ok_or_else(|| "missing dht response body".to_owned())?;
+    let node_id =
+        dht_dict_get_bytes(payload, b"id").ok_or_else(|| "missing dht response id".to_owned())?;
+    let token = dht_dict_get_bytes(payload, b"token");
+    let nodes = dht_dict_get_bytes(payload, b"nodes");
+    let values = dht_dict_get_list(payload, b"values")
+        .unwrap_or_default()
+        .into_iter()
+        .map(|value| match value {
+            DhtBencodeValue::Bytes(bytes) => Ok(bytes),
+            _ => Err("dht response values entries must be byte strings".to_owned()),
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+
+    let response = if token.is_some() || !values.is_empty() {
+        DhtResponseModel::GetPeers(DhtGetPeersResponseModel {
+            node_id,
+            token,
+            nodes,
+            values,
+        })
+    } else if let Some(nodes) = nodes {
+        DhtResponseModel::FindNode(DhtFindNodeResponseModel {
+            node_id,
+            nodes: decode_compact_dht_nodes(&nodes)?,
+        })
+    } else {
+        DhtResponseModel::Ping(DhtPingResponseModel { node_id })
+    };
+
+    Ok(DhtMessageModel {
+        transaction_id,
+        body: DhtMessageBody::Response(response),
+    })
+}
+
+/// Parses a DHT error message body from the decoded root dictionary.
+fn parse_dht_error_message(
+    transaction_id: Vec<u8>,
+    root: &BTreeMap<Vec<u8>, DhtBencodeValue>,
+) -> Result<DhtMessageModel, String> {
+    let errors =
+        dht_dict_get_list(root, b"e").ok_or_else(|| "missing dht error payload".to_owned())?;
+    if errors.len() != 2 {
+        return Err("dht error payload must have [code, message]".to_owned());
+    }
+    let code = match errors.first() {
+        Some(DhtBencodeValue::Int(value)) => *value,
+        _ => return Err("dht error code must be an integer".to_owned()),
+    };
+    let message = match errors.get(1) {
+        Some(DhtBencodeValue::Bytes(bytes)) => String::from_utf8_lossy(bytes).into_owned(),
+        _ => return Err("dht error message must be bytes".to_owned()),
+    };
+    Ok(DhtMessageModel {
+        transaction_id,
+        body: DhtMessageBody::Error(DhtErrorModel { code, message }),
+    })
+}