Refactor asset state

src/internet_identity/src/assets.rs

@@ -6,9 +6,33 @@ use crate::{http, state};
 use base64::engine::general_purpose::STANDARD as BASE64;
 use base64::Engine;
 use ic_cdk::api;
+use ic_certified_map::{AsHashTree, Hash, HashTree, RbTree};
 use include_dir::{include_dir, Dir, File};
+use internet_identity_interface::http_gateway::HeaderField;
 use lazy_static::lazy_static;
 use sha2::Digest;
+use std::collections::HashMap;
+
+const LABEL_ASSETS: &[u8] = b"http_assets";
+
+#[derive(Debug, Default, Clone)]
+pub struct CertifiedAssets {
+    pub assets: HashMap<String, (Vec<HeaderField>, Vec<u8>)>,
+    pub certification_v1: RbTree<String, Hash>,
+}
+
+impl CertifiedAssets {
+    /// Returns the root_hash of the asset certification tree.
+    /// In the future, this will include both certification v1 and v2.
+    pub fn root_hash(&self) -> Hash {
+        ic_certified_map::labeled_hash(LABEL_ASSETS, &self.certification_v1.root_hash())
+    }
+
+    pub fn witness(&self, path: &str) -> HashTree {
+        let witness = self.certification_v1.witness(path.as_bytes());
+        ic_certified_map::labeled(LABEL_ASSETS, witness)
+    }
+}
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub enum ContentEncoding {
@@ -63,9 +87,11 @@ lazy_static! {
 
 // used both in init and post_upgrade
 pub fn init_assets() {
-    state::assets_and_hashes_mut(|assets, asset_hashes| {
+    state::assets_mut(|certified_assets| {
         for (path, content, content_encoding, content_type) in get_static_assets() {
-            asset_hashes.insert(path.clone(), sha2::Sha256::digest(&content).into());
+            certified_assets
+                .certification_v1
+                .insert(path.clone(), sha2::Sha256::digest(&content).into());
             let mut headers = match content_encoding {
                 ContentEncoding::Identity => vec![],
                 ContentEncoding::GZip => {
@@ -76,7 +102,7 @@ pub fn init_assets() {
                 "Content-Type".to_string(),
                 content_type.to_mime_type_string(),
             ));
-            assets.insert(path, (headers, content));
+            certified_assets.assets.insert(path, (headers, content));
         }
     });
 }

src/internet_identity/src/delegation.rs

@@ -1,10 +1,10 @@
 use crate::active_anchor_stats::IIDomain;
-use crate::state::{persistent_state_mut, AssetHashes};
-use crate::{hash, state, update_root_hash, DAY_NS, LABEL_ASSETS, LABEL_SIG, MINUTE_NS};
+use crate::assets::CertifiedAssets;
+use crate::state::persistent_state_mut;
+use crate::{hash, state, update_root_hash, DAY_NS, LABEL_SIG, MINUTE_NS};
 use candid::Principal;
 use ic_cdk::api::{data_certificate, time};
 use ic_cdk::{id, trap};
-use ic_certified_map::AsHashTree;
 use ic_certified_map::{Hash, HashTree};
 use internet_identity::signature_map::SignatureMap;
 use internet_identity_interface::internet_identity::types::*;
@@ -127,7 +127,7 @@ pub fn get_delegation(
 ) -> GetDelegationResponse {
     check_frontend_length(&frontend);
 
-    state::asset_hashes_and_sigs(|asset_hashes, sigs| {
+    state::assets_and_signatures(|asset_hashes, sigs| {
         match get_signature(
             asset_hashes,
             sigs,
@@ -217,7 +217,7 @@ fn delegation_signature_msg_hash(d: &Delegation) -> Hash {
 }
 
 fn get_signature(
-    asset_hashes: &AssetHashes,
+    assets: &CertifiedAssets,
     sigs: &SignatureMap,
     pk: PublicKey,
     seed: Hash,
@@ -244,10 +244,7 @@ fn get_signature(
     }
 
     let tree = ic_certified_map::fork(
-        HashTree::Pruned(ic_certified_map::labeled_hash(
-            LABEL_ASSETS,
-            &asset_hashes.root_hash(),
-        )),
+        HashTree::Pruned(assets.root_hash()),
         ic_certified_map::labeled(LABEL_SIG, witness),
     );
 

src/internet_identity/src/http.rs

@@ -1,6 +1,6 @@
 use crate::archive::ArchiveState;
 use crate::assets::ContentType;
-use crate::{assets, state, IC0_APP_DOMAIN, INTERNETCOMPUTER_ORG_DOMAIN, LABEL_ASSETS, LABEL_SIG};
+use crate::{assets, state, IC0_APP_DOMAIN, INTERNETCOMPUTER_ORG_DOMAIN, LABEL_SIG};
 use base64::engine::general_purpose::STANDARD as BASE64;
 use base64::Engine;
 use ic_cdk::api::stable::stable64_size;
@@ -84,26 +84,28 @@ pub fn http_request(req: HttpRequest) -> HttpResponse {
             let mut headers = security_headers();
             headers.push(certificate_header);
 
-            state::assets(|a| match a.get(probably_an_asset) {
-                Some((asset_headers, value)) => {
-                    headers.append(&mut asset_headers.clone());
+            state::assets(
+                |certified_assets| match certified_assets.assets.get(probably_an_asset) {
+                    Some((asset_headers, data)) => {
+                        headers.append(&mut asset_headers.clone());
 
-                    HttpResponse {
-                        status_code: 200,
+                        HttpResponse {
+                            status_code: 200,
+                            headers,
+                            body: ByteBuf::from(data.clone()),
+                            upgrade: None,
+                            streaming_strategy: None,
+                        }
+                    }
+                    None => HttpResponse {
+                        status_code: 404,
                         headers,
-                        body: ByteBuf::from(value.clone()),
+                        body: ByteBuf::from(format!("Asset {probably_an_asset} not found.")),
                         upgrade: None,
                         streaming_strategy: None,
-                    }
-                }
-                None => HttpResponse {
-                    status_code: 404,
-                    headers,
-                    body: ByteBuf::from(format!("Asset {probably_an_asset} not found.")),
-                    upgrade: None,
-                    streaming_strategy: None,
+                    },
                 },
-            })
+            )
         }
     }
 }
@@ -414,10 +416,9 @@ fn make_asset_certificate_header(asset_name: &str) -> (String, String) {
     let certificate = data_certificate().unwrap_or_else(|| {
         trap("data certificate is only available in query calls");
     });
-    state::asset_hashes_and_sigs(|asset_hashes, sigs| {
-        let witness = asset_hashes.witness(asset_name.as_bytes());
+    state::assets_and_signatures(|assets, sigs| {
         let tree = ic_certified_map::fork(
-            ic_certified_map::labeled(LABEL_ASSETS, witness),
+            assets.witness(asset_name),
             HashTree::Pruned(ic_certified_map::labeled_hash(LABEL_SIG, &sigs.root_hash())),
         );
         let mut serializer = serde_cbor::ser::Serializer::new(vec![]);

src/internet_identity/src/main.rs

@@ -6,7 +6,6 @@ use crate::storage::anchor::Anchor;
 use candid::{candid_method, Principal};
 use ic_cdk::api::{caller, set_certified_data, trap};
 use ic_cdk_macros::{init, post_upgrade, pre_upgrade, query, update};
-use ic_certified_map::AsHashTree;
 use internet_identity_interface::archive::types::{BufferedEntry, Operation};
 use internet_identity_interface::http_gateway::{HttpRequest, HttpResponse};
 use internet_identity_interface::internet_identity::types::*;
@@ -31,7 +30,6 @@ const MINUTE_NS: u64 = secs_to_nanos(60);
 const HOUR_NS: u64 = 60 * MINUTE_NS;
 const DAY_NS: u64 = 24 * HOUR_NS;
 
-const LABEL_ASSETS: &[u8] = b"http_assets";
 const LABEL_SIG: &[u8] = b"sig";
 
 // Note: concatenating const &str is a hassle in rust. It seemed easiest to just repeat.
@@ -414,10 +412,10 @@ fn save_persistent_state() {
 
 fn update_root_hash() {
     use ic_certified_map::{fork_hash, labeled_hash};
-    state::asset_hashes_and_sigs(|asset_hashes, sigs| {
+    state::assets_and_signatures(|assets, sigs| {
         let prefixed_root_hash = fork_hash(
-            // NB: Labels added in lexicographic order
-            &labeled_hash(LABEL_ASSETS, &asset_hashes.root_hash()),
+            &assets.root_hash(),
+            // NB: sigs have to be added last due to lexical ordering of labels
             &labeled_hash(LABEL_SIG, &sigs.root_hash()),
         );
         set_certified_data(&prefixed_root_hash[..]);

src/internet_identity/src/state.rs

@@ -1,32 +1,28 @@
 use crate::archive::{ArchiveData, ArchiveState, ArchiveStatusCache};
+use crate::assets::CertifiedAssets;
 use crate::state::temp_keys::TempKeys;
 use crate::storage::anchor::Anchor;
 use crate::storage::{StableMemory, DEFAULT_RANGE_SIZE};
 use crate::{Salt, Storage};
 use candid::{CandidType, Deserialize, Principal};
 use ic_cdk::api::time;
 use ic_cdk::{call, trap};
-use ic_certified_map::{Hash, RbTree};
 use ic_stable_structures::DefaultMemoryImpl;
 use internet_identity::signature_map::SignatureMap;
-use internet_identity_interface::http_gateway::HeaderField;
 use internet_identity_interface::internet_identity::types::*;
 use std::cell::{Cell, RefCell};
 use std::collections::HashMap;
 use std::ops::{Deref, DerefMut};
 use std::time::Duration;
 
-pub type Assets = HashMap<String, (Vec<HeaderField>, Vec<u8>)>;
-pub type AssetHashes = RbTree<String, Hash>;
-
 mod temp_keys;
 
 // Default value for max number of delegation origins to store in the list of latest used delegation origins
 const MAX_NUM_DELEGATION_ORIGINS: u64 = 1000;
 
 thread_local! {
     static STATE: State = State::default();
-    static ASSETS: RefCell<Assets> = RefCell::new(HashMap::default());
+    static ASSETS: RefCell<CertifiedAssets> = RefCell::new(CertifiedAssets::default());
 }
 
 pub struct TentativeDeviceRegistration {
@@ -126,7 +122,6 @@ struct State {
     sigs: RefCell<SignatureMap>,
     // Temporary keys that can be used in lieu of a particular device
     temp_keys: RefCell<TempKeys>,
-    asset_hashes: RefCell<AssetHashes>,
     last_upgrade_timestamp: Cell<Timestamp>,
     // note: we COULD persist this through upgrades, although this is currently NOT persisted
     // through upgrades
@@ -153,7 +148,6 @@ impl Default for State {
             storage_state: RefCell::new(StorageState::Uninitialised),
             sigs: RefCell::new(SignatureMap::default()),
             temp_keys: RefCell::new(TempKeys::default()),
-            asset_hashes: RefCell::new(AssetHashes::default()),
             last_upgrade_timestamp: Cell::new(0),
             inflight_challenges: RefCell::new(HashMap::new()),
             tentative_device_registrations: RefCell::new(HashMap::new()),
@@ -309,18 +303,16 @@ pub fn tentative_device_registrations_mut<R>(
     STATE.with(|s| f(&mut s.tentative_device_registrations.borrow_mut()))
 }
 
-pub fn assets<R>(f: impl FnOnce(&Assets) -> R) -> R {
+pub fn assets<R>(f: impl FnOnce(&CertifiedAssets) -> R) -> R {
     ASSETS.with(|assets| f(&assets.borrow()))
 }
 
-pub fn assets_and_hashes_mut<R>(f: impl FnOnce(&mut Assets, &mut AssetHashes) -> R) -> R {
-    ASSETS.with(|assets| {
-        STATE.with(|s| f(&mut assets.borrow_mut(), &mut s.asset_hashes.borrow_mut()))
-    })
+pub fn assets_mut<R>(f: impl FnOnce(&mut CertifiedAssets) -> R) -> R {
+    ASSETS.with(|assets| f(&mut assets.borrow_mut()))
 }
 
-pub fn asset_hashes_and_sigs<R>(f: impl FnOnce(&AssetHashes, &SignatureMap) -> R) -> R {
-    STATE.with(|s| f(&s.asset_hashes.borrow(), &s.sigs.borrow()))
+pub fn assets_and_signatures<R>(f: impl FnOnce(&CertifiedAssets, &SignatureMap) -> R) -> R {
+    ASSETS.with(|assets| STATE.with(|s| f(&assets.borrow(), &s.sigs.borrow())))
 }
 
 pub fn signature_map<R>(f: impl FnOnce(&SignatureMap) -> R) -> R {