PEX: Announce support and message definition

This commit is contained in:
Ivan 2024-08-24 18:18:07 +02:00 committed by Igor Katson
parent c758d0022d
commit bbc951733f
No known key found for this signature in database
GPG key ID: B4EC22B66D61A3F5
7 changed files with 213 additions and 35 deletions

View file

@ -9,6 +9,7 @@ use buffers::{ByteBuf, ByteBufOwned};
use byteorder::{ByteOrder, BE};
use bytes::Bytes;
use clone_to_owned::CloneToOwned;
use extended::PeerExtendedMessageIds;
use librqbit_core::{constants::CHUNK_SIZE, hash_id::Id20, lengths::ChunkInfo};
use serde::{Deserialize, Serialize};
@ -47,6 +48,9 @@ const MSGID_EXTENDED: u8 = 20;
pub const EXTENDED_UT_METADATA_KEY: &[u8] = b"ut_metadata";
pub const MY_EXTENDED_UT_METADATA: u8 = 3;
pub const EXTENDED_UT_PEX_KEY: &[u8] = b"ut_pex";
pub const MY_EXTENDED_UT_PEX: u8 = 1;
#[derive(Debug)]
pub enum MessageDeserializeError {
NotEnoughData(usize, &'static str),
@ -276,7 +280,7 @@ where
pub fn serialize(
&self,
out: &mut Vec<u8>,
extended_handshake_ut_metadata: &dyn Fn() -> Option<u8>,
peer_extended_messages: &dyn Fn() -> PeerExtendedMessageIds,
) -> anyhow::Result<usize> {
let (lp, msg_id) = self.len_prefix_and_msg_id();
@ -326,7 +330,7 @@ where
Ok(msg_len)
}
Message::Extended(e) => {
e.serialize(out, extended_handshake_ut_metadata)?;
e.serialize(out, peer_extended_messages)?;
let msg_size = out.len();
// no fucking idea why +1, but I tweaked that for it all to match up
// with real messages.
@ -636,7 +640,7 @@ mod tests {
fn test_extended_serialize() {
let msg = Message::Extended(ExtendedMessage::Handshake(ExtendedHandshake::new()));
let mut out = Vec::new();
msg.serialize(&mut out, &|| None).unwrap();
msg.serialize(&mut out, &|| Default::default()).unwrap();
dbg!(out);
}
@ -652,7 +656,7 @@ mod tests {
let (msg, size) = MessageBorrowed::deserialize(&buf).unwrap();
assert_eq!(size, buf.len());
let mut write_buf = Vec::new();
msg.serialize(&mut write_buf, &|| None).unwrap();
msg.serialize(&mut write_buf, &|| Default::default()).unwrap();
if buf != write_buf {
{
use std::io::Write;