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}
32
33impl KeyEndorsementInfo {
34    /// Create a new key endorsement information for RAK endorsed by given node directly.
35    pub fn for_rak(node_id: CorePublicKey) -> Self {
36        Self {
37            node_id,
38            ..Default::default()
39        }
40    }
41
42    /// Create a new key endorsement information for extra key endorsed by RAK.
43    pub fn for_extra_key(node_id: CorePublicKey, rak: CorePublicKey) -> Self {
44        Self {
45            node_id,
46            rak: Some(rak),
47        }
48    }
49}
50
51fn apps<S: Store>(store: S) -> storage::TypedStore<impl Store> {
52    let store = storage::PrefixStore::new(store, &MODULE_NAME);
53    storage::TypedStore::new(storage::PrefixStore::new(store, &APPS))
54}
55
56/// Retrieves an application configuration.
57pub fn get_app(app_id: AppId) -> Option<types::AppConfig> {
58    CurrentState::with_store(|store| apps(store).get(app_id))
59}
60
61/// Retrieves all application configurations.
62pub fn get_apps() -> Vec<types::AppConfig> {
63    CurrentState::with_store(|store| {
64        let store = storage::PrefixStore::new(store, &MODULE_NAME);
65        let apps = storage::TypedStore::new(storage::PrefixStore::new(store, &APPS));
66        apps.iter()
67            .map(|(_, cfg): (AppId, types::AppConfig)| cfg)
68            .collect()
69    })
70}
71
72/// Updates an application configuration.
73pub fn set_app(cfg: types::AppConfig) {
74    CurrentState::with_store(|store| {
75        apps(store).insert(cfg.id, cfg);
76    })
77}
78
79/// Removes an application configuration.
80pub fn remove_app(app_id: AppId) {
81    CurrentState::with_store(|store| {
82        apps(store).remove(app_id);
83    })
84}
85
86/// Updates registration of the given ROFL enclave.
87pub fn update_registration(registration: types::Registration) -> Result<(), Error> {
88    let hrak = hash_rak(&registration.rak);
89
90    // Update expiration queue.
91    if let Some(existing) = get_registration_hrak(registration.app, hrak) {
92        // Disallow modification of extra keys.
93        if existing.extra_keys != registration.extra_keys {
94            return Err(Error::ExtraKeyUpdateNotAllowed);
95        }
96
97        remove_expiration_queue(existing.expiration, registration.app, hrak);
98    }
99    insert_expiration_queue(registration.expiration, registration.app, hrak);
100
101    // Update registration.
102    CurrentState::with_store(|mut root_store| {
103        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
104        let mut endorsers = storage::TypedStore::new(storage::PrefixStore::new(store, &ENDORSERS));
105        endorsers.insert(hrak, KeyEndorsementInfo::for_rak(registration.node_id));
106
107        for pk in &registration.extra_keys {
108            endorsers.insert(
109                hash_pk(pk),
110                KeyEndorsementInfo::for_extra_key(registration.node_id, registration.rak),
111            );
112        }
113
114        let app_id = registration.app;
115        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
116        let registrations = storage::PrefixStore::new(store, &REGISTRATIONS);
117        let mut app = storage::TypedStore::new(storage::PrefixStore::new(registrations, app_id));
118        app.insert(hrak, registration);
119    });
120
121    Ok(())
122}
123
124fn remove_registration_hrak(app_id: AppId, hrak: Hash) {
125    let registration = match get_registration_hrak(app_id, hrak) {
126        Some(registration) => registration,
127        None => return,
128    };
129
130    // Remove from expiration queue if present.
131    remove_expiration_queue(registration.expiration, registration.app, hrak);
132
133    // Remove registration.
134    CurrentState::with_store(|mut root_store| {
135        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
136        let mut endorsers = storage::TypedStore::new(storage::PrefixStore::new(store, &ENDORSERS));
137        endorsers.remove(hrak);
138
139        for pk in &registration.extra_keys {
140            endorsers.remove(hash_pk(pk));
141        }
142
143        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
144        let registrations = storage::PrefixStore::new(store, &REGISTRATIONS);
145        let mut app = storage::TypedStore::new(storage::PrefixStore::new(registrations, app_id));
146        app.remove(hrak);
147    });
148}
149
150/// Removes an existing registration of the given ROFL enclave.
151pub fn remove_registration(app_id: AppId, rak: &CorePublicKey) {
152    remove_registration_hrak(app_id, hash_rak(rak))
153}
154
155fn get_registration_hrak(app_id: AppId, hrak: Hash) -> Option<types::Registration> {
156    CurrentState::with_store(|store| {
157        let store = storage::PrefixStore::new(store, &MODULE_NAME);
158        let registrations = storage::PrefixStore::new(store, &REGISTRATIONS);
159        let app = storage::TypedStore::new(storage::PrefixStore::new(registrations, app_id));
160        app.get(hrak)
161    })
162}
163
164/// Retrieves registration of the given ROFL enclave. In case enclave is not registered, returns
165/// `None`.
166pub fn get_registration(app_id: AppId, rak: &CorePublicKey) -> Option<types::Registration> {
167    get_registration_hrak(app_id, hash_rak(rak))
168}
169
170/// Retrieves all registrations for the given ROFL application.
171pub fn get_registrations_for_app(app_id: AppId) -> Vec<types::Registration> {
172    CurrentState::with_store(|mut root_store| {
173        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
174        let registrations = storage::PrefixStore::new(store, &REGISTRATIONS);
175        let app = storage::TypedStore::new(storage::PrefixStore::new(registrations, app_id));
176
177        app.iter()
178            .map(|(_, registration): (Hash, types::Registration)| registration)
179            .collect()
180    })
181}
182
183/// Retrieves endorser of the given ROFL enclave. In case enclave is not registered, returns `None`.
184pub fn get_endorser(pk: &PublicKey) -> Option<KeyEndorsementInfo> {
185    let hpk = hash_pk(pk);
186
187    CurrentState::with_store(|store| {
188        let store = storage::PrefixStore::new(store, &MODULE_NAME);
189        let endorsers = storage::TypedStore::new(storage::PrefixStore::new(store, &ENDORSERS));
190        endorsers.get(hpk)
191    })
192}
193
194fn hash_rak(rak: &CorePublicKey) -> Hash {
195    hash_pk(&PublicKey::Ed25519(rak.into()))
196}
197
198fn hash_pk(pk: &PublicKey) -> Hash {
199    Hash::digest_bytes_list(&[pk.key_type().as_bytes(), pk.as_ref()])
200}
201
202fn queue_entry_key(epoch: EpochTime, app_id: AppId, hrak: Hash) -> Vec<u8> {
203    [&epoch.to_be_bytes(), app_id.as_ref(), hrak.as_ref()].concat()
204}
205
206fn insert_expiration_queue(epoch: EpochTime, app_id: AppId, hrak: Hash) {
207    CurrentState::with_store(|store| {
208        let store = storage::PrefixStore::new(store, &MODULE_NAME);
209        let mut queue = storage::PrefixStore::new(store, &EXPIRATION_QUEUE);
210        queue.insert(&queue_entry_key(epoch, app_id, hrak), &[]);
211    })
212}
213
214fn remove_expiration_queue(epoch: EpochTime, app_id: AppId, hrak: Hash) {
215    CurrentState::with_store(|store| {
216        let store = storage::PrefixStore::new(store, &MODULE_NAME);
217        let mut queue = storage::PrefixStore::new(store, &EXPIRATION_QUEUE);
218        queue.remove(&queue_entry_key(epoch, app_id, hrak));
219    })
220}
221
222struct ExpirationQueueEntry {
223    epoch: EpochTime,
224    app_id: AppId,
225    hrak: Hash,
226}
227
228impl<'a> TryFrom<&'a [u8]> for ExpirationQueueEntry {
229    type Error = anyhow::Error;
230
231    fn try_from(value: &'a [u8]) -> Result<Self, Self::Error> {
232        // Decode a storage key of the format (epoch, hrak).
233        if value.len() != 8 + AppId::SIZE + Hash::len() {
234            anyhow::bail!("incorrect expiration queue key size");
235        }
236
237        Ok(Self {
238            epoch: EpochTime::from_be_bytes(value[..8].try_into()?),
239            app_id: value[8..8 + AppId::SIZE].try_into()?,
240            hrak: value[8 + AppId::SIZE..].into(),
241        })
242    }
243}
244
245/// Removes all expired registrations, e.g. those that expire in epochs earlier than or equal to the
246/// passed epoch.
247pub fn expire_registrations(epoch: EpochTime, limit: usize) {
248    let expired: Vec<_> = CurrentState::with_store(|store| {
249        let store = storage::PrefixStore::new(store, &MODULE_NAME);
250        let queue = storage::TypedStore::new(storage::PrefixStore::new(store, &EXPIRATION_QUEUE));
251
252        queue
253            .iter()
254            .take_while(|(e, _): &(ExpirationQueueEntry, CorePublicKey)| e.epoch <= epoch)
255            .map(|(e, _)| (e.app_id, e.hrak))
256            .take(limit)
257            .collect()
258    });
259
260    for (app_id, hrak) in expired {
261        remove_registration_hrak(app_id, hrak);
262    }
263}
264
265#[cfg(test)]
266mod test {
267    use super::*;
268    use crate::testing::{keys, mock};
269
270    #[test]
271    fn test_app_cfg() {
272        let _mock = mock::Mock::default();
273
274        let app_id = AppId::from_creator_round_index(keys::alice::address(), 0, 0);
275        let app = get_app(app_id);
276        assert!(app.is_none());
277
278        let cfg = types::AppConfig {
279            id: app_id,
280            admin: Some(keys::alice::address()),
281            ..Default::default()
282        };
283        set_app(cfg.clone());
284        let app = get_app(app_id).expect("application config should be created");
285        assert_eq!(app, cfg);
286
287        let cfg = types::AppConfig { admin: None, ..cfg };
288        set_app(cfg.clone());
289        let app = get_app(app_id).expect("application config should be updated");
290        assert_eq!(app, cfg);
291
292        let apps = get_apps();
293        assert_eq!(apps.len(), 1);
294        assert_eq!(apps[0], cfg);
295
296        remove_app(app_id);
297        let app = get_app(app_id);
298        assert!(app.is_none(), "application should have been removed");
299        let apps = get_apps();
300        assert_eq!(apps.len(), 0);
301    }
302
303    #[test]
304    fn test_registration() {
305        let _mock = mock::Mock::default();
306        let app_id = Default::default();
307        let rak = keys::alice::pk().try_into().unwrap(); // Fake RAK.
308        let rak_pk = keys::alice::pk();
309
310        let registration = get_registration(app_id, &rak);
311        assert!(registration.is_none());
312        let endorser = get_endorser(&rak_pk);
313        assert!(endorser.is_none());
314        let endorser = get_endorser(&keys::dave::pk());
315        assert!(endorser.is_none());
316
317        let new_registration = types::Registration {
318            app: app_id,
319            rak,
320            expiration: 42,
321            extra_keys: vec![
322                keys::dave::pk(), // Add dave as an extra endorsed key.
323            ],
324            ..Default::default()
325        };
326        update_registration(new_registration.clone()).expect("registration update should work");
327
328        // Ensure extra endorsed keys cannot be updated later.
329        let bad_registration = types::Registration {
330            app: app_id,
331            extra_keys: vec![],
332            ..new_registration.clone()
333        };
334        update_registration(bad_registration.clone())
335            .expect_err("extra endorsed key update should not be allowed");
336
337        let registration = get_registration(app_id, &rak).expect("registration should be present");
338        assert_eq!(registration, new_registration);
339        let endorser = get_endorser(&rak_pk).expect("endorser should be present");
340        assert_eq!(endorser.node_id, new_registration.node_id);
341        assert!(endorser.rak.is_none());
342        let endorser = get_endorser(&keys::dave::pk()).expect("extra keys should be endorsed");
343        assert_eq!(endorser.node_id, new_registration.node_id);
344        assert_eq!(endorser.rak, Some(rak));
345        let registrations = get_registrations_for_app(new_registration.app);
346        assert_eq!(registrations.len(), 1);
347
348        expire_registrations(42, 128);
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        let registrations = get_registrations_for_app(new_registration.app);
357        assert_eq!(registrations.len(), 0);
358    }
359}