oasis_runtime_sdk/modules/rofl/
state.rs

1use crate::{
2    core::{
3        common::crypto::{hash::Hash, signature::PublicKey as CorePublicKey},
4        consensus::beacon::EpochTime,
5    },
6    crypto::signature::PublicKey,
7    state::CurrentState,
8    storage::{self, Store},
9};
10
11use super::{app_id::AppId, types, Error, MODULE_NAME};
12
13/// Map of application identifiers to their configs.
14const APPS: &[u8] = &[0x01];
15/// Map of (application identifier, H(RAK)) tuples to their registrations.
16const REGISTRATIONS: &[u8] = &[0x02];
17/// Map of H(pk)s to KeyEndorsementInfos. This is used when just the public key is needed to avoid
18/// fetching entire registrations from storage.
19const ENDORSERS: &[u8] = &[0x03];
20/// A queue of registration expirations.
21const EXPIRATION_QUEUE: &[u8] = &[0x04];
22
23/// Information about an endorsed key.
24#[derive(Clone, Debug, Default, PartialEq, Eq, cbor::Encode, cbor::Decode)]
25#[cbor(as_array)]
26pub struct KeyEndorsementInfo {
27    /// Identifier of node that endorsed the enclave.
28    pub node_id: CorePublicKey,
29    /// RAK of the enclave that endorsed the key. This is only set for endorsements of extra keys.
30    pub rak: Option<CorePublicKey>,
31    /// Identifier of the app this key is for.
32    pub app_id: AppId,
33}
34
35/// Information about an endorsed key.
36#[derive(Clone, Debug, Default, PartialEq, Eq, cbor::Encode, cbor::Decode)]
37#[cbor(as_array)]
38struct LegacyKeyEndorsementInfo {
39    /// Identifier of node that endorsed the enclave.
40    pub node_id: CorePublicKey,
41    /// RAK of the enclave that endorsed the key. This is only set for endorsements of extra keys.
42    pub rak: Option<CorePublicKey>,
43}
44
45impl KeyEndorsementInfo {
46    /// Create a new key endorsement information for RAK endorsed by given node directly.
47    pub fn for_rak(node_id: CorePublicKey, app_id: AppId) -> Self {
48        Self {
49            node_id,
50            app_id,
51            ..Default::default()
52        }
53    }
54
55    /// Create a new key endorsement information for extra key endorsed by RAK.
56    pub fn for_extra_key(node_id: CorePublicKey, app_id: AppId, rak: CorePublicKey) -> Self {
57        Self {
58            node_id,
59            rak: Some(rak),
60            app_id,
61        }
62    }
63}
64
65impl From<LegacyKeyEndorsementInfo> for KeyEndorsementInfo {
66    fn from(value: LegacyKeyEndorsementInfo) -> Self {
67        Self {
68            node_id: value.node_id,
69            rak: value.rak,
70            // Using a default app ID is fine because this app ID cannot be generated by any scheme.
71            app_id: AppId::default(),
72        }
73    }
74}
75
76fn apps<S: Store>(store: S) -> storage::TypedStore<impl Store> {
77    let store = storage::PrefixStore::new(store, &MODULE_NAME);
78    storage::TypedStore::new(storage::PrefixStore::new(store, &APPS))
79}
80
81/// Retrieves an application configuration.
82pub fn get_app(app_id: AppId) -> Option<types::AppConfig> {
83    CurrentState::with_store(|store| apps(store).get(app_id))
84}
85
86/// Retrieves all application configurations.
87pub fn get_apps() -> Vec<types::AppConfig> {
88    CurrentState::with_store(|store| {
89        let store = storage::PrefixStore::new(store, &MODULE_NAME);
90        let apps = storage::TypedStore::new(storage::PrefixStore::new(store, &APPS));
91        apps.iter()
92            .map(|(_, cfg): (AppId, types::AppConfig)| cfg)
93            .collect()
94    })
95}
96
97/// Updates an application configuration.
98pub fn set_app(cfg: types::AppConfig) {
99    CurrentState::with_store(|store| {
100        apps(store).insert(cfg.id, cfg);
101    })
102}
103
104/// Removes an application configuration.
105pub fn remove_app(app_id: AppId) {
106    CurrentState::with_store(|store| {
107        apps(store).remove(app_id);
108    })
109}
110
111/// Updates registration of the given ROFL enclave.
112pub fn update_registration(registration: types::Registration) -> Result<(), Error> {
113    let hrak = hash_rak(&registration.rak);
114
115    // Update expiration queue.
116    if let Some(existing) = get_registration_hrak(registration.app, hrak) {
117        // Disallow modification of extra keys.
118        if existing.extra_keys != registration.extra_keys {
119            return Err(Error::ExtraKeyUpdateNotAllowed);
120        }
121
122        remove_expiration_queue(existing.expiration, registration.app, hrak);
123    }
124    insert_expiration_queue(registration.expiration, registration.app, hrak);
125
126    // Update registration.
127    CurrentState::with_store(|mut root_store| {
128        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
129        let mut endorsers = storage::TypedStore::new(storage::PrefixStore::new(store, &ENDORSERS));
130        endorsers.insert(
131            hrak,
132            KeyEndorsementInfo::for_rak(registration.node_id, registration.app),
133        );
134
135        for pk in &registration.extra_keys {
136            endorsers.insert(
137                hash_pk(pk),
138                KeyEndorsementInfo::for_extra_key(
139                    registration.node_id,
140                    registration.app,
141                    registration.rak,
142                ),
143            );
144        }
145
146        let app_id = registration.app;
147        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
148        let registrations = storage::PrefixStore::new(store, &REGISTRATIONS);
149        let mut app = storage::TypedStore::new(storage::PrefixStore::new(registrations, app_id));
150        app.insert(hrak, registration);
151    });
152
153    Ok(())
154}
155
156fn remove_registration_hrak(app_id: AppId, hrak: Hash) {
157    let registration = match get_registration_hrak(app_id, hrak) {
158        Some(registration) => registration,
159        None => return,
160    };
161
162    // Remove from expiration queue if present.
163    remove_expiration_queue(registration.expiration, registration.app, hrak);
164
165    // Remove registration.
166    CurrentState::with_store(|mut root_store| {
167        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
168        let mut endorsers = storage::TypedStore::new(storage::PrefixStore::new(store, &ENDORSERS));
169        endorsers.remove(hrak);
170
171        for pk in &registration.extra_keys {
172            endorsers.remove(hash_pk(pk));
173        }
174
175        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
176        let registrations = storage::PrefixStore::new(store, &REGISTRATIONS);
177        let mut app = storage::TypedStore::new(storage::PrefixStore::new(registrations, app_id));
178        app.remove(hrak);
179    });
180}
181
182/// Removes an existing registration of the given ROFL enclave.
183pub fn remove_registration(app_id: AppId, rak: &CorePublicKey) {
184    remove_registration_hrak(app_id, hash_rak(rak))
185}
186
187fn get_registration_hrak(app_id: AppId, hrak: Hash) -> Option<types::Registration> {
188    CurrentState::with_store(|store| {
189        let store = storage::PrefixStore::new(store, &MODULE_NAME);
190        let registrations = storage::PrefixStore::new(store, &REGISTRATIONS);
191        let app = storage::TypedStore::new(storage::PrefixStore::new(registrations, app_id));
192        app.get(hrak)
193    })
194}
195
196/// Retrieves registration of the given ROFL enclave. In case enclave is not registered, returns
197/// `None`.
198pub fn get_registration(app_id: AppId, rak: &CorePublicKey) -> Option<types::Registration> {
199    get_registration_hrak(app_id, hash_rak(rak))
200}
201
202/// Retrieves all registrations for the given ROFL application.
203pub fn get_registrations_for_app(app_id: AppId) -> Vec<types::Registration> {
204    CurrentState::with_store(|mut root_store| {
205        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
206        let registrations = storage::PrefixStore::new(store, &REGISTRATIONS);
207        let app = storage::TypedStore::new(storage::PrefixStore::new(registrations, app_id));
208
209        app.iter()
210            .map(|(_, registration): (Hash, types::Registration)| registration)
211            .collect()
212    })
213}
214
215/// Retrieves endorser of the given ROFL enclave. In case enclave is not registered, returns `None`.
216pub fn get_endorser(pk: &PublicKey) -> Option<KeyEndorsementInfo> {
217    let hpk = hash_pk(pk);
218
219    CurrentState::with_store(|store| {
220        let store = storage::PrefixStore::new(store, &MODULE_NAME);
221        let endorsers = storage::TypedStore::new(storage::PrefixStore::new(store, &ENDORSERS));
222        let kei: Option<cbor::Value> = endorsers.get(hpk);
223        kei.map(|kei| match kei {
224            cbor::Value::Array(ref items) if items.len() == 2 => {
225                // Legacy variant without app ID.
226                let legacy: LegacyKeyEndorsementInfo = cbor::from_value(kei).unwrap();
227                legacy.into()
228            }
229            _ => cbor::from_value(kei).unwrap(),
230        })
231    })
232}
233
234fn hash_rak(rak: &CorePublicKey) -> Hash {
235    hash_pk(&PublicKey::Ed25519(rak.into()))
236}
237
238fn hash_pk(pk: &PublicKey) -> Hash {
239    Hash::digest_bytes_list(&[pk.key_type().as_bytes(), pk.as_ref()])
240}
241
242fn queue_entry_key(epoch: EpochTime, app_id: AppId, hrak: Hash) -> Vec<u8> {
243    [&epoch.to_be_bytes(), app_id.as_ref(), hrak.as_ref()].concat()
244}
245
246fn insert_expiration_queue(epoch: EpochTime, app_id: AppId, hrak: Hash) {
247    CurrentState::with_store(|store| {
248        let store = storage::PrefixStore::new(store, &MODULE_NAME);
249        let mut queue = storage::PrefixStore::new(store, &EXPIRATION_QUEUE);
250        queue.insert(&queue_entry_key(epoch, app_id, hrak), &[]);
251    })
252}
253
254fn remove_expiration_queue(epoch: EpochTime, app_id: AppId, hrak: Hash) {
255    CurrentState::with_store(|store| {
256        let store = storage::PrefixStore::new(store, &MODULE_NAME);
257        let mut queue = storage::PrefixStore::new(store, &EXPIRATION_QUEUE);
258        queue.remove(&queue_entry_key(epoch, app_id, hrak));
259    })
260}
261
262struct ExpirationQueueEntry {
263    epoch: EpochTime,
264    app_id: AppId,
265    hrak: Hash,
266}
267
268impl<'a> TryFrom<&'a [u8]> for ExpirationQueueEntry {
269    type Error = anyhow::Error;
270
271    fn try_from(value: &'a [u8]) -> Result<Self, Self::Error> {
272        // Decode a storage key of the format (epoch, hrak).
273        if value.len() != 8 + AppId::SIZE + Hash::len() {
274            anyhow::bail!("incorrect expiration queue key size");
275        }
276
277        Ok(Self {
278            epoch: EpochTime::from_be_bytes(value[..8].try_into()?),
279            app_id: value[8..8 + AppId::SIZE].try_into()?,
280            hrak: value[8 + AppId::SIZE..].into(),
281        })
282    }
283}
284
285/// Removes all expired registrations, e.g. those that expire in epochs earlier than or equal to the
286/// passed epoch.
287pub fn expire_registrations(epoch: EpochTime, limit: usize) {
288    let expired: Vec<_> = CurrentState::with_store(|store| {
289        let store = storage::PrefixStore::new(store, &MODULE_NAME);
290        let queue = storage::TypedStore::new(storage::PrefixStore::new(store, &EXPIRATION_QUEUE));
291
292        queue
293            .iter()
294            .take_while(|(e, _): &(ExpirationQueueEntry, CorePublicKey)| e.epoch <= epoch)
295            .map(|(e, _)| (e.app_id, e.hrak))
296            .take(limit)
297            .collect()
298    });
299
300    for (app_id, hrak) in expired {
301        remove_registration_hrak(app_id, hrak);
302    }
303}
304
305#[cfg(test)]
306mod test {
307    use super::*;
308    use crate::testing::{keys, mock};
309
310    #[test]
311    fn test_app_cfg() {
312        let _mock = mock::Mock::default();
313
314        let app_id = AppId::from_creator_round_index(keys::alice::address(), 0, 0);
315        let app = get_app(app_id);
316        assert!(app.is_none());
317
318        let cfg = types::AppConfig {
319            id: app_id,
320            admin: Some(keys::alice::address()),
321            ..Default::default()
322        };
323        set_app(cfg.clone());
324        let app = get_app(app_id).expect("application config should be created");
325        assert_eq!(app, cfg);
326
327        let cfg = types::AppConfig { admin: None, ..cfg };
328        set_app(cfg.clone());
329        let app = get_app(app_id).expect("application config should be updated");
330        assert_eq!(app, cfg);
331
332        let apps = get_apps();
333        assert_eq!(apps.len(), 1);
334        assert_eq!(apps[0], cfg);
335
336        remove_app(app_id);
337        let app = get_app(app_id);
338        assert!(app.is_none(), "application should have been removed");
339        let apps = get_apps();
340        assert_eq!(apps.len(), 0);
341    }
342
343    #[test]
344    fn test_registration() {
345        let _mock = mock::Mock::default();
346        let app_id = Default::default();
347        let rak = keys::alice::pk().try_into().unwrap(); // Fake RAK.
348        let rak_pk = keys::alice::pk();
349
350        let registration = get_registration(app_id, &rak);
351        assert!(registration.is_none());
352        let endorser = get_endorser(&rak_pk);
353        assert!(endorser.is_none());
354        let endorser = get_endorser(&keys::dave::pk());
355        assert!(endorser.is_none());
356
357        let new_registration = types::Registration {
358            app: app_id,
359            rak,
360            expiration: 42,
361            extra_keys: vec![
362                keys::dave::pk(), // Add dave as an extra endorsed key.
363            ],
364            ..Default::default()
365        };
366        update_registration(new_registration.clone()).expect("registration update should work");
367
368        // Ensure extra endorsed keys cannot be updated later.
369        let bad_registration = types::Registration {
370            app: app_id,
371            extra_keys: vec![],
372            ..new_registration.clone()
373        };
374        update_registration(bad_registration.clone())
375            .expect_err("extra endorsed key update should not be allowed");
376
377        let registration = get_registration(app_id, &rak).expect("registration should be present");
378        assert_eq!(registration, new_registration);
379        let endorser = get_endorser(&rak_pk).expect("endorser should be present");
380        assert_eq!(endorser.node_id, new_registration.node_id);
381        assert!(endorser.rak.is_none());
382        let endorser = get_endorser(&keys::dave::pk()).expect("extra keys should be endorsed");
383        assert_eq!(endorser.node_id, new_registration.node_id);
384        assert_eq!(endorser.rak, Some(rak));
385        let registrations = get_registrations_for_app(new_registration.app);
386        assert_eq!(registrations.len(), 1);
387
388        expire_registrations(42, 128);
389
390        let registration = get_registration(app_id, &rak);
391        assert!(registration.is_none());
392        let endorser = get_endorser(&rak_pk);
393        assert!(endorser.is_none());
394        let endorser = get_endorser(&keys::dave::pk());
395        assert!(endorser.is_none());
396        let registrations = get_registrations_for_app(new_registration.app);
397        assert_eq!(registrations.len(), 0);
398    }
399
400    #[test]
401    fn test_legacy_kei() {
402        let _mock = mock::Mock::default();
403        let rak_pk = keys::alice::pk();
404        let hpk = hash_pk(&rak_pk);
405
406        CurrentState::with_store(|store| {
407            let store = storage::PrefixStore::new(store, &MODULE_NAME);
408            let mut endorsers =
409                storage::TypedStore::new(storage::PrefixStore::new(store, &ENDORSERS));
410            endorsers.insert(
411                hpk,
412                LegacyKeyEndorsementInfo {
413                    node_id: keys::bob::pk_ed25519().into(),
414                    rak: None,
415                },
416            );
417        });
418
419        let kei = get_endorser(&rak_pk).unwrap();
420        assert_eq!(kei.node_id, keys::bob::pk_ed25519().into());
421        assert!(kei.rak.is_none());
422        assert_eq!(kei.app_id, AppId::default());
423    }
424}
oasis_runtime_sdk/modules/rofl/state.rs

oasis_runtime_sdk/modules/rofl/
state.rs
