1use std::error::Error;
4use std::fmt;
5use std::num::{NonZeroU32, NonZeroUsize};
6
7use serde_json::{Map, Number, Value};
8use sha2::{Digest, Sha256};
9
10const FORMAT_VERSION: u16 = 1;
11const MAX_SCHEMA_ID_BYTES: usize = 128;
12const DEFAULT_MAXIMUM_BYTES: usize = 64 * 1024;
13const DEFAULT_MAXIMUM_DEPTH: usize = 16;
14const MAXIMUM_BYTES: usize = 1024 * 1024;
15const MAXIMUM_DEPTH: usize = 64;
16const MAX_UPGRADE_CHAIN: usize = 64;
22
23#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
25#[non_exhaustive]
26pub enum DurableStateKind {
27 Checkpoint,
29 ExecutionContext,
31}
32
33impl DurableStateKind {
34 const fn format(self) -> &'static str {
35 match self {
36 Self::Checkpoint => "oxide-batch.checkpoint",
37 Self::ExecutionContext => "oxide-batch.execution-context",
38 }
39 }
40}
41
42impl fmt::Display for DurableStateKind {
43 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
44 formatter.write_str(match self {
45 Self::Checkpoint => "checkpoint",
46 Self::ExecutionContext => "execution context",
47 })
48 }
49}
50
51#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
53pub struct StateSchemaId(String);
54
55impl StateSchemaId {
56 pub fn new(value: impl Into<String>) -> Result<Self, StateError> {
64 let value = value.into();
65 if value.is_empty() {
66 return Err(StateError::EmptySchemaId);
67 }
68 if value.len() > MAX_SCHEMA_ID_BYTES {
69 return Err(StateError::SchemaIdTooLong {
70 max_bytes: MAX_SCHEMA_ID_BYTES,
71 });
72 }
73 if value.trim() != value {
74 return Err(StateError::SchemaIdHasSurroundingWhitespace);
75 }
76 if value.chars().any(char::is_control) {
77 return Err(StateError::SchemaIdContainsControl);
78 }
79 Ok(Self(value))
80 }
81
82 #[must_use]
84 pub fn as_str(&self) -> &str {
85 &self.0
86 }
87}
88
89impl fmt::Debug for StateSchemaId {
90 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
91 formatter.write_str("StateSchemaId(<redacted>)")
92 }
93}
94
95impl fmt::Display for StateSchemaId {
96 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
97 formatter.write_str(self.as_str())
98 }
99}
100
101#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
103pub struct StateSchemaVersion(NonZeroU32);
104
105impl StateSchemaVersion {
106 pub fn new(value: u32) -> Result<Self, StateError> {
112 NonZeroU32::new(value)
113 .map(Self)
114 .ok_or(StateError::ZeroSchemaVersion)
115 }
116
117 #[must_use]
119 pub const fn get(self) -> u32 {
120 self.0.get()
121 }
122}
123
124#[derive(Clone, Copy, Debug, Eq, PartialEq)]
126pub struct StateLimits {
127 maximum_bytes: NonZeroUsize,
128 maximum_depth: NonZeroUsize,
129}
130
131impl StateLimits {
132 pub fn new(maximum_bytes: usize, maximum_depth: usize) -> Result<Self, StateError> {
143 if maximum_bytes == 0 || maximum_bytes > MAXIMUM_BYTES {
144 return Err(StateError::InvalidByteLimit {
145 maximum: MAXIMUM_BYTES,
146 });
147 }
148 if maximum_depth == 0 || maximum_depth > MAXIMUM_DEPTH {
149 return Err(StateError::InvalidDepthLimit {
150 maximum: MAXIMUM_DEPTH,
151 });
152 }
153 let Some(maximum_bytes) = NonZeroUsize::new(maximum_bytes) else {
154 return Err(StateError::InvalidByteLimit {
155 maximum: MAXIMUM_BYTES,
156 });
157 };
158 let Some(maximum_depth) = NonZeroUsize::new(maximum_depth) else {
159 return Err(StateError::InvalidDepthLimit {
160 maximum: MAXIMUM_DEPTH,
161 });
162 };
163 Ok(Self {
164 maximum_bytes,
165 maximum_depth,
166 })
167 }
168
169 #[must_use]
171 pub const fn maximum_bytes(self) -> usize {
172 self.maximum_bytes.get()
173 }
174
175 #[must_use]
177 pub const fn maximum_depth(self) -> usize {
178 self.maximum_depth.get()
179 }
180}
181
182impl Default for StateLimits {
183 fn default() -> Self {
184 Self {
185 maximum_bytes: NonZeroUsize::new(DEFAULT_MAXIMUM_BYTES).unwrap_or(NonZeroUsize::MIN),
186 maximum_depth: NonZeroUsize::new(DEFAULT_MAXIMUM_DEPTH).unwrap_or(NonZeroUsize::MIN),
187 }
188 }
189}
190
191#[derive(Clone, Copy)]
201pub struct StateSchemaUpgrade {
202 from: StateSchemaVersion,
203 to: StateSchemaVersion,
204 apply: fn(&[u8]) -> Result<Vec<u8>, StateCodecError>,
205}
206
207impl StateSchemaUpgrade {
208 pub fn new(
219 from: StateSchemaVersion,
220 to: StateSchemaVersion,
221 apply: fn(&[u8]) -> Result<Vec<u8>, StateCodecError>,
222 ) -> Result<Self, StateError> {
223 if to <= from {
224 return Err(StateError::NonIncreasingUpgrade { from, to });
225 }
226 Ok(Self { from, to, apply })
227 }
228
229 #[must_use]
231 pub const fn from(&self) -> StateSchemaVersion {
232 self.from
233 }
234
235 #[must_use]
237 pub const fn to(&self) -> StateSchemaVersion {
238 self.to
239 }
240}
241
242impl fmt::Debug for StateSchemaUpgrade {
243 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
244 formatter
245 .debug_struct("StateSchemaUpgrade")
246 .field("from", &self.from)
247 .field("to", &self.to)
248 .finish_non_exhaustive()
249 }
250}
251
252pub trait VersionedStateCodec<T>: Send + Sync {
265 fn schema_id(&self) -> &StateSchemaId;
267
268 fn current_version(&self) -> StateSchemaVersion;
270
271 fn upgrades(&self) -> &[StateSchemaUpgrade] {
278 &[]
279 }
280
281 fn encode(&self, value: &T) -> Result<Vec<u8>, StateCodecError>;
287
288 fn decode(&self, payload: &[u8]) -> Result<T, StateCodecError>;
296}
297
298#[derive(Clone, Copy, Debug, Eq, PartialEq)]
300#[non_exhaustive]
301pub enum StateCodecError {
302 InvalidPayload,
304 UnsupportedSchemaVersion,
306}
307
308impl fmt::Display for StateCodecError {
309 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
310 formatter.write_str(match self {
311 Self::InvalidPayload => "durable state payload is invalid",
312 Self::UnsupportedSchemaVersion => "durable state schema version is unsupported",
313 })
314 }
315}
316
317impl Error for StateCodecError {}
318
319#[derive(Clone, Eq, PartialEq)]
320struct VersionedState {
321 schema_id: StateSchemaId,
322 schema_version: StateSchemaVersion,
323 payload: Value,
324 encoded_bytes: usize,
325}
326
327impl VersionedState {
328 fn encode<T>(
329 kind: DurableStateKind,
330 value: &T,
331 codec: &(impl VersionedStateCodec<T> + ?Sized),
332 limits: StateLimits,
333 ) -> Result<Self, StateError> {
334 let payload_bytes = codec.encode(value).map_err(StateError::Codec)?;
335 let payload: Value =
336 serde_json::from_slice(&payload_bytes).map_err(|_| StateError::InvalidPayload)?;
337 if !payload.is_object() {
338 return Err(StateError::PayloadNotObject);
339 }
340 Self::from_parts(
341 kind,
342 codec.schema_id().clone(),
343 codec.current_version(),
344 payload,
345 limits,
346 )
347 }
348
349 fn from_json(
350 kind: DurableStateKind,
351 bytes: &[u8],
352 limits: StateLimits,
353 ) -> Result<Self, StateError> {
354 if bytes.len() > limits.maximum_bytes() {
355 return Err(StateError::TooLarge {
356 kind,
357 max_bytes: limits.maximum_bytes(),
358 });
359 }
360 let value: Value =
361 serde_json::from_slice(bytes).map_err(|_| StateError::Malformed { kind })?;
362 if json_depth(&value) > limits.maximum_depth() {
363 return Err(StateError::TooDeep {
364 kind,
365 max_depth: limits.maximum_depth(),
366 });
367 }
368 let object = value.as_object().ok_or(StateError::Malformed { kind })?;
369 let format = object
370 .get("format")
371 .and_then(Value::as_str)
372 .ok_or(StateError::Malformed { kind })?;
373 if format != kind.format() {
374 return Err(StateError::FormatMismatch { kind });
375 }
376 let format_version = object
377 .get("format_version")
378 .and_then(Value::as_u64)
379 .and_then(|version| u16::try_from(version).ok())
380 .ok_or(StateError::Malformed { kind })?;
381 if format_version != FORMAT_VERSION {
382 return Err(StateError::UnsupportedFormatVersion {
383 kind,
384 version: format_version,
385 });
386 }
387 let schema_id = object
388 .get("schema")
389 .and_then(Value::as_str)
390 .ok_or(StateError::Malformed { kind })?;
391 let schema_id = StateSchemaId::new(schema_id)?;
392 let schema_version = object
393 .get("schema_version")
394 .and_then(Value::as_u64)
395 .and_then(|version| u32::try_from(version).ok())
396 .ok_or(StateError::Malformed { kind })?;
397 let schema_version = StateSchemaVersion::new(schema_version)?;
398 let payload = object
399 .get("payload")
400 .cloned()
401 .ok_or(StateError::Malformed { kind })?;
402 if !payload.is_object() {
403 return Err(StateError::PayloadNotObject);
404 }
405 Ok(Self {
406 schema_id,
407 schema_version,
408 payload,
409 encoded_bytes: bytes.len(),
410 })
411 }
412
413 fn from_parts(
414 kind: DurableStateKind,
415 schema_id: StateSchemaId,
416 schema_version: StateSchemaVersion,
417 payload: Value,
418 limits: StateLimits,
419 ) -> Result<Self, StateError> {
420 let envelope = envelope(kind, &schema_id, schema_version, payload.clone());
421 let bytes = serde_json::to_vec(&envelope).map_err(|_| StateError::Malformed { kind })?;
422 if bytes.len() > limits.maximum_bytes() {
423 return Err(StateError::TooLarge {
424 kind,
425 max_bytes: limits.maximum_bytes(),
426 });
427 }
428 if json_depth(&envelope) > limits.maximum_depth() {
429 return Err(StateError::TooDeep {
430 kind,
431 max_depth: limits.maximum_depth(),
432 });
433 }
434 Ok(Self {
435 schema_id,
436 schema_version,
437 payload,
438 encoded_bytes: bytes.len(),
439 })
440 }
441
442 fn decode<T>(
443 &self,
444 kind: DurableStateKind,
445 codec: &(impl VersionedStateCodec<T> + ?Sized),
446 ) -> Result<T, StateError> {
447 if &self.schema_id != codec.schema_id() {
448 return Err(StateError::SchemaMismatch { kind });
449 }
450 let current = codec.current_version();
451 if self.schema_version > current {
452 return Err(StateError::UnsupportedSchemaVersion {
453 kind,
454 found: self.schema_version,
455 current,
456 });
457 }
458 let payload =
459 serde_json::to_vec(&self.payload).map_err(|_| StateError::Malformed { kind })?;
460 let payload = self.upgrade(kind, codec, payload)?;
461 codec.decode(&payload).map_err(StateError::Codec)
462 }
463
464 fn upgrade<T>(
470 &self,
471 kind: DurableStateKind,
472 codec: &(impl VersionedStateCodec<T> + ?Sized),
473 mut payload: Vec<u8>,
474 ) -> Result<Vec<u8>, StateError> {
475 let current = codec.current_version();
476 let upgrades = codec.upgrades();
477 let mut version = self.schema_version;
478 let mut applied = 0_usize;
479 while version < current {
480 let mut edges = upgrades.iter().filter(|upgrade| upgrade.from == version);
481 let edge = edges.next().ok_or(StateError::NoUpgradePath {
482 kind,
483 found: version,
484 current,
485 })?;
486 if edges.next().is_some() {
487 return Err(StateError::AmbiguousUpgrade {
488 kind,
489 from: version,
490 });
491 }
492 if edge.to > current {
493 return Err(StateError::UpgradeOvershootsCurrent {
494 kind,
495 to: edge.to,
496 current,
497 });
498 }
499 applied += 1;
500 if applied > MAX_UPGRADE_CHAIN {
501 return Err(StateError::UpgradeChainTooLong {
502 kind,
503 max_upgrades: MAX_UPGRADE_CHAIN,
504 });
505 }
506 payload = (edge.apply)(&payload).map_err(StateError::Codec)?;
507 check_upgraded(kind, &payload)?;
508 version = edge.to;
509 }
510 Ok(payload)
511 }
512
513 fn to_json(&self, kind: DurableStateKind) -> Result<Vec<u8>, StateError> {
514 serde_json::to_vec(&envelope(
515 kind,
516 &self.schema_id,
517 self.schema_version,
518 self.payload.clone(),
519 ))
520 .map_err(|_| StateError::Malformed { kind })
521 }
522
523 fn payload_json(&self, kind: DurableStateKind) -> Result<Vec<u8>, StateError> {
524 serde_json::to_vec(&self.payload).map_err(|_| StateError::Malformed { kind })
525 }
526}
527
528fn envelope(
529 kind: DurableStateKind,
530 schema_id: &StateSchemaId,
531 schema_version: StateSchemaVersion,
532 payload: Value,
533) -> Value {
534 let mut object = Map::new();
535 object.insert(
536 String::from("format"),
537 Value::String(String::from(kind.format())),
538 );
539 object.insert(
540 String::from("format_version"),
541 Value::Number(Number::from(FORMAT_VERSION)),
542 );
543 object.insert(
544 String::from("schema"),
545 Value::String(String::from(schema_id.as_str())),
546 );
547 object.insert(
548 String::from("schema_version"),
549 Value::Number(Number::from(schema_version.get())),
550 );
551 object.insert(String::from("payload"), payload);
552 Value::Object(object)
553}
554
555fn check_upgraded(kind: DurableStateKind, payload: &[u8]) -> Result<(), StateError> {
565 if payload.len() > MAXIMUM_BYTES {
566 return Err(StateError::TooLarge {
567 kind,
568 max_bytes: MAXIMUM_BYTES,
569 });
570 }
571 let value: Value = serde_json::from_slice(payload)
572 .map_err(|_| StateError::UpgradeProducedInvalidJson { kind })?;
573 if !value.is_object() {
574 return Err(StateError::PayloadNotObject);
575 }
576 if json_depth(&value) > MAXIMUM_DEPTH {
577 return Err(StateError::TooDeep {
578 kind,
579 max_depth: MAXIMUM_DEPTH,
580 });
581 }
582 Ok(())
583}
584
585fn json_depth(value: &Value) -> usize {
586 match value {
587 Value::Array(values) => 1 + values.iter().map(json_depth).max().unwrap_or_default(),
588 Value::Object(values) => 1 + values.values().map(json_depth).max().unwrap_or_default(),
589 _ => 1,
590 }
591}
592
593macro_rules! durable_state {
594 ($name:ident, $kind:expr, $docs:literal) => {
595 #[doc = $docs]
596 #[derive(Clone, Eq, PartialEq)]
597 pub struct $name(VersionedState);
598
599 impl $name {
600 pub fn encode<T>(
607 value: &T,
608 codec: &(impl VersionedStateCodec<T> + ?Sized),
609 limits: StateLimits,
610 ) -> Result<Self, StateError> {
611 VersionedState::encode($kind, value, codec, limits).map(Self)
612 }
613
614 pub fn from_json(bytes: &[u8], limits: StateLimits) -> Result<Self, StateError> {
621 VersionedState::from_json($kind, bytes, limits).map(Self)
622 }
623
624 pub fn decode<T>(
630 &self,
631 codec: &(impl VersionedStateCodec<T> + ?Sized),
632 ) -> Result<T, StateError> {
633 self.0.decode($kind, codec)
634 }
635
636 #[must_use]
638 pub const fn format_version(&self) -> u16 {
639 FORMAT_VERSION
640 }
641
642 #[must_use]
644 pub const fn schema_id(&self) -> &StateSchemaId {
645 &self.0.schema_id
646 }
647
648 #[must_use]
650 pub const fn schema_version(&self) -> StateSchemaVersion {
651 self.0.schema_version
652 }
653
654 #[must_use]
656 pub const fn encoded_len(&self) -> usize {
657 self.0.encoded_bytes
658 }
659
660 pub fn to_json(&self) -> Result<Vec<u8>, StateError> {
667 self.0.to_json($kind)
668 }
669
670 pub fn payload_json(&self) -> Result<Vec<u8>, StateError> {
681 self.0.payload_json($kind)
682 }
683 }
684
685 impl fmt::Debug for $name {
686 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
687 formatter
688 .debug_struct(stringify!($name))
689 .field("format_version", &FORMAT_VERSION)
690 .field("schema_version", &self.schema_version())
691 .field("encoded_bytes", &self.encoded_len())
692 .field("payload", &"<redacted>")
693 .finish()
694 }
695 }
696 };
697}
698
699durable_state!(
700 Checkpoint,
701 DurableStateKind::Checkpoint,
702 "A bounded, versioned reader position committed with a chunk."
703);
704durable_state!(
705 ExecutionContext,
706 DurableStateKind::ExecutionContext,
707 "Bounded, versioned application restart state committed with a chunk."
708);
709
710impl Checkpoint {
711 #[must_use]
716 pub fn generation_digest(&self) -> [u8; 32] {
717 self.to_json()
718 .map_or([0; 32], |bytes| Sha256::digest(&bytes).into())
719 }
720}
721
722#[derive(Clone, Debug, Eq, PartialEq)]
724#[non_exhaustive]
725pub enum StateError {
726 EmptySchemaId,
728 SchemaIdTooLong {
730 max_bytes: usize,
732 },
733 SchemaIdHasSurroundingWhitespace,
735 SchemaIdContainsControl,
737 ZeroSchemaVersion,
739 InvalidByteLimit {
741 maximum: usize,
743 },
744 InvalidDepthLimit {
746 maximum: usize,
748 },
749 TooLarge {
751 kind: DurableStateKind,
753 max_bytes: usize,
755 },
756 TooDeep {
758 kind: DurableStateKind,
760 max_depth: usize,
762 },
763 Malformed {
765 kind: DurableStateKind,
767 },
768 FormatMismatch {
770 kind: DurableStateKind,
772 },
773 UnsupportedFormatVersion {
775 kind: DurableStateKind,
777 version: u16,
779 },
780 SchemaMismatch {
782 kind: DurableStateKind,
784 },
785 NonIncreasingUpgrade {
787 from: StateSchemaVersion,
789 to: StateSchemaVersion,
791 },
792 NoUpgradePath {
794 kind: DurableStateKind,
796 found: StateSchemaVersion,
798 current: StateSchemaVersion,
800 },
801 AmbiguousUpgrade {
804 kind: DurableStateKind,
806 from: StateSchemaVersion,
808 },
809 UpgradeOvershootsCurrent {
811 kind: DurableStateKind,
813 to: StateSchemaVersion,
815 current: StateSchemaVersion,
817 },
818 UpgradeChainTooLong {
820 kind: DurableStateKind,
822 max_upgrades: usize,
824 },
825 UpgradeProducedInvalidJson {
827 kind: DurableStateKind,
829 },
830 UnsupportedSchemaVersion {
832 kind: DurableStateKind,
834 found: StateSchemaVersion,
836 current: StateSchemaVersion,
838 },
839 InvalidPayload,
841 PayloadNotObject,
843 Codec(StateCodecError),
845}
846
847impl fmt::Display for StateError {
848 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
849 match self {
850 Self::EmptySchemaId => formatter.write_str("state schema identifier must not be empty"),
851 Self::SchemaIdTooLong { max_bytes } => {
852 write!(
853 formatter,
854 "state schema identifier exceeds {max_bytes} UTF-8 bytes"
855 )
856 }
857 Self::SchemaIdHasSurroundingWhitespace => {
858 formatter.write_str("state schema identifier has surrounding whitespace")
859 }
860 Self::SchemaIdContainsControl => {
861 formatter.write_str("state schema identifier contains a control character")
862 }
863 Self::ZeroSchemaVersion => formatter.write_str("state schema version must be nonzero"),
864 Self::InvalidByteLimit { maximum } => {
865 write!(
866 formatter,
867 "state byte limit must be between 1 and {maximum}"
868 )
869 }
870 Self::InvalidDepthLimit { maximum } => {
871 write!(
872 formatter,
873 "state depth limit must be between 1 and {maximum}"
874 )
875 }
876 Self::TooLarge { kind, max_bytes } => {
877 write!(formatter, "{kind} exceeds {max_bytes} bytes")
878 }
879 Self::TooDeep { kind, max_depth } => {
880 write!(formatter, "{kind} exceeds JSON depth {max_depth}")
881 }
882 Self::Malformed { kind } => write!(formatter, "{kind} is malformed"),
883 Self::FormatMismatch { kind } => {
884 write!(formatter, "durable state is not a {kind}")
885 }
886 Self::UnsupportedFormatVersion { kind, .. } => {
887 write!(formatter, "{kind} format version is unsupported")
888 }
889 Self::SchemaMismatch { kind } => {
890 write!(formatter, "{kind} schema does not match the component")
891 }
892 Self::NonIncreasingUpgrade { .. } => {
893 formatter.write_str("state schema upgrade must increase the version")
894 }
895 Self::NoUpgradePath { kind, .. } => {
896 write!(formatter, "{kind} schema version has no upgrade path")
897 }
898 Self::AmbiguousUpgrade { kind, .. } => {
899 write!(formatter, "{kind} schema upgrade is ambiguous")
900 }
901 Self::UpgradeOvershootsCurrent { kind, .. } => {
902 write!(
903 formatter,
904 "{kind} schema upgrade passes the current version"
905 )
906 }
907 Self::UpgradeChainTooLong { kind, max_upgrades } => {
908 write!(
909 formatter,
910 "{kind} schema upgrade chain exceeds {max_upgrades} upgrades"
911 )
912 }
913 Self::UpgradeProducedInvalidJson { kind } => {
914 write!(formatter, "{kind} schema upgrade produced invalid JSON")
915 }
916 Self::UnsupportedSchemaVersion { kind, .. } => {
917 write!(formatter, "{kind} schema version is unsupported")
918 }
919 Self::InvalidPayload => formatter.write_str("durable state payload is not valid JSON"),
920 Self::PayloadNotObject => {
921 formatter.write_str("durable state payload must be a JSON object")
922 }
923 Self::Codec(error) => error.fmt(formatter),
924 }
925 }
926}
927
928impl Error for StateError {
929 fn source(&self) -> Option<&(dyn Error + 'static)> {
930 match self {
931 Self::Codec(error) => Some(error),
932 _ => None,
933 }
934 }
935}