489 lines
19 KiB
Rust
489 lines
19 KiB
Rust
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 }),
|
|
})
|
|
}
|