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

//! Peer connection messages: single-byte-code initialization handshake,
//! then u32-code peer messages.

use bytes::Bytes;

use crate::types::{Directory, File, SearchResponse};
use crate::wire::{Reader, Writer};
use crate::{ProtoError, Result};

pub mod init_code {
    pub const PIERCE_FIREWALL: u8 = 0;
    pub const PEER_INIT: u8 = 1;
}

pub mod code {
    pub const BROWSE_REQUEST: u32 = 4;
    pub const BROWSE_RESPONSE: u32 = 5;
    pub const SEARCH_RESPONSE: u32 = 9;
    pub const USER_INFO_REQUEST: u32 = 15;
    pub const USER_INFO_RESPONSE: u32 = 16;
    pub const FOLDER_CONTENTS_REQUEST: u32 = 36;
    pub const FOLDER_CONTENTS_RESPONSE: u32 = 37;
    pub const TRANSFER_REQUEST: u32 = 40;
    pub const TRANSFER_RESPONSE: u32 = 41;
    pub const QUEUE_DOWNLOAD: u32 = 43;
    pub const PLACE_IN_QUEUE_RESPONSE: u32 = 44;
    pub const UPLOAD_FAILED: u32 = 46;
    pub const UPLOAD_DENIED: u32 = 50;
    pub const PLACE_IN_QUEUE_REQUEST: u32 = 51;
}

/// Connection types carried by PeerInit / ConnectToPeer.
pub const CONN_TYPE_PEER: &str = "P";
pub const CONN_TYPE_TRANSFER: &str = "F";
pub const CONN_TYPE_DISTRIBUTED: &str = "D";

/// TransferRequest direction values.
pub const DIRECTION_DOWNLOAD: u32 = 0;
pub const DIRECTION_UPLOAD: u32 = 1;

// ---- initialization (u8 code) ----

#[derive(Debug, Clone)]
pub enum InitMessage {
    PierceFirewall {
        token: u32,
    },
    PeerInit {
        username: String,
        conn_type: String,
        token: u32,
    },
}

pub fn pierce_firewall(token: u32) -> Vec<u8> {
    Writer::new()
        .u32(token)
        .finish_u8(init_code::PIERCE_FIREWALL)
}

pub fn peer_init(username: &str, conn_type: &str, token: u32) -> Vec<u8> {
    Writer::new()
        .string(username)
        .string(conn_type)
        .u32(token)
        .finish_u8(init_code::PEER_INIT)
}

/// Decode the first frame on a fresh incoming connection.
/// `frame` is the full frame body (code byte included, length header not).
pub fn decode_init(frame: Bytes) -> Result<InitMessage> {
    let mut r = Reader::new(frame);
    match r.u8()? {
        init_code::PIERCE_FIREWALL => Ok(InitMessage::PierceFirewall { token: r.u32()? }),
        init_code::PEER_INIT => Ok(InitMessage::PeerInit {
            username: r.string()?,
            conn_type: r.string()?,
            token: r.u32()?,
        }),
        _ => Err(ProtoError::UnexpectedEof),
    }
}

// ---- peer messages (u32 code) ----

/// Ask a peer to send us a file (modern flow; direction=download, size 0).
pub fn transfer_request_download(token: u32, filename: &str) -> Vec<u8> {
    Writer::new()
        .u32(DIRECTION_DOWNLOAD)
        .u32(token)
        .string(filename)
        .u64(0)
        .finish_u32(code::TRANSFER_REQUEST)
}

/// Offer to upload one of our files to a peer (direction=upload).
pub fn transfer_request_upload(token: u32, filename: &str, size: u64) -> Vec<u8> {
    Writer::new()
        .u32(DIRECTION_UPLOAD)
        .u32(token)
        .string(filename)
        .u64(size)
        .finish_u32(code::TRANSFER_REQUEST)
}

/// Accept a peer's upload offer (their TransferRequest with direction=upload).
pub fn transfer_response_allowed(token: u32, size: u64) -> Vec<u8> {
    Writer::new()
        .u32(token)
        .u8(1)
        .u64(size)
        .finish_u32(code::TRANSFER_RESPONSE)
}

pub fn transfer_response_rejected(token: u32, reason: &str) -> Vec<u8> {
    Writer::new()
        .u32(token)
        .u8(0)
        .string(reason)
        .finish_u32(code::TRANSFER_RESPONSE)
}

/// Legacy enqueue message; some clients expect it before TransferRequest.
pub fn queue_download(filename: &str) -> Vec<u8> {
    Writer::new()
        .string(filename)
        .finish_u32(code::QUEUE_DOWNLOAD)
}

