use std::error::Error;
use std::fmt::{self, Display, Formatter};
use serde::{Deserialize, Serialize};
use crate::RequestId;
pub const INTERNAL_RPC_FORMAT_VERSION: u16 = 1;
pub const MAX_INTERNAL_RPC_FRAME_BYTES: usize = 16 * 1024 * 1024;
pub const MAX_METASTORE_RPC_PAYLOAD_BYTES: usize = 2 * 1024 * 1024;
const INTERNAL_RPC_MAGIC: &[u8; 8] = b"LDSRPC01";
const INTERNAL_RPC_HEADER_BYTES: usize = 8 + 2 + 1 + 2 + 2 + 4 + 4;
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum InternalRpcMethod {
RaftAppend,
RaftVote,
RaftInstallSnapshot,
MetastoreCommand,
MetastoreSnapshot,
WriterIngest,
WriterFence,
Search,
MountManifest,
Health,
}
impl InternalRpcMethod {
const fn as_wire(self) -> &'static str {
match self {
Self::RaftAppend => "raft.append",
Self::RaftVote => "raft.vote",
Self::RaftInstallSnapshot => "raft.install_snapshot",
Self::MetastoreCommand => "metastore.command",
Self::MetastoreSnapshot => "metastore.snapshot",
Self::WriterIngest => "writer.ingest",
Self::WriterFence => "writer.fence",
Self::Search => "search.execute",
Self::MountManifest => "search.mount",
Self::Health => "node.health",
}
}
fn parse(value: &str) -> Result<Self, InternalRpcError> {
match value {
"raft.append" => Ok(Self::RaftAppend),
"raft.vote" => Ok(Self::RaftVote),
"raft.install_snapshot" => Ok(Self::RaftInstallSnapshot),
"metastore.command" => Ok(Self::MetastoreCommand),
"metastore.snapshot" => Ok(Self::MetastoreSnapshot),
"writer.ingest" => Ok(Self::WriterIngest),
"writer.fence" => Ok(Self::WriterFence),
"search.execute" => Ok(Self::Search),
"search.mount" => Ok(Self::MountManifest),
"node.health" => Ok(Self::Health),
_ => Err(InternalRpcError::UnknownMethod(value.to_owned())),
}
}
const fn payload_limit(self) -> usize {
match self {
Self::MetastoreCommand => MAX_METASTORE_RPC_PAYLOAD_BYTES,
_ => MAX_INTERNAL_RPC_FRAME_BYTES - INTERNAL_RPC_HEADER_BYTES,
}
}
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum InternalRpcStatus {
Ok,
Overloaded,
NotLeader,
StaleEpoch,
Unavailable,
Malformed,
UnknownVersion,
Unauthenticated,
PermissionDenied,
WrongCluster,
WrongRole,
StaleNodeIncarnation,
CredentialMismatch,
IncompatibleAuthProfile,
}
impl InternalRpcStatus {
const fn to_wire(self) -> u8 {
match self {
Self::Ok => 0,
Self::Overloaded => 1,
Self::NotLeader => 2,
Self::StaleEpoch => 3,
Self::Unavailable => 4,
Self::Malformed => 5,
Self::UnknownVersion => 6,
Self::Unauthenticated => 7,
Self::PermissionDenied => 8,
Self::WrongCluster => 9,
Self::WrongRole => 10,
Self::StaleNodeIncarnation => 11,
Self::CredentialMismatch => 12,
Self::IncompatibleAuthProfile => 13,
}
}
fn from_wire(value: u8) -> Result<Self, InternalRpcError> {
match value {
0 => Ok(Self::Ok),
1 => Ok(Self::Overloaded),
2 => Ok(Self::NotLeader),
3 => Ok(Self::StaleEpoch),
4 => Ok(Self::Unavailable),
5 => Ok(Self::Malformed),
6 => Ok(Self::UnknownVersion),
7 => Ok(Self::Unauthenticated),
8 => Ok(Self::PermissionDenied),
9 => Ok(Self::WrongCluster),
10 => Ok(Self::WrongRole),
11 => Ok(Self::StaleNodeIncarnation),
12 => Ok(Self::CredentialMismatch),
13 => Ok(Self::IncompatibleAuthProfile),
_ => Err(InternalRpcError::UnknownStatus(value)),
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct InternalRpcFrame {
pub format_version: u16,
pub status: InternalRpcStatus,
pub request_id: RequestId,
pub method: InternalRpcMethod,
pub payload: Vec<u8>,
}
impl InternalRpcFrame {
#[must_use]
pub fn request(request_id: RequestId, method: InternalRpcMethod, payload: Vec<u8>) -> Self {
Self {
format_version: INTERNAL_RPC_FORMAT_VERSION,
status: InternalRpcStatus::Ok,
request_id,
method,
payload,
}
}
pub fn encode(&self) -> Result<Vec<u8>, InternalRpcError> {
if self.format_version != INTERNAL_RPC_FORMAT_VERSION {
return Err(InternalRpcError::UnsupportedVersion(self.format_version));
}
let request_id = self.request_id.as_str().as_bytes();
let method = self.method.as_wire().as_bytes();
if self.payload.len() > self.method.payload_limit() {
return Err(InternalRpcError::PayloadTooLarge {
actual: self.payload.len(),
maximum: self.method.payload_limit(),
});
}
let frame_len = INTERNAL_RPC_HEADER_BYTES
.checked_add(request_id.len())
.and_then(|length| length.checked_add(method.len()))
.and_then(|length| length.checked_add(self.payload.len()))
.ok_or(InternalRpcError::FrameTooLarge(usize::MAX))?;
if frame_len > MAX_INTERNAL_RPC_FRAME_BYTES {
return Err(InternalRpcError::FrameTooLarge(frame_len));
}
let request_id_len =
u16::try_from(request_id.len()).map_err(|_| InternalRpcError::RequestIdTooLong(request_id.len()))?;
let method_len = u16::try_from(method.len()).map_err(|_| InternalRpcError::MethodTooLong(method.len()))?;
let payload_len = u32::try_from(self.payload.len()).map_err(|_| InternalRpcError::FrameTooLarge(frame_len))?;
let mut bytes = Vec::with_capacity(frame_len);
bytes.extend_from_slice(INTERNAL_RPC_MAGIC);
bytes.extend_from_slice(&self.format_version.to_le_bytes());
bytes.push(self.status.to_wire());
bytes.extend_from_slice(&request_id_len.to_le_bytes());
bytes.extend_from_slice(&method_len.to_le_bytes());
bytes.extend_from_slice(&payload_len.to_le_bytes());
bytes.extend_from_slice(&0_u32.to_le_bytes());
bytes.extend_from_slice(request_id);
bytes.extend_from_slice(method);
bytes.extend_from_slice(&self.payload);
let checksum = lodestone_types::crc32c(&bytes[INTERNAL_RPC_HEADER_BYTES..]);
bytes[INTERNAL_RPC_HEADER_BYTES - 4..INTERNAL_RPC_HEADER_BYTES].copy_from_slice(&checksum.to_le_bytes());
Ok(bytes)
}
pub fn decode(bytes: &[u8]) -> Result<Self, InternalRpcError> {
if bytes.len() > MAX_INTERNAL_RPC_FRAME_BYTES {
return Err(InternalRpcError::FrameTooLarge(bytes.len()));
}
if bytes.len() < INTERNAL_RPC_HEADER_BYTES {
return Err(InternalRpcError::Truncated);
}
if bytes.get(..8) != Some(INTERNAL_RPC_MAGIC.as_slice()) {
return Err(InternalRpcError::BadMagic);
}
let format_version = read_u16(bytes, 8)?;
if format_version != INTERNAL_RPC_FORMAT_VERSION {
return Err(InternalRpcError::UnsupportedVersion(format_version));
}
let status = InternalRpcStatus::from_wire(bytes[10])?;
let request_id_len = usize::from(read_u16(bytes, 11)?);
let method_len = usize::from(read_u16(bytes, 13)?);
let payload_len =
usize::try_from(read_u32(bytes, 15)?).map_err(|_| InternalRpcError::FrameTooLarge(bytes.len()))?;
let checksum = read_u32(bytes, 19)?;
let expected_len = INTERNAL_RPC_HEADER_BYTES
.checked_add(request_id_len)
.and_then(|length| length.checked_add(method_len))
.and_then(|length| length.checked_add(payload_len))
.ok_or(InternalRpcError::FrameTooLarge(usize::MAX))?;
if expected_len != bytes.len() {
return Err(InternalRpcError::LengthMismatch {
declared: expected_len,
actual: bytes.len(),
});
}
let actual_checksum = lodestone_types::crc32c(&bytes[INTERNAL_RPC_HEADER_BYTES..]);
if actual_checksum != checksum {
return Err(InternalRpcError::ChecksumMismatch {
expected: checksum,
actual: actual_checksum,
});
}
let request_id_end = INTERNAL_RPC_HEADER_BYTES + request_id_len;
let method_end = request_id_end + method_len;
let request_id = std::str::from_utf8(&bytes[INTERNAL_RPC_HEADER_BYTES..request_id_end])
.map_err(|error| InternalRpcError::InvalidUtf8(error.to_string()))?;
let request_id = RequestId::new(request_id).map_err(|error| InternalRpcError::Malformed(error.to_string()))?;
let method = std::str::from_utf8(&bytes[request_id_end..method_end])
.map_err(|error| InternalRpcError::InvalidUtf8(error.to_string()))?;
let method = InternalRpcMethod::parse(method)?;
if payload_len > method.payload_limit() {
return Err(InternalRpcError::PayloadTooLarge {
actual: payload_len,
maximum: method.payload_limit(),
});
}
Ok(Self {
format_version,
status,
request_id,
method,
payload: bytes[method_end..].to_vec(),
})
}
}
fn read_u16(bytes: &[u8], offset: usize) -> Result<u16, InternalRpcError> {
let value = bytes.get(offset..offset + 2).ok_or(InternalRpcError::Truncated)?;
let mut array = [0_u8; 2];
array.copy_from_slice(value);
Ok(u16::from_le_bytes(array))
}
fn read_u32(bytes: &[u8], offset: usize) -> Result<u32, InternalRpcError> {
let value = bytes.get(offset..offset + 4).ok_or(InternalRpcError::Truncated)?;
let mut array = [0_u8; 4];
array.copy_from_slice(value);
Ok(u32::from_le_bytes(array))
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum InternalRpcError {
BadMagic,
ChecksumMismatch { expected: u32, actual: u32 },
FrameTooLarge(usize),
InvalidUtf8(String),
LengthMismatch { declared: usize, actual: usize },
Malformed(String),
MethodTooLong(usize),
PayloadTooLarge { actual: usize, maximum: usize },
RequestIdTooLong(usize),
Truncated,
UnknownMethod(String),
UnknownStatus(u8),
UnsupportedVersion(u16),
}
impl Display for InternalRpcError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
match self {
Self::BadMagic => write!(formatter, "internal RPC magic is invalid"),
Self::ChecksumMismatch { expected, actual } => {
write!(
formatter,
"internal RPC checksum mismatch: expected {expected}, got {actual}"
)
}
Self::FrameTooLarge(actual) => write!(formatter, "internal RPC frame is too large: {actual} bytes"),
Self::InvalidUtf8(reason) => write!(formatter, "internal RPC string is not UTF-8: {reason}"),
Self::LengthMismatch { declared, actual } => {
write!(
formatter,
"internal RPC length declares {declared} bytes but received {actual}"
)
}
Self::Malformed(reason) => write!(formatter, "internal RPC is malformed: {reason}"),
Self::MethodTooLong(actual) => write!(formatter, "internal RPC method is too long: {actual} bytes"),
Self::PayloadTooLarge { actual, maximum } => {
write!(
formatter,
"internal RPC payload is {actual} bytes; maximum is {maximum}"
)
}
Self::RequestIdTooLong(actual) => {
write!(formatter, "internal RPC request ID is too long: {actual} bytes")
}
Self::Truncated => write!(formatter, "internal RPC frame is truncated"),
Self::UnknownMethod(method) => write!(formatter, "internal RPC method is unknown: {method}"),
Self::UnknownStatus(status) => write!(formatter, "internal RPC status is unknown: {status}"),
Self::UnsupportedVersion(version) => {
write!(formatter, "internal RPC version {version} is unsupported")
}
}
}
}
impl Error for InternalRpcError {}
#[cfg(test)]
mod tests {
use super::{InternalRpcError, InternalRpcFrame, InternalRpcMethod};
use crate::RequestId;
#[test]
fn round_trip_and_checksum_precede_payload_decode() -> Result<(), Box<dyn std::error::Error>> {
let frame = InternalRpcFrame::request(
RequestId::new("rpc-7")?,
InternalRpcMethod::WriterIngest,
br#"{"document":"value"}"#.to_vec(),
);
let bytes = frame.encode()?;
assert_eq!(InternalRpcFrame::decode(&bytes)?, frame);
let mut corrupt = bytes;
let Some(last) = corrupt.last_mut() else {
return Err("encoded frame was empty".into());
};
*last ^= 0xff;
assert!(matches!(
InternalRpcFrame::decode(&corrupt),
Err(InternalRpcError::ChecksumMismatch { .. })
));
Ok(())
}
#[test]
fn malformed_inputs_are_bounded_typed_refusals() -> Result<(), Box<dyn std::error::Error>> {
for length in 0..23 {
assert!(matches!(
InternalRpcFrame::decode(&vec![0_u8; length]),
Err(InternalRpcError::Truncated) | Err(InternalRpcError::BadMagic)
));
}
let mut unknown_version =
InternalRpcFrame::request(RequestId::new("v")?, InternalRpcMethod::Health, Vec::new()).encode()?;
unknown_version[8..10].copy_from_slice(&2_u16.to_le_bytes());
assert!(matches!(
InternalRpcFrame::decode(&unknown_version),
Err(InternalRpcError::UnsupportedVersion(2))
));
Ok(())
}
}