use crate::bittorrent_dht::bittorrent_dht::BitTorrentDHTTransaction;
use bendy::decoding::{Decoder, Error, FromBencode, Object, ResultExt};
use nom7::bytes::complete::take;
use nom7::number::complete::be_u16;
use nom7::IResult;
#[derive(Debug, Eq, PartialEq)]
pub struct BitTorrentDHTRequest {
pub id: Vec<u8>,
pub target: Option<Vec<u8>>,
pub info_hash: Option<Vec<u8>>,
pub token: Option<Vec<u8>>,
pub implied_port: Option<u8>,
pub port: Option<u16>,
}
#[derive(Debug, Eq, PartialEq)]
pub struct BitTorrentDHTResponse {
pub id: Vec<u8>,
pub nodes: Option<Vec<Node>>,
pub nodes6: Option<Vec<Node>>,
pub values: Option<Vec<Peer>>,
pub token: Option<Vec<u8>>,
}
#[derive(Debug, Eq, PartialEq)]
pub struct BitTorrentDHTError {
pub num: u16,
pub msg: String,
}
#[derive(Debug, Eq, PartialEq)]
pub struct Node {
pub id: Vec<u8>,
pub ip: Vec<u8>,
pub port: u16,
}
#[derive(Debug, Eq, PartialEq)]
pub struct Peer {
pub ip: Vec<u8>,
pub port: u16,
}
pub fn parse_node(i: &[u8]) -> IResult<&[u8], Node> {
let (i, id) = take(20usize)(i)?;
let (i, ip) = take(4usize)(i)?;
let (i, port) = be_u16(i)?;
Ok((
i,
Node {
id: id.to_vec(),
ip: ip.to_vec(),
port,
},
))
}
pub fn parse_node6(i: &[u8]) -> IResult<&[u8], Node> {
let (i, id) = take(20usize)(i)?;
let (i, ip) = take(16usize)(i)?;
let (i, port) = be_u16(i)?;
Ok((
i,
Node {
id: id.to_vec(),
ip: ip.to_vec(),
port,
},
))
}
fn parse_peer(i: &[u8]) -> IResult<&[u8], Peer> {
let (i, ip) = if i.len() < 18 {
take(4usize)(i)
} else {
take(16usize)(i)
}?;
let (i, port) = be_u16(i)?;
Ok((
i,
Peer {
ip: ip.to_vec(),
port,
},
))
}
impl FromBencode for BitTorrentDHTRequest {
const EXPECTED_RECURSION_DEPTH: usize = 1;
fn decode_bencode_object(object: Object) -> Result<Self, Error>
where
Self: Sized,
{
let mut id = None;
let mut target = None;
let mut info_hash = None;
let mut token = None;
let mut implied_port = None;
let mut port = None;
let mut dict_dec = object.try_into_dictionary()?;
while let Some(pair) = dict_dec.next_pair()? {
match pair {
(b"id", value) => {
id = value.try_into_bytes().context("id").map(Some)?;
}
(b"target", value) => {
target = value
.try_into_bytes()
.context("target")
.map(|v| Some(v.to_vec()))?;
}
(b"info_hash", value) => {
info_hash = value
.try_into_bytes()
.context("info_hash")
.map(|v| Some(v.to_vec()))?;
}
(b"token", value) => {
token = value
.try_into_bytes()
.context("token")
.map(|v| Some(v.to_vec()))?;
}
(b"implied_port", value) => {
implied_port = u8::decode_bencode_object(value)
.context("implied_port")
.map(Some)?
}
(b"port", value) => {
port = u16::decode_bencode_object(value)
.context("port")
.map(Some)?
}
(_unknown_field, _) => {}
}
}
let id = id.ok_or_else(|| Error::missing_field("id"))?;
Ok(BitTorrentDHTRequest {
id: id.to_vec(),
target,
info_hash,
token,
implied_port,
port,
})
}
}
impl FromBencode for BitTorrentDHTResponse {
const EXPECTED_RECURSION_DEPTH: usize = 2;
fn decode_bencode_object(object: Object) -> Result<Self, Error>
where
Self: Sized,
{
let mut id = None;
let mut nodes = None;
let mut nodes6 = None;
let mut values = vec![];
let mut token = None;
let mut dict_dec = object.try_into_dictionary()?;
while let Some(pair) = dict_dec.next_pair()? {
match pair {
(b"id", value) => {
id = value.try_into_bytes().context("id").map(Some)?;
}
(b"nodes", value) => {
let (_, decoded_nodes) =
nom7::multi::many0(parse_node)(value.try_into_bytes().context("nodes")?)
.map_err(|_| Error::malformed_content("nodes.node"))?;
if !decoded_nodes.is_empty() {
nodes = Some(decoded_nodes);
}
}
(b"nodes6", value) => {
let (_, decoded_nodes) =
nom7::multi::many0(parse_node6)(value.try_into_bytes().context("nodes6")?)
.map_err(|_| Error::malformed_content("nodes6.nodes6"))?;
if !decoded_nodes.is_empty() {
nodes6 = Some(decoded_nodes);
}
}
(b"values", value) => {
if let Object::List(mut list) = value {
while let Some(entry) = list.next_object()? {
let (_, peer) =
parse_peer(entry.try_into_bytes().context("values.entry")?)
.map_err(|_| Error::malformed_content("values.entry.peer"))?;
values.push(peer);
}
}
}
(b"token", value) => {
token = value
.try_into_bytes()
.context("token")
.map(|v| Some(v.to_vec()))?;
}
(_unknown_field, _) => {}
}
}
let id = id.ok_or_else(|| Error::missing_field("id"))?;
Ok(BitTorrentDHTResponse {
id: id.to_vec(),
nodes,
nodes6,
values: if values.is_empty() {
None
} else {
Some(values)
},
token,
})
}
}
impl FromBencode for BitTorrentDHTError {
const EXPECTED_RECURSION_DEPTH: usize = 1;
fn decode_bencode_object(object: Object) -> Result<Self, Error>
where
Self: Sized,
{
let mut num = None;
let mut msg = None;
let mut list_dec = object.try_into_list()?;
while let Some(object) = list_dec.next_object()? {
match object {
Object::Integer(_) => {
num = u16::decode_bencode_object(object)
.context("num")
.map(Some)?;
}
Object::Bytes(_) => {
msg = String::decode_bencode_object(object)
.context("msg")
.map(Some)?;
}
_ => {}
}
}
let num = num.ok_or_else(|| Error::missing_field("num"))?;
let msg = msg.ok_or_else(|| Error::missing_field("msg"))?;
Ok(BitTorrentDHTError { num, msg })
}
}
pub fn parse_bittorrent_dht_packet(
bytes: &[u8], tx: &mut BitTorrentDHTTransaction,
) -> Result<(), Error> {
let mut decoder = Decoder::new(bytes).with_max_depth(3);
let object = decoder.next_object()?;
let mut packet_type = None;
let mut query_type = None;
let mut query_arguments = None;
let mut response = None;
let mut error = None;
let mut transaction_id = None;
let mut client_version = None;
let mut dict_dec = object
.ok_or_else(|| Error::unexpected_token("Dict", "EOF"))?
.try_into_dictionary()?;
while let Some(pair) = dict_dec.next_pair()? {
match pair {
(b"y", value) => {
packet_type = String::decode_bencode_object(value)
.context("packet_type")
.map(Some)?;
}
(b"q", value) => {
query_type = String::decode_bencode_object(value)
.context("query_type")
.map(Some)?;
}
(b"a", value) => {
query_arguments = BitTorrentDHTRequest::decode_bencode_object(value)
.context("query_arguments")
.map(Some)?;
}
(b"r", value) => {
response = BitTorrentDHTResponse::decode_bencode_object(value)
.context("response")
.map(Some)?;
}
(b"e", value) => {
error = BitTorrentDHTError::decode_bencode_object(value)
.context("error")
.map(Some)?;
}
(b"t", value) => {
transaction_id = value.try_into_bytes().context("transaction_id").map(Some)?;
}
(b"v", value) => {
client_version = value
.try_into_bytes()
.context("client_version")
.map(|v| Some(v.to_vec()))?;
}
(_unknown_field, _) => {}
}
}
if let Some(t) = packet_type {
match t.as_str() {
"q" => {
tx.request_type =
Some(query_type.ok_or_else(|| Error::missing_field("query_type"))?);
tx.request =
Some(query_arguments.ok_or_else(|| Error::missing_field("query_arguments"))?);
}
"r" => {
tx.response = Some(response.ok_or_else(|| Error::missing_field("response"))?);
}
"e" => {
tx.error = Some(error.ok_or_else(|| Error::missing_field("error"))?);
}
v => {
return Err(Error::unexpected_token("packet_type q, r, or e", v));
}
}
} else {
return Err(Error::missing_field("packet_type"));
}
tx.transaction_id = transaction_id
.ok_or_else(|| Error::missing_field("transaction_id"))?
.to_vec();
tx.client_version = client_version;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::direction::Direction;
use test_case::test_case;
#[test_case(
b"d2:id20:abcdefghij0123456789e",
BitTorrentDHTRequest { id: b"abcdefghij0123456789".to_vec(), implied_port: None, info_hash: None, port: None, token: None, target: None } ;
"test request from bencode 2")]
#[test_case(
b"d2:id20:abcdefghij01234567899:info_hash20:mnopqrstuvwxyz123456e",
BitTorrentDHTRequest { id: b"abcdefghij0123456789".to_vec(), implied_port: None, info_hash: Some(b"mnopqrstuvwxyz123456".to_vec()), port: None, token: None, target: None } ;
"test request from bencode 4")]
fn test_request_from_bencode(encoded: &[u8], expected: BitTorrentDHTRequest) {
let decoded = BitTorrentDHTRequest::from_bencode(encoded).unwrap();
assert_eq!(expected, decoded);
}
#[test_case(
b"d12:implied_porti1e9:info_hash20:mnopqrstuvwxyz1234564:porti6881e5:token8:aoeusnthe",
"Error: missing field: id" ;
"test request from bencode err 1")]
#[test_case(
b"d2:id20:abcdefghij012345678912:implied_porti9999e9:info_hash20:mnopqrstuvwxyz1234564:porti6881e5:token8:aoeusnthe",
"Error: malformed content discovered in implied_port" ;
"test request from bencode err 2")]
#[test_case(
b"d2:id20:abcdefghij012345678912:implied_porti-1e9:info_hash20:mnopqrstuvwxyz1234564:porti6881e5:token8:aoeusnthe",
"Error: malformed content discovered in implied_port" ;
"test request from bencode err 3")]
#[test_case(
b"d2:id20:abcdefghij012345678912:implied_porti1e9:info_hash20:mnopqrstuvwxyz1234564:porti9999999e5:token8:aoeusnthe",
"Error: malformed content discovered in port" ;
"test request from bencode err 4")]
#[test_case(
b"d2:id20:abcdefghij012345678912:implied_porti1e9:info_hash20:mnopqrstuvwxyz1234564:porti-1e5:token8:aoeusnthe",
"Error: malformed content discovered in port" ;
"test request from bencode err 5")]
#[test_case(
b"i123e",
"Error: discovered Dict but expected Num" ;
"test request from bencode err 6")]
fn test_request_from_bencode_err(encoded: &[u8], expected_error: &str) {
let err = BitTorrentDHTRequest::from_bencode(encoded).unwrap_err();
assert_eq!(expected_error, err.to_string());
}
#[test_case(
b"d5:token8:aoeusnth6:valuesl6:axje.u6:idhtnmee",
"Error: missing field: id" ;
"test response from bencode err 1")]
#[test_case(
b"i123e",
"Error: discovered Dict but expected Num" ;
"test response from bencode err 2")]
fn test_response_from_bencode_err(encoded: &[u8], expected_error: &str) {
let err = BitTorrentDHTResponse::from_bencode(encoded).unwrap_err();
assert_eq!(expected_error, err.to_string());
}
#[test_case(
b"li201e23:A Generic Error Ocurrede",
BitTorrentDHTError { num: 201u16, msg: "A Generic Error Ocurred".to_string() } ;
"test error from bencode 1")]
#[test_case(
b"li202e12:Server Errore",
BitTorrentDHTError { num: 202u16, msg: "Server Error".to_string() } ;
"test error from bencode 2")]
#[test_case(
b"li203e14:Protocol Errore",
BitTorrentDHTError { num: 203u16, msg: "Protocol Error".to_string() } ;
"test error from bencode 3")]
#[test_case(
b"li204e14:Method Unknowne",
BitTorrentDHTError { num: 204u16, msg: "Method Unknown".to_string() } ;
"test error from bencode 4")]
fn test_error_from_bencode(encoded: &[u8], expected: BitTorrentDHTError) {
let decoded = BitTorrentDHTError::from_bencode(encoded).unwrap();
assert_eq!(expected, decoded);
}
#[test_case(
b"l23:A Generic Error Ocurrede",
"Error: missing field: num" ;
"test error from bencode err 1")]
#[test_case(
b"li201ee",
"Error: missing field: msg" ;
"test error from bencode err 2")]
#[test_case(
b"li999999ee",
"Error: malformed content discovered in num" ;
"test error from bencode err 3")]
#[test_case(
b"li-1ee",
"Error: malformed content discovered in num" ;
"test error from bencode err 4")]
#[test_case(
b"i123e",
"Error: discovered List but expected Num" ;
"test error from bencode err 5")]
fn test_error_from_bencode_err(encoded: &[u8], expected_error: &str) {
let err = BitTorrentDHTError::from_bencode(encoded).unwrap_err();
assert_eq!(expected_error, err.to_string());
}
#[test_case(
b"d1:ad2:id20:abcdefghij0123456789e1:q4:ping1:t2:aa1:v4:UT011:y1:qe",
Some("ping".to_string()),
Some(BitTorrentDHTRequest { id: b"abcdefghij0123456789".to_vec(), implied_port: None, info_hash: None, port: None, token: None, target: None }),
None,
None,
b"aa".to_vec(),
Some(b"UT01".to_vec()) ;
"test parse bittorrent dht packet 1"
)]
#[test_case(
b"d1:eli201e23:A Generic Error Ocurrede1:t2:aa1:v4:UT011:y1:ee",
None,
None,
None,
Some(BitTorrentDHTError { num: 201u16, msg: "A Generic Error Ocurred".to_string() }),
b"aa".to_vec(),
Some(b"UT01".to_vec()) ;
"test parse bittorrent dht packet 3"
)]
fn test_parse_bittorrent_dht_packet(
encoded: &[u8], request_type: Option<String>,
expected_request: Option<BitTorrentDHTRequest>,
expected_response: Option<BitTorrentDHTResponse>,
expected_error: Option<BitTorrentDHTError>, expected_transaction_id: Vec<u8>,
expected_client_version: Option<Vec<u8>>,
) {
let mut tx = BitTorrentDHTTransaction::new(Direction::ToServer);
parse_bittorrent_dht_packet(encoded, &mut tx).unwrap();
assert_eq!(request_type, tx.request_type);
assert_eq!(expected_request, tx.request);
assert_eq!(expected_response, tx.response);
assert_eq!(expected_error, tx.error);
assert_eq!(expected_transaction_id, tx.transaction_id);
assert_eq!(expected_client_version, tx.client_version);
}
#[test_case(
b"",
"Error: discovered Dict but expected EOF" ;
"test parse bittorrent dht packet err 1"
)]
#[test_case(
b"li2123ei321ee",
"Error: discovered Dict but expected List" ;
"test parse bittorrent dht packet err 2"
)]
#[test_case(
b"d1:ad2:id20:abcdefghij0123456789e1:q4:ping1:t2:aae",
"Error: missing field: packet_type" ;
"test parse bittorrent dht packet err 3"
)]
#[test_case(
b"d1:ad2:id20:abcdefghij0123456789e1:q4:ping1:t2:aa1:y1:Fe",
"Error: discovered packet_type q, r, or e but expected F" ;
"test parse bittorrent dht packet err 4"
)]
#[test_case(
b"d1:ad2:id20:abcdefghij0123456789e1:t2:aa1:y1:qe",
"Error: missing field: query_type" ;
"test parse bittorrent dht packet err 5"
)]
#[test_case(
b"d1:q4:ping1:t2:aa1:y1:qe",
"Error: missing field: query_arguments" ;
"test parse bittorrent dht packet err 6"
)]
#[test_case(
b"d1:t2:aa1:y1:re",
"Error: missing field: response" ;
"test parse bittorrent dht packet err 7"
)]
#[test_case(
b"d1:t2:aa1:y1:ee",
"Error: missing field: error" ;
"test parse bittorrent dht packet err 8"
)]
#[test_case(
b"d1:ade1:q4:ping1:t2:aa1:y1:qe",
"Error: missing field: id in query_arguments" ;
"test parse bittorrent dht packet err 9"
)]
#[test_case(
b"d1:ad2:id20:abcdefghij0123456789e1:q4:ping1:y1:qe",
"Error: missing field: transaction_id" ;
"test parse bittorrent dht packet err 10"
)]
fn test_parse_bittorrent_dht_packet_err(encoded: &[u8], expected_error: &str) {
let mut tx = BitTorrentDHTTransaction::new(Direction::ToServer);
let err = parse_bittorrent_dht_packet(encoded, &mut tx).unwrap_err();
assert_eq!(expected_error, err.to_string());
}
#[test]
fn test_parse_node() {
let bytes = b"aaaaaaaaaaaaaaaaaaaa\x00\x00\x00\x00\x00\x01";
let (_rem, node) = parse_node(bytes).unwrap();
assert_eq!(node.id, b"aaaaaaaaaaaaaaaaaaaa");
assert_eq!(node.ip, b"\x00\x00\x00\x00");
assert_eq!(node.port, 1);
let bytes = b"aaaaaaaaaaaaaaaaaaa\x00\x00\x00\x00\x00\x01";
assert!(parse_node(bytes).is_err());
let bytes = b"aaaaaaaaaaaaaaaaaaaa\x00\x00\x00\x00\x00\x01bb";
let (rem, _node) = parse_node(bytes).unwrap();
assert_eq!(rem, b"bb");
}
}