1 files changed, 255 insertions, 0 deletions
diff --git a/toolkit/components/xulstore/src/statics.rs b/toolkit/components/xulstore/src/statics.rs
new file mode 100644
index 0000000000..4b988b6062
--- /dev/null
+++ b/toolkit/components/xulstore/src/statics.rs
@@ -0,0 +1,255 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+use crate::{
+    error::{XULStoreError, XULStoreResult},
+    ffi::ProfileChangeObserver,
+    make_key, SEPARATOR,
+};
+use moz_task::is_main_thread;
+use nsstring::nsString;
+use once_cell::sync::Lazy;
+use rkv::backend::{SafeMode, SafeModeDatabase, SafeModeEnvironment};
+use rkv::{StoreOptions, Value};
+use std::{
+    collections::BTreeMap,
+    fs::{create_dir_all, remove_file, File},
+    path::PathBuf,
+    str,
+    sync::{Arc, Mutex, RwLock},
+};
+use xpcom::{
+    interfaces::{nsIFile, nsIObserverService, nsIProperties, nsIXULRuntime},
+    RefPtr, XpCom,
+};
+
+type Manager = rkv::Manager<SafeModeEnvironment>;
+type Rkv = rkv::Rkv<SafeModeEnvironment>;
+type SingleStore = rkv::SingleStore<SafeModeDatabase>;
+type XULStoreCache = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
+
+pub struct Database {
+    pub rkv: Arc<RwLock<Rkv>>,
+    pub store: SingleStore,
+}
+
+impl Database {
+    fn new(rkv: Arc<RwLock<Rkv>>, store: SingleStore) -> Database {
+        Database { rkv, store }
+    }
+}
+
+static PROFILE_DIR: Lazy<Mutex<Option<PathBuf>>> = Lazy::new(|| {
+    observe_profile_change();
+    Mutex::new(get_profile_dir().ok())
+});
+
+pub(crate) static DATA_CACHE: Lazy<Mutex<Option<XULStoreCache>>> =
+    Lazy::new(|| Mutex::new(cache_data().ok()));
+
+pub(crate) fn get_database() -> XULStoreResult<Database> {
+    let mut manager = Manager::singleton().write()?;
+    let xulstore_dir = get_xulstore_dir()?;
+    let xulstore_path = xulstore_dir.as_path();
+    let rkv = manager.get_or_create(xulstore_path, Rkv::new::<SafeMode>)?;
+    let store = rkv.read()?.open_single("db", StoreOptions::create())?;
+    Ok(Database::new(rkv, store))
+}
+
+pub(crate) fn update_profile_dir() {
+    // Failure to update the dir isn't fatal (although it means that we won't
+    // persist XULStore data for this session), so we don't return a result.
+    // But we use a closure returning a result to enable use of the ? operator.
+    (|| -> XULStoreResult<()> {
+        {
+            let mut profile_dir_guard = PROFILE_DIR.lock()?;
+            *profile_dir_guard = get_profile_dir().ok();
+        }
+
+        let mut cache_guard = DATA_CACHE.lock()?;
+        *cache_guard = cache_data().ok();
+
+        Ok(())
+    })()
+    .unwrap_or_else(|err| error!("error updating profile dir: {}", err));
+}
+
+fn get_profile_dir() -> XULStoreResult<PathBuf> {
+    // We can't use getter_addrefs() here because get_DirectoryService()
+    // returns its nsIProperties interface, and its Get() method returns
+    // a directory via its nsQIResult out param, which gets translated to
+    // a `*mut *mut libc::c_void` in Rust, whereas getter_addrefs() expects
+    // a closure with a `*mut *const T` parameter.
+
+    let dir_svc: RefPtr<nsIProperties> =
+        xpcom::components::Directory::service().map_err(|_| XULStoreError::Unavailable)?;
+    let mut profile_dir = xpcom::GetterAddrefs::<nsIFile>::new();
+    unsafe {
+        dir_svc
+            .Get(
+                cstr!("ProfD").as_ptr(),
+                &nsIFile::IID,
+                profile_dir.void_ptr(),
+            )
+            .to_result()
+            .or_else(|_| {
+                dir_svc
+                    .Get(
+                        cstr!("ProfDS").as_ptr(),
+                        &nsIFile::IID,
+                        profile_dir.void_ptr(),
+                    )
+                    .to_result()
+            })?;
+    }
+    let profile_dir = profile_dir.refptr().ok_or(XULStoreError::Unavailable)?;
+
+    let mut profile_path = nsString::new();
+    unsafe {
+        profile_dir.GetPath(&mut *profile_path).to_result()?;
+    }
+
+    let path = String::from_utf16(&profile_path[..])?;
+    Ok(PathBuf::from(&path))
+}
+
+fn get_xulstore_dir() -> XULStoreResult<PathBuf> {
+    let mut xulstore_dir = PROFILE_DIR
+        .lock()?
+        .clone()
+        .ok_or(XULStoreError::Unavailable)?;
+    xulstore_dir.push("xulstore");
+
+    create_dir_all(xulstore_dir.clone())?;
+
+    Ok(xulstore_dir)
+}
+
+fn observe_profile_change() {
+    assert!(is_main_thread());
+
+    // Failure to observe the change isn't fatal (although it means we won't
+    // persist XULStore data for this session), so we don't return a result.
+    // But we use a closure returning a result to enable use of the ? operator.
+    (|| -> XULStoreResult<()> {
+        // Observe profile changes so we can update this directory accordingly.
+        let obs_svc: RefPtr<nsIObserverService> =
+            xpcom::components::Observer::service().map_err(|_| XULStoreError::Unavailable)?;
+        let observer = ProfileChangeObserver::new();
+        unsafe {
+            obs_svc
+                .AddObserver(
+                    observer.coerce(),
+                    cstr!("profile-after-change").as_ptr(),
+                    false,
+                )
+                .to_result()?
+        };
+        Ok(())
+    })()
+    .unwrap_or_else(|err| error!("error observing profile change: {}", err));
+}
+
+fn in_safe_mode() -> XULStoreResult<bool> {
+    let xul_runtime: RefPtr<nsIXULRuntime> =
+        xpcom::components::XULRuntime::service().map_err(|_| XULStoreError::Unavailable)?;
+    let mut in_safe_mode = false;
+    unsafe {
+        xul_runtime.GetInSafeMode(&mut in_safe_mode).to_result()?;
+    }
+    Ok(in_safe_mode)
+}
+
+fn cache_data() -> XULStoreResult<XULStoreCache> {
+    let db = get_database()?;
+    maybe_migrate_data(&db, db.store);
+
+    let mut all = XULStoreCache::default();
+    if in_safe_mode()? {
+        return Ok(all);
+    }
+
+    let env = db.rkv.read()?;
+    let reader = env.read()?;
+    let iterator = db.store.iter_start(&reader)?;
+
+    for result in iterator {
+        let (key, value): (&str, String) = match result {
+            Ok((key, value)) => match (str::from_utf8(&key), unwrap_value(&value)) {
+                (Ok(key), Ok(value)) => (key, value),
+                (Err(err), _) => return Err(err.into()),
+                (_, Err(err)) => return Err(err),
+            },
+            Err(err) => return Err(err.into()),
+        };
+
+        let parts = key.split(SEPARATOR).collect::<Vec<&str>>();
+        if parts.len() != 3 {
+            return Err(XULStoreError::UnexpectedKey(key.to_string()));
+        }
+        let (doc, id, attr) = (
+            parts[0].to_string(),
+            parts[1].to_string(),
+            parts[2].to_string(),
+        );
+
+        all.entry(doc)
+            .or_default()
+            .entry(id)
+            .or_default()
+            .entry(attr)
+            .or_insert(value);
+    }
+
+    Ok(all)
+}
+
+fn maybe_migrate_data(db: &Database, store: SingleStore) {
+    // Failure to migrate data isn't fatal, so we don't return a result.
+    // But we use a closure returning a result to enable use of the ? operator.
+    (|| -> XULStoreResult<()> {
+        let mut old_datastore = PROFILE_DIR
+            .lock()?
+            .clone()
+            .ok_or(XULStoreError::Unavailable)?;
+        old_datastore.push("xulstore.json");
+        if !old_datastore.exists() {
+            debug!("old datastore doesn't exist: {:?}", old_datastore);
+            return Ok(());
+        }
+
+        let file = File::open(old_datastore.clone())?;
+        let json: XULStoreCache = serde_json::from_reader(file)?;
+
+        let env = db.rkv.read()?;
+        let mut writer = env.write()?;
+
+        for (doc, ids) in json {
+            for (id, attrs) in ids {
+                for (attr, value) in attrs {
+                    let key = make_key(&doc, &id, &attr);
+                    store.put(&mut writer, &key, &Value::Str(&value))?;
+                }
+            }
+        }
+
+        writer.commit()?;
+
+        remove_file(old_datastore)?;
+
+        Ok(())
+    })()
+    .unwrap_or_else(|err| error!("error migrating data: {}", err));
+}
+
+fn unwrap_value(value: &Value) -> XULStoreResult<String> {
+    match value {
+        Value::Str(val) => Ok(val.to_string()),
+
+        // This should never happen, but it could happen in theory
+        // if someone writes a different kind of value into the store
+        // using a more general API (kvstore, rkv, LMDB).
+        _ => Err(XULStoreError::UnexpectedValue),
+    }
+}