pub fn place_in_queue_request(filename: &str) -> Vec<u8> {
    Writer::new()
        .string(filename)
        .finish_u32(code::PLACE_IN_QUEUE_REQUEST)
}

pub fn place_in_queue_response(filename: &str, place: u32) -> Vec<u8> {
    Writer::new()
        .string(filename)
        .u32(place)
        .finish_u32(code::PLACE_IN_QUEUE_RESPONSE)
}

pub fn upload_denied(filename: &str, reason: &str) -> Vec<u8> {
    Writer::new()
        .string(filename)
        .string(reason)
        .finish_u32(code::UPLOAD_DENIED)
}

// ---- responder-side encoders (we answer other peers' requests) ----

/// A search reply we send to a searcher (code 9, zlib-compressed).
pub fn search_response(
    username: &str,
    token: u32,
    files: &[File],
    has_free_upload_slot: bool,
    upload_speed: u32,
    queue_length: u32,
) -> Vec<u8> {
    let mut w = Writer::new()
        .string(username)
        .u32(token)
        .u32(files.len() as u32);
    for f in files {
        w = f.write(w);
    }
    w = w
        .bool(has_free_upload_slot)
        .u32(upload_speed)
        .u32(queue_length)
        .u32(0) // unknown, for compatibility
        .u32(0); // locked file count
    w.finish_u32_compressed(code::SEARCH_RESPONSE)
}

/// Our full shared directory listing (code 5, zlib-compressed).
pub fn browse_response(dirs: &[Directory]) -> Vec<u8> {
    let mut w = Writer::new().u32(dirs.len() as u32);
    for d in dirs {
        w = d.write(w);
    }
    w = w.u32(0).u32(0); // unknown + locked directory count
    w.finish_u32_compressed(code::BROWSE_RESPONSE)
}

/// Contents of a single requested folder (code 37, zlib-compressed).
pub fn folder_contents_response(token: u32, folder: &str, dir: &Directory) -> Vec<u8> {
    let w = Writer::new().u32(token).string(folder).u32(1);
    let w = dir.write(w);
    w.finish_u32_compressed(code::FOLDER_CONTENTS_RESPONSE)
}

/// Our user-info card (code 16).
pub fn user_info_response(
    description: &str,
    upload_slots: u32,
    queue_length: u32,
    has_free_upload_slot: bool,
) -> Vec<u8> {
    Writer::new()
        .string(description)
        .u8(0) // has picture = false
        .u32(upload_slots)
        .u32(queue_length)
        .bool(has_free_upload_slot)
        .finish_u32(code::USER_INFO_RESPONSE)
}

#[derive(Debug, Clone)]
pub enum PeerEvent {
    SearchResponse(SearchResponse),
    /// Peer is ready to send us a file (direction=upload) or wants one
    /// from us (direction=download → an upload request we service).
    TransferRequest {
        direction: u32,
        token: u32,
        filename: String,
        size: u64,
    },
    /// Answer to our TransferRequest.
    TransferResponse {
        token: u32,
        allowed: bool,
        size: u64,
        message: String,
    },
    PlaceInQueueResponse {
        filename: String,
        place: u32,
    },
    UploadDenied {
        filename: String,
        reason: String,
    },
    UploadFailed {
        filename: String,
    },
    // Requests other peers make of us (M5 responder):
    BrowseRequest,
    UserInfoRequest,
    FolderContentsRequest {
        token: u32,
        folder: String,
    },
    /// Peer asks us to queue an upload of one of our files to them.
    QueueDownload {
        filename: String,
    },
    PlaceInQueueRequest {
        filename: String,
    },
    Unhandled {
        code: u32,
    },
}

