use std::error::Error;
use std::fmt;
use std::num::{NonZeroU32, NonZeroUsize};
use serde_json::{Map, Number, Value};
use sha2::{Digest, Sha256};
const FORMAT_VERSION: u16 = 1;
const MAX_SCHEMA_ID_BYTES: usize = 128;
const DEFAULT_MAXIMUM_BYTES: usize = 64 * 1024;
const DEFAULT_MAXIMUM_DEPTH: usize = 16;
const MAXIMUM_BYTES: usize = 1024 * 1024;
const MAXIMUM_DEPTH: usize = 64;
const MAX_UPGRADE_CHAIN: usize = 64;
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[non_exhaustive]
pub enum DurableStateKind {
Checkpoint,
ExecutionContext,
}
impl DurableStateKind {
const fn format(self) -> &'static str {
match self {
Self::Checkpoint => "oxide-batch.checkpoint",
Self::ExecutionContext => "oxide-batch.execution-context",
}
}
}
impl fmt::Display for DurableStateKind {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::Checkpoint => "checkpoint",
Self::ExecutionContext => "execution context",
})
}
}
#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct StateSchemaId(String);
impl StateSchemaId {
pub fn new(value: impl Into<String>) -> Result<Self, StateError> {
let value = value.into();
if value.is_empty() {
return Err(StateError::EmptySchemaId);
}
if value.len() > MAX_SCHEMA_ID_BYTES {
return Err(StateError::SchemaIdTooLong {
max_bytes: MAX_SCHEMA_ID_BYTES,
});
}
if value.trim() != value {
return Err(StateError::SchemaIdHasSurroundingWhitespace);
}
if value.chars().any(char::is_control) {
return Err(StateError::SchemaIdContainsControl);
}
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Debug for StateSchemaId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("StateSchemaId(<redacted>)")
}
}
impl fmt::Display for StateSchemaId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(self.as_str())
}
}
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct StateSchemaVersion(NonZeroU32);
impl StateSchemaVersion {
pub fn new(value: u32) -> Result<Self, StateError> {
NonZeroU32::new(value)
.map(Self)
.ok_or(StateError::ZeroSchemaVersion)
}
#[must_use]
pub const fn get(self) -> u32 {
self.0.get()
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct StateLimits {
maximum_bytes: NonZeroUsize,
maximum_depth: NonZeroUsize,
}
impl StateLimits {
pub fn new(maximum_bytes: usize, maximum_depth: usize) -> Result<Self, StateError> {
if maximum_bytes == 0 || maximum_bytes > MAXIMUM_BYTES {
return Err(StateError::InvalidByteLimit {
maximum: MAXIMUM_BYTES,
});
}
if maximum_depth == 0 || maximum_depth > MAXIMUM_DEPTH {
return Err(StateError::InvalidDepthLimit {
maximum: MAXIMUM_DEPTH,
});
}
let Some(maximum_bytes) = NonZeroUsize::new(maximum_bytes) else {
return Err(StateError::InvalidByteLimit {
maximum: MAXIMUM_BYTES,
});
};
let Some(maximum_depth) = NonZeroUsize::new(maximum_depth) else {
return Err(StateError::InvalidDepthLimit {
maximum: MAXIMUM_DEPTH,
});
};
Ok(Self {
maximum_bytes,
maximum_depth,
})
}
#[must_use]
pub const fn maximum_bytes(self) -> usize {
self.maximum_bytes.get()
}
#[must_use]
pub const fn maximum_depth(self) -> usize {
self.maximum_depth.get()
}
}
impl Default for StateLimits {
fn default() -> Self {
Self {
maximum_bytes: NonZeroUsize::new(DEFAULT_MAXIMUM_BYTES).unwrap_or(NonZeroUsize::MIN),
maximum_depth: NonZeroUsize::new(DEFAULT_MAXIMUM_DEPTH).unwrap_or(NonZeroUsize::MIN),
}
}
}
#[derive(Clone, Copy)]
pub struct StateSchemaUpgrade {
from: StateSchemaVersion,
to: StateSchemaVersion,
apply: fn(&[u8]) -> Result<Vec<u8>, StateCodecError>,
}
impl StateSchemaUpgrade {
pub fn new(
from: StateSchemaVersion,
to: StateSchemaVersion,
apply: fn(&[u8]) -> Result<Vec<u8>, StateCodecError>,
) -> Result<Self, StateError> {
if to <= from {
return Err(StateError::NonIncreasingUpgrade { from, to });
}
Ok(Self { from, to, apply })
}
#[must_use]
pub const fn from(&self) -> StateSchemaVersion {
self.from
}
#[must_use]
pub const fn to(&self) -> StateSchemaVersion {
self.to
}
}
impl fmt::Debug for StateSchemaUpgrade {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("StateSchemaUpgrade")
.field("from", &self.from)
.field("to", &self.to)
.finish_non_exhaustive()
}
}
pub trait VersionedStateCodec<T>: Send + Sync {
fn schema_id(&self) -> &StateSchemaId;
fn current_version(&self) -> StateSchemaVersion;
fn upgrades(&self) -> &[StateSchemaUpgrade] {
&[]
}
fn encode(&self, value: &T) -> Result<Vec<u8>, StateCodecError>;
fn decode(&self, payload: &[u8]) -> Result<T, StateCodecError>;
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum StateCodecError {
InvalidPayload,
UnsupportedSchemaVersion,
}
impl fmt::Display for StateCodecError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::InvalidPayload => "durable state payload is invalid",
Self::UnsupportedSchemaVersion => "durable state schema version is unsupported",
})
}
}
impl Error for StateCodecError {}
#[derive(Clone, Eq, PartialEq)]
struct VersionedState {
schema_id: StateSchemaId,
schema_version: StateSchemaVersion,
payload: Value,
encoded_bytes: usize,
}
impl VersionedState {
fn encode<T>(
kind: DurableStateKind,
value: &T,
codec: &(impl VersionedStateCodec<T> + ?Sized),
limits: StateLimits,
) -> Result<Self, StateError> {
let payload_bytes = codec.encode(value).map_err(StateError::Codec)?;
let payload: Value =
serde_json::from_slice(&payload_bytes).map_err(|_| StateError::InvalidPayload)?;
if !payload.is_object() {
return Err(StateError::PayloadNotObject);
}
Self::from_parts(
kind,
codec.schema_id().clone(),
codec.current_version(),
payload,
limits,
)
}
fn from_json(
kind: DurableStateKind,
bytes: &[u8],
limits: StateLimits,
) -> Result<Self, StateError> {
if bytes.len() > limits.maximum_bytes() {
return Err(StateError::TooLarge {
kind,
max_bytes: limits.maximum_bytes(),
});
}
let value: Value =
serde_json::from_slice(bytes).map_err(|_| StateError::Malformed { kind })?;
if json_depth(&value) > limits.maximum_depth() {
return Err(StateError::TooDeep {
kind,
max_depth: limits.maximum_depth(),
});
}
let object = value.as_object().ok_or(StateError::Malformed { kind })?;
let format = object
.get("format")
.and_then(Value::as_str)
.ok_or(StateError::Malformed { kind })?;
if format != kind.format() {
return Err(StateError::FormatMismatch { kind });
}
let format_version = object
.get("format_version")
.and_then(Value::as_u64)
.and_then(|version| u16::try_from(version).ok())
.ok_or(StateError::Malformed { kind })?;
if format_version != FORMAT_VERSION {
return Err(StateError::UnsupportedFormatVersion {
kind,
version: format_version,
});
}
let schema_id = object
.get("schema")
.and_then(Value::as_str)
.ok_or(StateError::Malformed { kind })?;
let schema_id = StateSchemaId::new(schema_id)?;
let schema_version = object
.get("schema_version")
.and_then(Value::as_u64)
.and_then(|version| u32::try_from(version).ok())
.ok_or(StateError::Malformed { kind })?;
let schema_version = StateSchemaVersion::new(schema_version)?;
let payload = object
.get("payload")
.cloned()
.ok_or(StateError::Malformed { kind })?;
if !payload.is_object() {
return Err(StateError::PayloadNotObject);
}
Ok(Self {
schema_id,
schema_version,
payload,
encoded_bytes: bytes.len(),
})
}
fn from_parts(
kind: DurableStateKind,
schema_id: StateSchemaId,
schema_version: StateSchemaVersion,
payload: Value,
limits: StateLimits,
) -> Result<Self, StateError> {
let envelope = envelope(kind, &schema_id, schema_version, payload.clone());
let bytes = serde_json::to_vec(&envelope).map_err(|_| StateError::Malformed { kind })?;
if bytes.len() > limits.maximum_bytes() {
return Err(StateError::TooLarge {
kind,
max_bytes: limits.maximum_bytes(),
});
}
if json_depth(&envelope) > limits.maximum_depth() {
return Err(StateError::TooDeep {
kind,
max_depth: limits.maximum_depth(),
});
}
Ok(Self {
schema_id,
schema_version,
payload,
encoded_bytes: bytes.len(),
})
}
fn decode<T>(
&self,
kind: DurableStateKind,
codec: &(impl VersionedStateCodec<T> + ?Sized),
) -> Result<T, StateError> {
if &self.schema_id != codec.schema_id() {
return Err(StateError::SchemaMismatch { kind });
}
let current = codec.current_version();
if self.schema_version > current {
return Err(StateError::UnsupportedSchemaVersion {
kind,
found: self.schema_version,
current,
});
}
let payload =
serde_json::to_vec(&self.payload).map_err(|_| StateError::Malformed { kind })?;
let payload = self.upgrade(kind, codec, payload)?;
codec.decode(&payload).map_err(StateError::Codec)
}
fn upgrade<T>(
&self,
kind: DurableStateKind,
codec: &(impl VersionedStateCodec<T> + ?Sized),
mut payload: Vec<u8>,
) -> Result<Vec<u8>, StateError> {
let current = codec.current_version();
let upgrades = codec.upgrades();
let mut version = self.schema_version;
let mut applied = 0_usize;
while version < current {
let mut edges = upgrades.iter().filter(|upgrade| upgrade.from == version);
let edge = edges.next().ok_or(StateError::NoUpgradePath {
kind,
found: version,
current,
})?;
if edges.next().is_some() {
return Err(StateError::AmbiguousUpgrade {
kind,
from: version,
});
}
if edge.to > current {
return Err(StateError::UpgradeOvershootsCurrent {
kind,
to: edge.to,
current,
});
}
applied += 1;
if applied > MAX_UPGRADE_CHAIN {
return Err(StateError::UpgradeChainTooLong {
kind,
max_upgrades: MAX_UPGRADE_CHAIN,
});
}
payload = (edge.apply)(&payload).map_err(StateError::Codec)?;
check_upgraded(kind, &payload)?;
version = edge.to;
}
Ok(payload)
}
fn to_json(&self, kind: DurableStateKind) -> Result<Vec<u8>, StateError> {
serde_json::to_vec(&envelope(
kind,
&self.schema_id,
self.schema_version,
self.payload.clone(),
))
.map_err(|_| StateError::Malformed { kind })
}
fn payload_json(&self, kind: DurableStateKind) -> Result<Vec<u8>, StateError> {
serde_json::to_vec(&self.payload).map_err(|_| StateError::Malformed { kind })
}
}
fn envelope(
kind: DurableStateKind,
schema_id: &StateSchemaId,
schema_version: StateSchemaVersion,
payload: Value,
) -> Value {
let mut object = Map::new();
object.insert(
String::from("format"),
Value::String(String::from(kind.format())),
);
object.insert(
String::from("format_version"),
Value::Number(Number::from(FORMAT_VERSION)),
);
object.insert(
String::from("schema"),
Value::String(String::from(schema_id.as_str())),
);
object.insert(
String::from("schema_version"),
Value::Number(Number::from(schema_version.get())),
);
object.insert(String::from("payload"), payload);
Value::Object(object)
}
fn check_upgraded(kind: DurableStateKind, payload: &[u8]) -> Result<(), StateError> {
if payload.len() > MAXIMUM_BYTES {
return Err(StateError::TooLarge {
kind,
max_bytes: MAXIMUM_BYTES,
});
}
let value: Value = serde_json::from_slice(payload)
.map_err(|_| StateError::UpgradeProducedInvalidJson { kind })?;
if !value.is_object() {
return Err(StateError::PayloadNotObject);
}
if json_depth(&value) > MAXIMUM_DEPTH {
return Err(StateError::TooDeep {
kind,
max_depth: MAXIMUM_DEPTH,
});
}
Ok(())
}
fn json_depth(value: &Value) -> usize {
match value {
Value::Array(values) => 1 + values.iter().map(json_depth).max().unwrap_or_default(),
Value::Object(values) => 1 + values.values().map(json_depth).max().unwrap_or_default(),
_ => 1,
}
}
macro_rules! durable_state {
($name:ident, $kind:expr, $docs:literal) => {
#[doc = $docs]
#[derive(Clone, Eq, PartialEq)]
pub struct $name(VersionedState);
impl $name {
pub fn encode<T>(
value: &T,
codec: &(impl VersionedStateCodec<T> + ?Sized),
limits: StateLimits,
) -> Result<Self, StateError> {
VersionedState::encode($kind, value, codec, limits).map(Self)
}
pub fn from_json(bytes: &[u8], limits: StateLimits) -> Result<Self, StateError> {
VersionedState::from_json($kind, bytes, limits).map(Self)
}
pub fn decode<T>(
&self,
codec: &(impl VersionedStateCodec<T> + ?Sized),
) -> Result<T, StateError> {
self.0.decode($kind, codec)
}
#[must_use]
pub const fn format_version(&self) -> u16 {
FORMAT_VERSION
}
#[must_use]
pub const fn schema_id(&self) -> &StateSchemaId {
&self.0.schema_id
}
#[must_use]
pub const fn schema_version(&self) -> StateSchemaVersion {
self.0.schema_version
}
#[must_use]
pub const fn encoded_len(&self) -> usize {
self.0.encoded_bytes
}
pub fn to_json(&self) -> Result<Vec<u8>, StateError> {
self.0.to_json($kind)
}
pub fn payload_json(&self) -> Result<Vec<u8>, StateError> {
self.0.payload_json($kind)
}
}
impl fmt::Debug for $name {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct(stringify!($name))
.field("format_version", &FORMAT_VERSION)
.field("schema_version", &self.schema_version())
.field("encoded_bytes", &self.encoded_len())
.field("payload", &"<redacted>")
.finish()
}
}
};
}
durable_state!(
Checkpoint,
DurableStateKind::Checkpoint,
"A bounded, versioned reader position committed with a chunk."
);
durable_state!(
ExecutionContext,
DurableStateKind::ExecutionContext,
"Bounded, versioned application restart state committed with a chunk."
);
impl Checkpoint {
#[must_use]
pub fn generation_digest(&self) -> [u8; 32] {
self.to_json()
.map_or([0; 32], |bytes| Sha256::digest(&bytes).into())
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum StateError {
EmptySchemaId,
SchemaIdTooLong {
max_bytes: usize,
},
SchemaIdHasSurroundingWhitespace,
SchemaIdContainsControl,
ZeroSchemaVersion,
InvalidByteLimit {
maximum: usize,
},
InvalidDepthLimit {
maximum: usize,
},
TooLarge {
kind: DurableStateKind,
max_bytes: usize,
},
TooDeep {
kind: DurableStateKind,
max_depth: usize,
},
Malformed {
kind: DurableStateKind,
},
FormatMismatch {
kind: DurableStateKind,
},
UnsupportedFormatVersion {
kind: DurableStateKind,
version: u16,
},
SchemaMismatch {
kind: DurableStateKind,
},
NonIncreasingUpgrade {
from: StateSchemaVersion,
to: StateSchemaVersion,
},
NoUpgradePath {
kind: DurableStateKind,
found: StateSchemaVersion,
current: StateSchemaVersion,
},
AmbiguousUpgrade {
kind: DurableStateKind,
from: StateSchemaVersion,
},
UpgradeOvershootsCurrent {
kind: DurableStateKind,
to: StateSchemaVersion,
current: StateSchemaVersion,
},
UpgradeChainTooLong {
kind: DurableStateKind,
max_upgrades: usize,
},
UpgradeProducedInvalidJson {
kind: DurableStateKind,
},
UnsupportedSchemaVersion {
kind: DurableStateKind,
found: StateSchemaVersion,
current: StateSchemaVersion,
},
InvalidPayload,
PayloadNotObject,
Codec(StateCodecError),
}
impl fmt::Display for StateError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::EmptySchemaId => formatter.write_str("state schema identifier must not be empty"),
Self::SchemaIdTooLong { max_bytes } => {
write!(
formatter,
"state schema identifier exceeds {max_bytes} UTF-8 bytes"
)
}
Self::SchemaIdHasSurroundingWhitespace => {
formatter.write_str("state schema identifier has surrounding whitespace")
}
Self::SchemaIdContainsControl => {
formatter.write_str("state schema identifier contains a control character")
}
Self::ZeroSchemaVersion => formatter.write_str("state schema version must be nonzero"),
Self::InvalidByteLimit { maximum } => {
write!(
formatter,
"state byte limit must be between 1 and {maximum}"
)
}
Self::InvalidDepthLimit { maximum } => {
write!(
formatter,
"state depth limit must be between 1 and {maximum}"
)
}
Self::TooLarge { kind, max_bytes } => {
write!(formatter, "{kind} exceeds {max_bytes} bytes")
}
Self::TooDeep { kind, max_depth } => {
write!(formatter, "{kind} exceeds JSON depth {max_depth}")
}
Self::Malformed { kind } => write!(formatter, "{kind} is malformed"),
Self::FormatMismatch { kind } => {
write!(formatter, "durable state is not a {kind}")
}
Self::UnsupportedFormatVersion { kind, .. } => {
write!(formatter, "{kind} format version is unsupported")
}
Self::SchemaMismatch { kind } => {
write!(formatter, "{kind} schema does not match the component")
}
Self::NonIncreasingUpgrade { .. } => {
formatter.write_str("state schema upgrade must increase the version")
}
Self::NoUpgradePath { kind, .. } => {
write!(formatter, "{kind} schema version has no upgrade path")
}
Self::AmbiguousUpgrade { kind, .. } => {
write!(formatter, "{kind} schema upgrade is ambiguous")
}
Self::UpgradeOvershootsCurrent { kind, .. } => {
write!(
formatter,
"{kind} schema upgrade passes the current version"
)
}
Self::UpgradeChainTooLong { kind, max_upgrades } => {
write!(
formatter,
"{kind} schema upgrade chain exceeds {max_upgrades} upgrades"
)
}
Self::UpgradeProducedInvalidJson { kind } => {
write!(formatter, "{kind} schema upgrade produced invalid JSON")
}
Self::UnsupportedSchemaVersion { kind, .. } => {
write!(formatter, "{kind} schema version is unsupported")
}
Self::InvalidPayload => formatter.write_str("durable state payload is not valid JSON"),
Self::PayloadNotObject => {
formatter.write_str("durable state payload must be a JSON object")
}
Self::Codec(error) => error.fmt(formatter),
}
}
}
impl Error for StateError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::Codec(error) => Some(error),
_ => None,
}
}
}