templar_common/versioned_state/
core.rs

1use std::io::{Error, ErrorKind};
2
3use borsh::{BorshDeserialize, BorshSerialize};
4use near_sdk::{env, serde::de::DeserializeOwned, serde_json};
5
6const VERSION_KEY: &[u8] = b"__v";
7
8pub fn write_state_version(version: u32) {
9    env::storage_write(VERSION_KEY, &version.to_le_bytes());
10}
11
12pub fn read_state_version() -> Result<u32, std::io::Error> {
13    let Some(bytes) = env::storage_read(VERSION_KEY) else {
14        return Ok(0);
15    };
16
17    borsh::from_slice(&bytes)
18}
19
20#[derive(Debug)]
21#[near_sdk::near(serializers = [borsh])]
22pub struct VersionedState<T: StateVersion>(T);
23
24impl<T: StateVersion> VersionedState<T> {
25    pub fn new(state: T) -> Self {
26        write_state_version(T::VERSION);
27        Self(state)
28    }
29
30    pub fn version(&self) -> u32 {
31        T::VERSION
32    }
33}
34
35impl<T: StateVersion> std::ops::Deref for VersionedState<T> {
36    type Target = T;
37
38    fn deref(&self) -> &Self::Target {
39        &self.0
40    }
41}
42
43impl<T: StateVersion> std::ops::DerefMut for VersionedState<T> {
44    fn deref_mut(&mut self) -> &mut Self::Target {
45        &mut self.0
46    }
47}
48
49pub trait StateVersion {
50    const VERSION: u32;
51    type NewArgs;
52
53    fn new(args: Self::NewArgs) -> VersionedState<Self>
54    where
55        Self: Sized;
56
57    fn needs_migration() -> Result<bool, std::io::Error> {
58        let stored = read_state_version()?;
59        if stored > Self::VERSION {
60            return Err(Error::new(
61                ErrorKind::InvalidData,
62                format!(
63                    "Stored state version {stored} is newer than supported version {}",
64                    Self::VERSION
65                ),
66            ));
67        }
68
69        Ok(stored < Self::VERSION)
70    }
71}
72
73pub trait StateTransformer {
74    type Input: StateVersion + BorshDeserialize;
75    type Output: StateVersion + BorshSerialize;
76    type Error;
77
78    fn input_version(&self) -> u32 {
79        Self::Input::VERSION
80    }
81
82    fn output_version(&self) -> u32 {
83        Self::Output::VERSION
84    }
85
86    fn run(&self) -> Result<Self::Output, MigrationError<Self::Error>> {
87        let stored = read_state_version()?;
88        let expected = self.input_version();
89        if stored != expected {
90            return Err(MigrationError::StoredVersionMismatch { stored, expected });
91        }
92        let old_state =
93            env::state_read::<Self::Input>().ok_or(MigrationError::FailedToDeserializeOldState)?;
94        let new_state = self
95            .transform(old_state)
96            .map_err(MigrationError::Transformation)?;
97        env::state_write(&new_state);
98        write_state_version(self.output_version());
99        Ok(new_state)
100    }
101
102    fn transform(&self, input: Self::Input) -> Result<Self::Output, Self::Error>;
103}
104
105#[derive(thiserror::Error, Debug)]
106pub enum MigrationError<E> {
107    #[error("Failed to deserialize stored state version: {0}")]
108    StoredVersionDeserialization(#[from] std::io::Error),
109    #[error("Stored state version {stored} != args `from_version` {expected}")]
110    StoredVersionMismatch { stored: u32, expected: u32 },
111    #[error("Failed to deserialize old state")]
112    FailedToDeserializeOldState,
113    #[error("Failed to transform old state")]
114    Transformation(E),
115}
116
117pub trait Migrator {
118    /// Input state version this migration expects to read (the underlying
119    /// [`StateTransformer::input_version`]).
120    fn input_version(&self) -> u32;
121    /// Output state version this migration writes (the underlying
122    /// [`StateTransformer::output_version`]).
123    fn output_version(&self) -> u32;
124    /// Apply the migration, advancing the stored state version to
125    /// [`Migrator::output_version`]. Consumes `self`; run only after the chain is validated.
126    fn run(self);
127}
128
129/// Deserialize `migrate_args` as either a single migration (legacy wire shape, a JSON object) or an
130/// ordered list (a JSON array). A migration is JSON-tagged by `from_version`, so the first
131/// non-whitespace byte disambiguates: `[` is a list, anything else a single object wrapped into a
132/// one-element chain. Keeps the pre-list wire format working unchanged.
133pub fn parse_one_or_many<T: DeserializeOwned>(bytes: &[u8]) -> Result<Vec<T>, serde_json::Error> {
134    let is_array = bytes
135        .iter()
136        .find(|b| !b.is_ascii_whitespace())
137        .is_some_and(|b| *b == b'[');
138
139    if is_array {
140        serde_json::from_slice::<Vec<T>>(bytes)
141    } else {
142        serde_json::from_slice::<T>(bytes).map(|migration| vec![migration])
143    }
144}
145
146#[derive(thiserror::Error, Debug)]
147pub enum MigrationChainError {
148    #[error("Failed to read stored state version: {0}")]
149    StoredVersion(#[from] std::io::Error),
150    #[error("state migration is required but no migrations were provided")]
151    Empty,
152    #[error("first migration input version {input} != stored state version {stored}")]
153    StartMismatch { input: u32, stored: u32 },
154    #[error("migration output version {output} != next migration input version {input}")]
155    LinkMismatch { output: u32, input: u32 },
156    #[error("final migration output version {output} != target state version {target}")]
157    EndMismatch { output: u32, target: u32 },
158}
159
160/// Validate an ordered migration chain end-to-end, then run each step in sequence.
161///
162/// Every invariant is checked *before* any transform runs, so an invalid chain reverts without
163/// writing state:
164/// 1. the list is non-empty,
165/// 2. the first migration's input version equals the stored state version,
166/// 3. each migration's output version equals the next migration's input version, and
167/// 4. the last migration's output version equals `target`.
168///
169/// Only then are the migrations run in order; each [`Migrator::run`] independently re-asserts
170/// `stored == input_version()` and bumps the stored version to its output.
171pub fn run_migration_chain<M: Migrator>(
172    migrations: Vec<M>,
173    target: u32,
174) -> Result<(), MigrationChainError> {
175    let Some((first, rest)) = migrations.split_first() else {
176        return Err(MigrationChainError::Empty);
177    };
178
179    let stored = read_state_version()?;
180    if first.input_version() != stored {
181        return Err(MigrationChainError::StartMismatch {
182            input: first.input_version(),
183            stored,
184        });
185    }
186
187    let mut expected_input = first.output_version();
188    for migration in rest {
189        let input = migration.input_version();
190        if input != expected_input {
191            return Err(MigrationChainError::LinkMismatch {
192                output: expected_input,
193                input,
194            });
195        }
196        expected_input = migration.output_version();
197    }
198
199    // `expected_input` now holds the final migration's output version.
200    if expected_input != target {
201        return Err(MigrationChainError::EndMismatch {
202            output: expected_input,
203            target,
204        });
205    }
206
207    for migration in migrations {
208        migration.run();
209    }
210
211    Ok(())
212}
213
214#[cfg(test)]
215mod tests {
216    use near_sdk::{test_utils::VMContextBuilder, testing_env};
217    use rstest::rstest;
218
219    use super::*;
220
221    fn context() {
222        testing_env!(VMContextBuilder::new().build());
223    }
224
225    #[test]
226    fn stored_version_defaults_to_zero() {
227        context();
228        assert_eq!(read_state_version().unwrap(), 0);
229    }
230
231    #[test]
232    fn malformed_stored_version_errors() {
233        context();
234        write_state_version(7);
235        env::storage_write(VERSION_KEY, &[1, 2, 3]);
236
237        assert!(read_state_version().is_err());
238    }
239
240    #[test]
241    fn future_stored_version_errors() {
242        context();
243        write_state_version(9);
244
245        let error = TestState::needs_migration().unwrap_err();
246        assert_eq!(error.kind(), ErrorKind::InvalidData);
247        assert!(error
248            .to_string()
249            .contains("Stored state version 9 is newer"));
250    }
251
252    struct TestState;
253
254    impl StateVersion for TestState {
255        const VERSION: u32 = 2;
256        type NewArgs = ();
257
258        fn new((): Self::NewArgs) -> VersionedState<Self> {
259            VersionedState::new(Self)
260        }
261    }
262
263    // A `Migrator` standing in for a real transformer: it advances the stored version from `input`
264    // to `output`, re-asserting `stored == input` the way `StateTransformer::run` does.
265    struct MockMigration {
266        input: u32,
267        output: u32,
268    }
269
270    impl Migrator for MockMigration {
271        fn input_version(&self) -> u32 {
272            self.input
273        }
274
275        fn output_version(&self) -> u32 {
276            self.output
277        }
278
279        fn run(self) {
280            assert_eq!(
281                read_state_version().unwrap(),
282                self.input,
283                "mock migration ran against the wrong stored version",
284            );
285            write_state_version(self.output);
286        }
287    }
288
289    fn m(input: u32, output: u32) -> MockMigration {
290        MockMigration { input, output }
291    }
292
293    #[derive(Debug, PartialEq)]
294    #[near_sdk::near(serializers = [json])]
295    #[serde(tag = "from_version", rename_all = "snake_case")]
296    enum TestMigration {
297        V0,
298        V1,
299    }
300
301    #[test]
302    fn parse_single_object_is_one_element_chain() {
303        let parsed: Vec<TestMigration> = parse_one_or_many(br#"{"from_version":"v0"}"#).unwrap();
304        assert_eq!(parsed, vec![TestMigration::V0]);
305    }
306
307    #[test]
308    fn parse_array_is_multi_element_chain() {
309        let parsed: Vec<TestMigration> =
310            parse_one_or_many(br#"[{"from_version":"v0"},{"from_version":"v1"}]"#).unwrap();
311        assert_eq!(parsed, vec![TestMigration::V0, TestMigration::V1]);
312    }
313
314    #[test]
315    fn parse_tolerates_leading_whitespace_before_array() {
316        let parsed: Vec<TestMigration> =
317            parse_one_or_many(b"  \n [{\"from_version\":\"v0\"}]").unwrap();
318        assert_eq!(parsed, vec![TestMigration::V0]);
319    }
320
321    #[test]
322    fn chain_runs_multiple_steps_to_target() {
323        context();
324        write_state_version(0);
325
326        run_migration_chain(vec![m(0, 1), m(1, 2)], 2).unwrap();
327
328        assert_eq!(read_state_version().unwrap(), 2);
329    }
330
331    #[test]
332    fn chain_runs_single_step_to_target() {
333        context();
334        write_state_version(1);
335
336        run_migration_chain(vec![m(1, 2)], 2).unwrap();
337
338        assert_eq!(read_state_version().unwrap(), 2);
339    }
340
341    // Every invalid chain is rejected before any transform runs, so the stored version is left
342    // untouched. `matches` pins the exact error variant (and its versions) per case.
343    #[rstest]
344    #[case::empty(Vec::new(), 2, |e: &_| matches!(e, MigrationChainError::Empty))]
345    #[case::wrong_start(vec![m(1, 2)], 2, |e: &_| matches!(e, MigrationChainError::StartMismatch { input: 1, stored: 0 }))]
346    #[case::broken_link(vec![m(0, 1), m(2, 3)], 3, |e: &_| matches!(e, MigrationChainError::LinkMismatch { output: 1, input: 2 }))]
347    #[case::not_landing_on_target(vec![m(0, 1)], 2, |e: &_| matches!(e, MigrationChainError::EndMismatch { output: 1, target: 2 }))]
348    fn chain_rejects_invalid(
349        #[case] migrations: Vec<MockMigration>,
350        #[case] target: u32,
351        #[case] matches: fn(&MigrationChainError) -> bool,
352    ) {
353        context();
354        write_state_version(0);
355
356        let err = run_migration_chain(migrations, target).unwrap_err();
357        assert!(matches(&err), "unexpected error variant: {err:?}");
358        assert_eq!(
359            read_state_version().unwrap(),
360            0,
361            "no state written on invalid chain"
362        );
363    }
364}
templar_common/versioned_state/core.rs

templar_common/versioned_state/
core.rs