pub fn decode(code: u32, payload: Bytes) -> Result<PeerEvent> {
    // Search responses carry a zlib payload; handle before Reader dispatch.
    if code == code::SEARCH_RESPONSE {
        return Ok(PeerEvent::SearchResponse(decode_search_response(payload)?));
    }
    let mut r = Reader::new(payload);
    Ok(match code {
        code::TRANSFER_REQUEST => {
            let direction = r.u32()?;
            let token = r.u32()?;
            let filename = r.string()?;
            let size = if r.remaining() >= 8 { r.u64()? } else { 0 };
            PeerEvent::TransferRequest {
                direction,
                token,
                filename,
                size,
            }
        }
        code::TRANSFER_RESPONSE => {
            let token = r.u32()?;
            let allowed = r.bool()?;
            if allowed {
                let size = if r.remaining() >= 8 { r.u64()? } else { 0 };
                PeerEvent::TransferResponse {
                    token,
                    allowed,
                    size,
                    message: String::new(),
                }
            } else {
                let message = r.string().unwrap_or_default();
                PeerEvent::TransferResponse {
                    token,
                    allowed,
                    size: 0,
                    message,
                }
            }
        }
        code::PLACE_IN_QUEUE_RESPONSE => PeerEvent::PlaceInQueueResponse {
            filename: r.string()?,
            place: r.u32()?,
        },
        code::UPLOAD_DENIED => PeerEvent::UploadDenied {
            filename: r.string()?,
            reason: r.string().unwrap_or_default(),
        },
        code::UPLOAD_FAILED => PeerEvent::UploadFailed {
            filename: r.string()?,
        },
        code::BROWSE_REQUEST => PeerEvent::BrowseRequest,
        code::USER_INFO_REQUEST => PeerEvent::UserInfoRequest,
        code::FOLDER_CONTENTS_REQUEST => PeerEvent::FolderContentsRequest {
            token: r.u32()?,
            folder: r.string()?,
        },
        code::QUEUE_DOWNLOAD => PeerEvent::QueueDownload {
            filename: r.string()?,
        },
        code::PLACE_IN_QUEUE_REQUEST => PeerEvent::PlaceInQueueRequest {
            filename: r.string()?,
        },
        other => PeerEvent::Unhandled { code: other },
    })
}

/// Payload is zlib-compressed: username, token, files, slot/speed/queue,
/// optional unknown u32, optional locked file list.
fn decode_search_response(payload: Bytes) -> Result<SearchResponse> {
    let mut r = Reader::new(payload);
    r.decompress()?;

    let username = r.string()?;
    let token = r.u32()?;
    let file_count = r.u32()?.min(100_000);
    let mut files = Vec::with_capacity(file_count.min(4096) as usize);
    for _ in 0..file_count {
        files.push(File::read(&mut r)?);
    }
    let has_free_upload_slot = r.bool()?;
    let upload_speed = r.u32()?;
    let queue_length = r.u32()?;
    if r.has_more() {
        let _unknown = r.u32()?;
    }
    let mut locked_files = Vec::new();
    if r.has_more() {
        let locked_count = r.u32()?.min(100_000);
        for _ in 0..locked_count {
            locked_files.push(File::read(&mut r)?);
        }
    }

    Ok(SearchResponse {
        username,
        token,
        files,
        has_free_upload_slot,
        upload_speed,
        queue_length,
        locked_files,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::write::ZlibEncoder;
    use flate2::Compression;
    use std::io::Write as _;

    fn zlib(data: &[u8]) -> Bytes {
        let mut enc = ZlibEncoder::new(Vec::new(), Compression::default());
        enc.write_all(data).unwrap();
        Bytes::from(enc.finish().unwrap())
    }

    #[test]
    fn init_roundtrip() {
        let msg = peer_init("alice", "P", 42);
        let init = decode_init(Bytes::copy_from_slice(&msg[4..])).unwrap();
        match init {
            InitMessage::PeerInit {
                username,
                conn_type,
                token,
            } => {
                assert_eq!(
                    (username.as_str(), conn_type.as_str(), token),
                    ("alice", "P", 42)
                );
            }
            other => panic!("wrong init: {other:?}"),
        }

        let msg = pierce_firewall(7);
        match decode_init(Bytes::copy_from_slice(&msg[4..])).unwrap() {
            InitMessage::PierceFirewall { token } => assert_eq!(token, 7),
            other => panic!("wrong init: {other:?}"),
        }
    }

    #[test]
    fn search_response_decode() {
        // Build the uncompressed body by hand.
        let body = Writer::new()
            .string("bob")
            .u32(555)
            .u32(1)
            // file entry
            .u8(1)
            .string("Music\\Artist\\Album\\01 - Song.flac")
            .u64(31_337_000)
            .string("flac")
            .u32(2)
            .u32(0)
            .u32(1024) // bitrate
            .u32(1)
            .u32(213) // duration
            // trailer
            .bool(true)
            .u32(500_000)
            .u32(0)
            .finish_u32(0);
        let compressed = zlib(&body[8..]); // strip len+code, compress payload only

        let resp = decode_search_response(compressed).unwrap();
        assert_eq!(resp.username, "bob");
        assert_eq!(resp.token, 555);
        assert_eq!(resp.files.len(), 1);
        let f = &resp.files[0];
        assert_eq!(f.size, 31_337_000);
        assert_eq!(f.bitrate(), Some(1024));
        assert_eq!(f.duration_secs(), Some(213));
        assert!(resp.has_free_upload_slot);
    }
}