SoulStream/crates/soulseek-proto/src/dist.rs

//! Distributed network messages (connection type "D"), single-byte codes.
//! We join as a child: connect to a parent, receive relayed searches, and
//! report our branch level/root.

use bytes::Bytes;

use crate::wire::{Reader, Writer};
use crate::Result;

pub mod code {
    pub const PING: u8 = 0;
    pub const SEARCH_REQUEST: u8 = 3;
    pub const BRANCH_LEVEL: u8 = 4;
    pub const BRANCH_ROOT: u8 = 5;
    pub const CHILD_DEPTH: u8 = 7;
    pub const EMBEDDED_MESSAGE: u8 = 93;
}

// ---- outgoing (to parent / children) ----

pub fn branch_level(level: u32) -> Vec<u8> {
    Writer::new().u32(level).finish_u8(code::BRANCH_LEVEL)
}

pub fn branch_root(root_username: &str) -> Vec<u8> {
    Writer::new()
        .string(root_username)
        .finish_u8(code::BRANCH_ROOT)
}

pub fn ping() -> Vec<u8> {
    Writer::new().finish_u8(code::PING)
}

/// Relay a search to our children.
pub fn search_request(username: &str, token: u32, query: &str) -> Vec<u8> {
    Writer::new()
        .u32(0)
        .string(username)
        .u32(token)
        .string(query)
        .finish_u8(code::SEARCH_REQUEST)
}

// ---- incoming ----

#[derive(Debug, Clone)]
pub enum DistributedEvent {
    /// A search to match against our shares and reply to `username`.
    SearchRequest {
        username: String,
        token: u32,
        query: String,
    },
    BranchLevel(u32),
    BranchRoot(String),
    Ping,
    Unhandled {
        code: u8,
    },
}

/// Decode a distributed frame (code byte + payload). Handles the
/// `EmbeddedMessage` (93) wrapper by unwrapping to its inner distributed
/// message.
pub fn decode(frame: Bytes) -> Result<DistributedEvent> {
    let mut r = Reader::new(frame);
    let code = r.u8()?;
    decode_inner(code, &mut r)
}

fn decode_inner(code: u8, r: &mut Reader) -> Result<DistributedEvent> {
    Ok(match code {
        code::EMBEDDED_MESSAGE => {
            // [u8 inner_code][inner payload]
            let inner = r.u8()?;
            return decode_inner(inner, r);
        }
        code::SEARCH_REQUEST => {
            let _unknown = r.u32()?;
            let username = r.string()?;
            let token = r.u32()?;
            let query = r.string()?;
            DistributedEvent::SearchRequest {
                username,
                token,
                query,
            }
        }
        code::BRANCH_LEVEL => DistributedEvent::BranchLevel(r.u32()?),
        code::BRANCH_ROOT => DistributedEvent::BranchRoot(r.string()?),
        code::PING => DistributedEvent::Ping,
        other => DistributedEvent::Unhandled { code: other },
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn decode_search_request() {
        let msg = search_request("alice", 42, "the beatles");
        let ev = decode(Bytes::copy_from_slice(&msg[4..])).unwrap();
        match ev {
            DistributedEvent::SearchRequest {
                username,
                token,
                query,
            } => {
                assert_eq!(
                    (username.as_str(), token, query.as_str()),
                    ("alice", 42, "the beatles")
                );
            }
            other => panic!("wrong: {other:?}"),
        }
    }

    #[test]
    fn decode_embedded_wrapper() {
        // Server EmbeddedMessage delivers [93][3][search payload].
        let inner = search_request("bob", 7, "query");
        // inner[4] is the code byte (3); inner[5..] is the payload.
        let mut framed = vec![code::EMBEDDED_MESSAGE];
        framed.extend_from_slice(&inner[4..]); // code 3 + payload
        let ev = decode(Bytes::from(framed)).unwrap();
        match ev {
            DistributedEvent::SearchRequest {
                username, query, ..
            } => {
                assert_eq!((username.as_str(), query.as_str()), ("bob", "query"));
            }
            other => panic!("wrong: {other:?}"),
        }
    }
}