lanspread/crates/lanspread-proto/src/lib.rs

use std::net::SocketAddr;

use bytes::Bytes;
use lanspread_db::db::{Game, GameFileDescription};
use serde::{Deserialize, Serialize};

pub const PROTOCOL_VERSION: u32 = 5;

pub use lanspread_db::db::Availability;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub struct GameSummary {
    pub id: String,
    pub name: String,
    pub size: u64,
    pub downloaded: bool,
    pub installed: bool,
    pub eti_version: Option<String>,
    pub manifest_hash: u64,
    pub availability: Availability,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Hello {
    pub peer_id: String,
    pub proto_ver: u32,
    pub listen_addr: SocketAddr,
    pub library: LibrarySnapshot,
    pub features: Vec<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct HelloAck {
    pub peer_id: String,
    pub proto_ver: u32,
    pub listen_addr: SocketAddr,
    pub library: LibrarySnapshot,
    pub features: Vec<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct LibrarySnapshot {
    pub library_rev: u64,
    pub games: Vec<GameSummary>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct LibraryDelta {
    pub from_rev: u64,
    pub to_rev: u64,
    pub added: Vec<GameSummary>,
    pub updated: Vec<GameSummary>,
    pub removed: Vec<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub enum Request {
    Ping,
    ListGames,
    GetGame {
        id: String,
    },
    GetGameFileData(GameFileDescription),
    GetGameFileChunk {
        game_id: String,
        relative_path: String,
        offset: u64,
        length: u64,
    },
    StreamInstall {
        game_id: String,
    },
    Hello(Hello),
    LibraryDelta {
        peer_id: String,
        delta: LibraryDelta,
    },
    Goodbye {
        peer_id: String,
    },
    Invalid(Bytes, String),
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Response {
    Pong,
    ListGames(Vec<Game>),
    GetGame {
        id: String,
        file_descriptions: Vec<GameFileDescription>,
    },
    HelloAck(HelloAck),
    GameNotFound(String),
    InvalidRequest(Bytes, String),
    EncodingError(String),
    DecodingError(Bytes, String),
    InternalPeerError(String),
}

const STREAM_INSTALL_CONTROL_FRAME_TAG: u8 = 0;
const STREAM_INSTALL_FILE_CHUNK_FRAME_TAG: u8 = 1;
const STREAM_INSTALL_ENCODE_ERROR_FRAME: &[u8] =
    b"\0{\"Error\":{\"message\":\"stream install frame encoding error\"}}";

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum StreamInstallFrame {
    ArchiveBegin {
        archive_name: String,
        solid: bool,
        unpacked_size: u64,
    },
    Directory {
        relative_path: String,
    },
    FileBegin {
        relative_path: String,
        size: u64,
        crc32: u32,
    },
    FileChunk {
        bytes: Bytes,
    },
    FileEnd {
        relative_path: String,
    },
    ArchiveEnd {
        archive_name: String,
    },
    Complete,
    Error {
        message: String,
    },
}

// Add Message trait
pub trait Message {
    fn decode(bytes: Bytes) -> Self;
    fn encode(&self) -> Bytes;
}

// Implement for Request
impl Message for Request {
    fn decode(bytes: Bytes) -> Self {
        match serde_json::from_slice(&bytes) {
            Ok(t) => t,
            Err(e) => {
                tracing::error!(?e, "Request decoding error");
                Request::Invalid(bytes, e.to_string())
            }
        }
    }

    fn encode(&self) -> Bytes {
        match serde_json::to_vec(self) {
            Ok(s) => Bytes::from(s),
            Err(e) => {
                tracing::error!(?e, "Request encoding error");
                Bytes::from(format!(r#"{{"error": "encoding error: {e}"}}"#))
            }
        }
    }
}

// Implement for Response
impl Message for Response {
    fn decode(bytes: Bytes) -> Self {
        match serde_json::from_slice(&bytes) {
            Ok(t) => t,
            Err(e) => {
                tracing::error!(?e, "Response decoding error");
                Response::DecodingError(bytes, e.to_string())
            }
        }
    }

    fn encode(&self) -> Bytes {
        match serde_json::to_vec(self) {
            Ok(s) => Bytes::from(s),
            Err(e) => {
                tracing::error!(?e, "Response encoding error");
                Bytes::from(format!(r#"{{"error": "encoding error: {e}"}}"#))
            }
        }
    }
}

impl Message for StreamInstallFrame {
    fn decode(bytes: Bytes) -> Self {
        if bytes.is_empty() {
            return stream_install_decode_error("stream install frame is empty");
        }

        let tag = bytes[0];
        let payload = bytes.slice(1..);
        match tag {
            STREAM_INSTALL_CONTROL_FRAME_TAG => decode_stream_install_control_frame(&payload),
            STREAM_INSTALL_FILE_CHUNK_FRAME_TAG => StreamInstallFrame::FileChunk { bytes: payload },
            _ => stream_install_decode_error(format!("unknown stream install frame tag {tag}")),
        }
    }

    fn encode(&self) -> Bytes {
        match self {
            StreamInstallFrame::FileChunk { bytes } => {
                tagged_stream_install_frame(STREAM_INSTALL_FILE_CHUNK_FRAME_TAG, bytes)
            }
            _ => match serde_json::to_vec(self) {
                Ok(payload) => {
                    tagged_stream_install_frame(STREAM_INSTALL_CONTROL_FRAME_TAG, &payload)
                }
                Err(e) => {
                    tracing::error!(?e, "StreamInstallFrame encoding error");
                    Bytes::from_static(STREAM_INSTALL_ENCODE_ERROR_FRAME)
                }
            },
        }
    }
}

fn decode_stream_install_control_frame(payload: &[u8]) -> StreamInstallFrame {
    match serde_json::from_slice(payload) {
        Ok(StreamInstallFrame::FileChunk { .. }) => {
            stream_install_decode_error("stream install control frame cannot contain file bytes")
        }
        Ok(frame) => frame,
        Err(e) => {
            tracing::error!(?e, "StreamInstallFrame decoding error");
            stream_install_decode_error(format!("stream install frame decoding error: {e}"))
        }
    }
}

fn tagged_stream_install_frame(tag: u8, payload: &[u8]) -> Bytes {
    let mut frame = Vec::with_capacity(1 + payload.len());
    frame.push(tag);
    frame.extend_from_slice(payload);
    Bytes::from(frame)
}

fn stream_install_decode_error(message: impl Into<String>) -> StreamInstallFrame {
    StreamInstallFrame::Error {
        message: message.into(),
    }
}