smartdb/rust/crates/rustdb-wire/src/encoder.rs

use bson::Document;

use crate::opcodes::*;

/// Encode an OP_MSG response.
pub fn encode_op_msg_response(
    response_to: i32,
    response: &Document,
    request_id: i32,
) -> Vec<u8> {
    let body_bson = bson::to_vec(response).expect("failed to serialize BSON response");

    // Header (16) + flagBits (4) + section type (1) + body BSON
    let message_length = 16 + 4 + 1 + body_bson.len();

    let mut buf = Vec::with_capacity(message_length);

    // Header
    buf.extend_from_slice(&(message_length as i32).to_le_bytes());
    buf.extend_from_slice(&request_id.to_le_bytes());
    buf.extend_from_slice(&response_to.to_le_bytes());
    buf.extend_from_slice(&OP_MSG.to_le_bytes());

    // Flag bits (0 = no flags)
    buf.extend_from_slice(&0u32.to_le_bytes());

    // Section type 0 (body)
    buf.push(SECTION_BODY);

    // Body BSON
    buf.extend_from_slice(&body_bson);

    buf
}

/// Encode an OP_REPLY response (legacy, for OP_QUERY responses).
pub fn encode_op_reply_response(
    response_to: i32,
    documents: &[Document],
    request_id: i32,
    cursor_id: i64,
) -> Vec<u8> {
    let doc_buffers: Vec<Vec<u8>> = documents
        .iter()
        .map(|doc| bson::to_vec(doc).expect("failed to serialize BSON document"))
        .collect();
    let total_docs_size: usize = doc_buffers.iter().map(|b| b.len()).sum();

    // Header (16) + responseFlags (4) + cursorID (8) + startingFrom (4) + numberReturned (4) + docs
    let message_length = 16 + 4 + 8 + 4 + 4 + total_docs_size;

    let mut buf = Vec::with_capacity(message_length);

    // Header
    buf.extend_from_slice(&(message_length as i32).to_le_bytes());
    buf.extend_from_slice(&request_id.to_le_bytes());
    buf.extend_from_slice(&response_to.to_le_bytes());
    buf.extend_from_slice(&OP_REPLY.to_le_bytes());

    // OP_REPLY fields
    buf.extend_from_slice(&0i32.to_le_bytes()); // responseFlags
    buf.extend_from_slice(&cursor_id.to_le_bytes()); // cursorID
    buf.extend_from_slice(&0i32.to_le_bytes()); // startingFrom
    buf.extend_from_slice(&(documents.len() as i32).to_le_bytes()); // numberReturned

    // Documents
    for doc_buf in &doc_buffers {
        buf.extend_from_slice(doc_buf);
    }

    buf
}

/// Encode an error response as OP_MSG.
pub fn encode_error_response(
    response_to: i32,
    error_code: i32,
    error_message: &str,
    request_id: i32,
) -> Vec<u8> {
    let response = bson::doc! {
        "ok": 0,
        "errmsg": error_message,
        "code": error_code,
        "codeName": error_code_name(error_code),
    };
    encode_op_msg_response(response_to, &response, request_id)
}

/// Map error codes to their code names.
pub fn error_code_name(code: i32) -> &'static str {
    match code {
        0 => "OK",
        1 => "InternalError",
        2 => "BadValue",
        13 => "Unauthorized",
        26 => "NamespaceNotFound",
        27 => "IndexNotFound",
        48 => "NamespaceExists",
        59 => "CommandNotFound",
        66 => "ImmutableField",
        73 => "InvalidNamespace",
        85 => "IndexOptionsConflict",
        112 => "WriteConflict",
        121 => "DocumentValidationFailure",
        211 => "KeyNotFound",
        251 => "NoSuchTransaction",
        11000 => "DuplicateKey",
        11001 => "DuplicateKeyValue",
        _ => "UnknownError",
    }
}

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

    #[test]
    fn test_encode_op_msg_roundtrip() {
        let doc = bson::doc! { "ok": 1 };
        let encoded = encode_op_msg_response(1, &doc, 2);

        // Verify header
        let msg_len = i32::from_le_bytes([encoded[0], encoded[1], encoded[2], encoded[3]]);
        assert_eq!(msg_len as usize, encoded.len());

        let op_code = i32::from_le_bytes([encoded[12], encoded[13], encoded[14], encoded[15]]);
        assert_eq!(op_code, OP_MSG);
    }

    #[test]
    fn test_encode_op_reply() {
        let docs = vec![bson::doc! { "ok": 1 }];
        let encoded = encode_op_reply_response(1, &docs, 2, 0);

        let msg_len = i32::from_le_bytes([encoded[0], encoded[1], encoded[2], encoded[3]]);
        assert_eq!(msg_len as usize, encoded.len());

        let op_code = i32::from_le_bytes([encoded[12], encoded[13], encoded[14], encoded[15]]);
        assert_eq!(op_code, OP_REPLY);
    }
}