From 6a41b490f0762a1d27b0cc7aa9836f84e76d119c Mon Sep 17 00:00:00 2001
From: nazeh <ar.nazeh@gmail.com>
Date: Tue, 3 Dec 2024 14:20:48 +0300
Subject: [PATCH] feat(homeserver): mark unsafe methods caused by LMDB

---
 pubky-homeserver/src/core.rs     | 10 +++++-----
 pubky-homeserver/src/database.rs |  4 +++-
 pubky-homeserver/src/main.rs     | 18 ++++++++++--------
 pubky-homeserver/src/server.rs   | 17 ++++++-----------
 4 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/pubky-homeserver/src/core.rs b/pubky-homeserver/src/core.rs
index 7523fc2..58f3eb0 100644
--- a/pubky-homeserver/src/core.rs
+++ b/pubky-homeserver/src/core.rs
@@ -19,7 +19,6 @@ pub(crate) struct AppState {
     pub(crate) db: DB,
     pub(crate) pkarr_client: pkarr::Client,
     pub(crate) config: Config,
-    pub(crate) port: u16,
 }
 
 #[derive(Debug)]
@@ -30,10 +29,12 @@ pub struct HomeserverCore {
 }
 
 impl HomeserverCore {
-    pub fn new(config: &Config) -> Result<Self> {
+    /// # Safety
+    /// HomeserverCore uses LMDB, [opening][heed::EnvOpenOptions::open] which comes with some safety precautions.
+    pub unsafe fn new(config: &Config) -> Result<Self> {
         tracing::debug!(?config);
 
-        let db = DB::open(config.clone())?;
+        let db = unsafe { DB::open(config.clone())? };
 
         let mut dht_settings = pkarr::mainline::Settings::default();
 
@@ -53,7 +54,6 @@ impl HomeserverCore {
             db,
             pkarr_client: pkarr_client.clone(),
             config: config.clone(),
-            port: config.port(),
         };
 
         let router = crate::routes::create_app(state.clone());
@@ -66,7 +66,7 @@ impl HomeserverCore {
     pub fn test() -> Result<Self> {
         let testnet = pkarr::mainline::Testnet::new(0).expect("ignore");
 
-        HomeserverCore::new(&Config::test(&testnet))
+        unsafe { HomeserverCore::new(&Config::test(&testnet)) }
     }
 
     // TODO: move this logic to a common place.
diff --git a/pubky-homeserver/src/database.rs b/pubky-homeserver/src/database.rs
index 8ea4f48..4560d8d 100644
--- a/pubky-homeserver/src/database.rs
+++ b/pubky-homeserver/src/database.rs
@@ -19,7 +19,9 @@ pub struct DB {
 }
 
 impl DB {
-    pub fn open(config: Config) -> anyhow::Result<Self> {
+    /// # Safety
+    /// Opening [LMDB][heed::EnvOpenOptions::open] is backed by a memory map which comes with some safety precautions.
+    pub unsafe fn open(config: Config) -> anyhow::Result<Self> {
         let buffers_dir = config.storage().clone().join("buffers");
 
         // Cleanup buffers.
diff --git a/pubky-homeserver/src/main.rs b/pubky-homeserver/src/main.rs
index dad25df..3497c1a 100644
--- a/pubky-homeserver/src/main.rs
+++ b/pubky-homeserver/src/main.rs
@@ -31,14 +31,16 @@ async fn main() -> Result<()> {
         )
         .init();
 
-    let server = Homeserver::start(if args.testnet {
-        Config::testnet()
-    } else if let Some(config_path) = args.config {
-        Config::load(config_path).await?
-    } else {
-        Config::default()
-    })
-    .await?;
+    let server = unsafe {
+        Homeserver::start(if args.testnet {
+            Config::testnet()
+        } else if let Some(config_path) = args.config {
+            Config::load(config_path).await?
+        } else {
+            Config::default()
+        })
+        .await?
+    };
 
     server.run_until_done().await?;
 
diff --git a/pubky-homeserver/src/server.rs b/pubky-homeserver/src/server.rs
index 2d2d58d..424d5a1 100644
--- a/pubky-homeserver/src/server.rs
+++ b/pubky-homeserver/src/server.rs
@@ -24,7 +24,9 @@ pub struct Homeserver {
 }
 
 impl Homeserver {
-    pub async fn start(config: Config) -> Result<Self> {
+    /// # Safety
+    /// Homeserver uses LMDB, [opening][heed::EnvOpenOptions::open] which comes with some safety precautions.
+    pub async unsafe fn start(config: Config) -> Result<Self> {
         let mut tasks = JoinSet::new();
 
         let listener = TcpListener::bind(SocketAddr::from(([0, 0, 0, 0], config.port())))?;
@@ -33,16 +35,13 @@ impl Homeserver {
 
         let keypair = config.keypair().clone();
 
-        let mut core = HomeserverCore::new(&config)?;
-
-        // Update the port.
-        core.state.port = port;
-
         let acceptor = RustlsAcceptor::new(RustlsConfig::from_config(Arc::new(
             keypair.to_rpk_rustls_server_config(),
         )));
         let server = axum_server::from_tcp(listener).acceptor(acceptor);
 
+        let core = unsafe { HomeserverCore::new(&config)? };
+
         // Spawn http server task
         tasks.spawn(
             server.serve(
@@ -69,15 +68,11 @@ impl Homeserver {
     pub async fn start_test(testnet: &Testnet) -> Result<Self> {
         info!("Running testnet..");
 
-        Homeserver::start(Config::test(testnet)).await
+        unsafe { Homeserver::start(Config::test(testnet)).await }
     }
 
     // === Getters ===
 
-    pub fn port(&self) -> u16 {
-        self.state.port
-    }
-
     pub fn public_key(&self) -> PublicKey {
         self.state.config.keypair().public_key()
     }