Code cleanup

Cargo.toml

@@ -79,6 +79,7 @@ tower-http = "0.5"
 tracing = "0.1"
 tracing-subscriber = "0.3"
 ts-rs = "7"
+ulid = "1"
 unicode-segmentation = "1.10"
 url = "2.3"
 walkdir = "2.3"
@@ -108,6 +109,7 @@ tokio = { workspace = true, features = ["full"] }
 tokio-stream = { workspace = true }
 tower = { workspace = true }
 tower-http = { workspace = true, features = ["cors", "trace", "fs"] }
+ulid = { workspace = true }
 url = { workspace = true }
 idna = { workspace = true }
 
@@ -118,6 +120,7 @@ tracing = { workspace = true }
 
 cfg-if = { workspace = true }
 
+bytes = { workspace = true }
 compact_str = { workspace = true }
 either = { workspace = true }
 futures = { workspace = true }

ext_federation/src/ap_client.rs

@@ -11,9 +11,10 @@ use url::Url;
 use magnetar_core::web_model::content_type::ContentActivityStreams;
 
 use crate::{
-    ApClientService,
-    ApSignature,
-    ApSigningField, ApSigningHeaders, client::federation_client::{FederationClient, FederationClientError}, crypto::{ApSigningError, ApSigningKey, SigningAlgorithm}, SigningInput, SigningParts,
+    client::federation_client::{FederationClient, FederationClientError},
+    crypto::{ApSigningError, ApSigningKey, SigningAlgorithm},
+    ApClientService, ApSignature, ApSigningField,
+    ApSigningHeaders, SigningInput, SigningParts,
 };
 
 pub struct ApClientServiceDefaultProvider {
@@ -414,9 +415,9 @@ mod test {
 
     use crate::{
         ap_client::ApClientServiceDefaultProvider,
-        ApClientService,
         client::federation_client::FederationClient,
         crypto::{ApHttpPrivateKey, SigningAlgorithm},
+        ApClientService,
     };
 
     #[tokio::test]
@@ -440,7 +441,7 @@ mod test {
 
         let val = ap_client
             .signed_get(
-                ApHttpPrivateKey::Rsa(Box::new(Cow::Owned(rsa_key)))
+                ApHttpPrivateKey::Rsa(Cow::Owned(Box::new(rsa_key)))
                     .create_signing_key(&key_id, SigningAlgorithm::RsaSha256)
                     .into_diagnostic()?,
                 SigningAlgorithm::RsaSha256,

ext_federation/src/crypto.rs

@@ -1,12 +1,16 @@
-use std::{borrow::Cow, fmt::Display};
-
+use rsa::pkcs1::DecodeRsaPrivateKey;
+use rsa::pkcs1::DecodeRsaPublicKey;
+use rsa::pkcs8::DecodePrivateKey;
+use rsa::pkcs8::DecodePublicKey;
 use rsa::signature::Verifier;
 use rsa::{
     sha2::{Sha256, Sha512},
     signature::Signer,
 };
-
 use serde::{Deserialize, Serialize};
+use std::fmt::Formatter;
+use std::str::FromStr;
+use std::{borrow::Cow, fmt::Display};
 use strum::AsRefStr;
 use thiserror::Error;
 
@@ -35,7 +39,7 @@ pub enum SigningAlgorithm {
 }
 
 impl Display for SigningAlgorithm {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Hs2019 => write!(f, "hs2019"),
             Self::RsaSha256 => write!(f, "rsa-sha256"),
@@ -61,6 +65,59 @@ pub enum ApHttpPublicKey<'a> {
     Ed25519(Cow<'a, ed25519_dalek::VerifyingKey>),
 }
 
+#[derive(Debug, Copy, Clone, Error)]
+pub struct ApHttpPublicKeyParseError;
+
+impl Display for ApHttpPublicKeyParseError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "Failed to parse the public key: No available parser could decode the PEM string"
+        )
+    }
+}
+
+impl FromStr for ApHttpPublicKey<'_> {
+    type Err = ApHttpPublicKeyParseError;
+
+    fn from_str(input_pem: &str) -> Result<Self, Self::Err> {
+        let pem = input_pem.trim();
+
+        let parse_pkcs1_rsa: &dyn Fn(_) -> _ = &|p| {
+            Some(ApHttpPublicKey::Rsa(Cow::Owned(
+                rsa::RsaPublicKey::from_pkcs1_pem(p).ok()?,
+            )))
+        };
+        let parse_spki_rsa: &dyn Fn(_) -> _ = &|p| {
+            Some(ApHttpPublicKey::Rsa(Cow::Owned(
+                rsa::RsaPublicKey::from_public_key_pem(p).ok()?,
+            )))
+        };
+        let parse_spki_ed25519: &dyn Fn(_) -> _ = &|p| {
+            Some(ApHttpPublicKey::Ed25519(Cow::Owned(
+                ed25519_dalek::VerifyingKey::from_public_key_pem(p).ok()?,
+            )))
+        };
+
+        // Some heuristics
+        let parsers: &[_] = match pem {
+            p if p.starts_with("-----BEGIN PUBLIC KEY-----") => {
+                &[parse_spki_rsa, parse_spki_ed25519]
+            }
+            p if p.starts_with("-----BEGIN RSA PUBLIC KEY-----") => &[parse_pkcs1_rsa],
+            _ => &[parse_spki_rsa, parse_spki_ed25519, parse_pkcs1_rsa],
+        };
+
+        for parser in parsers {
+            if let Some(k) = parser(pem) {
+                return Ok(k);
+            }
+        }
+
+        Err(ApHttpPublicKeyParseError)
+    }
+}
+
 impl ApHttpVerificationKey<'_> {
     pub fn verify(&self, message: &[u8], signature: &[u8]) -> Result<(), ApVerificationError> {
         match self {
@@ -109,12 +166,10 @@ impl ApHttpPublicKey<'_> {
                 ));
                 Ok(verification_key.verify(message, signature)?)
             }
-            (_, SigningAlgorithm::RsaSha256) => {
-                return Err(ApVerificationError::KeyAlgorithmMismatch(
+            (_, SigningAlgorithm::RsaSha256) => Err(ApVerificationError::KeyAlgorithmMismatch(
                 algorithm,
                 self.as_ref().to_owned(),
-                ));
-            }
+            )),
             (Self::Ed25519(key), SigningAlgorithm::Hs2019) => {
                 let verification_key = ApHttpVerificationKey::Ed25519(Cow::Borrowed(key.as_ref()));
                 Ok(verification_key.verify(message, signature)?)
@@ -126,17 +181,72 @@ impl ApHttpPublicKey<'_> {
 #[derive(Debug, Clone, AsRefStr)]
 pub enum ApHttpPrivateKey<'a> {
     #[strum(serialize = "rsa")]
-    Rsa(Box<Cow<'a, rsa::RsaPrivateKey>>),
+    Rsa(Cow<'a, Box<rsa::RsaPrivateKey>>),
     #[strum(serialize = "ed25519")]
     Ed25519(Cow<'a, ed25519_dalek::SecretKey>),
 }
 
+#[derive(Debug, Copy, Clone, Error)]
+pub struct ApHttpPrivateKeyParseError;
+
+impl Display for ApHttpPrivateKeyParseError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "Failed to parse the private key: No available parser could decode the PEM string"
+        )
+    }
+}
+
+impl FromStr for ApHttpPrivateKey<'_> {
+    type Err = ApHttpPrivateKeyParseError;
+
+    fn from_str(input_pem: &str) -> Result<Self, Self::Err> {
+        let pem = input_pem.trim();
+
+        let parse_pkcs1_rsa: &dyn Fn(_) -> _ = &|p| {
+            Some(ApHttpPrivateKey::Rsa(Cow::Owned(Box::new(
+                rsa::RsaPrivateKey::from_pkcs1_pem(p).ok()?,
+            ))))
+        };
+        let parse_pkcs8_rsa: &dyn Fn(_) -> _ = &|p| {
+            Some(ApHttpPrivateKey::Rsa(Cow::Owned(Box::new(
+                rsa::RsaPrivateKey::from_pkcs8_pem(p).ok()?,
+            ))))
+        };
+        let parse_pkcs8_ed25519: &dyn Fn(_) -> _ = &|p| {
+            Some(ApHttpPrivateKey::Ed25519(Cow::Owned(
+                ed25519_dalek::SigningKey::from_pkcs8_pem(p)
+                    .ok()?
+                    .to_bytes(),
+            )))
+        };
+
+        // Some heuristics
+        let parsers: &[_] = match pem {
+            p if p.contains("-----BEGIN PRIVATE KEY-----") => {
+                &[parse_pkcs8_rsa, parse_pkcs8_ed25519]
+            }
+            p if p.contains("-----BEGIN RSA PRIVATE KEY-----") => &[parse_pkcs1_rsa],
+            _ => &[parse_pkcs8_rsa, parse_pkcs8_ed25519, parse_pkcs1_rsa],
+        };
+
+        for parser in parsers {
+            if let Some(k) = parser(pem) {
+                return Ok(k);
+            }
+        }
+
+        Err(ApHttpPrivateKeyParseError)
+    }
+}
+
 #[derive(Debug, Clone, AsRefStr)]
 pub enum ApHttpSigningKey<'a> {
     #[strum(serialize = "rsa-sha256")]
-    RsaSha256(Cow<'a, rsa::pkcs1v15::SigningKey<rsa::sha2::Sha256>>),
+    RsaSha256(Cow<'a, rsa::pkcs1v15::SigningKey<Sha256>>),
     #[strum(serialize = "rsa-sha512")]
-    RsaSha512(Cow<'a, rsa::pkcs1v15::SigningKey<rsa::sha2::Sha512>>),
+    RsaSha512(Cow<'a, rsa::pkcs1v15::SigningKey<Sha512>>),
     #[strum(serialize = "ed25519")]
     Ed25519(Cow<'a, ed25519_dalek::SigningKey>),
 }
@@ -192,7 +302,7 @@ impl ApHttpPrivateKey<'_> {
             key: match (self, algorithm) {
                 (Self::Rsa(key), SigningAlgorithm::RsaSha256 | SigningAlgorithm::Hs2019) => {
                     ApHttpSigningKey::RsaSha256(Cow::Owned(rsa::pkcs1v15::SigningKey::new(
-                        key.clone().into_owned(),
+                        *key.as_ref().to_owned(),
                     )))
                 }
                 (Self::Ed25519(key), SigningAlgorithm::Hs2019) => ApHttpSigningKey::Ed25519(

ext_model/src/lib.rs

@@ -9,15 +9,15 @@ use sea_orm::{
     ColumnTrait, ConnectOptions, DatabaseConnection, DbErr, EntityTrait, QueryFilter,
     TransactionTrait,
 };
-use serde::{Deserialize, Deserializer, Serialize};
 use serde::de::Error;
+use serde::{Deserialize, Deserializer, Serialize};
 use serde_json::Value;
 use strum::IntoStaticStr;
 use thiserror::Error;
 use tokio::select;
 use tokio_util::sync::CancellationToken;
-use tracing::{error, info, trace, warn};
 use tracing::log::LevelFilter;
+use tracing::{error, info, trace, warn};
 use url::Host;
 
 pub use ck;
@@ -122,13 +122,25 @@ impl CalckeyModel {
             .await?)
     }
 
-    pub async fn get_user_and_profile_by_id(&self, id: &str) -> Result<Option<(user::Model, user_profile::Model)>, CalckeyDbError> {
-        Ok(user::Entity::find()
+    pub async fn get_user_for_cache_by_id(&self, id: &str) -> Result<Option<(user::Model, user_profile::Model, user_keypair::Model)>, CalckeyDbError> {
+        let txn = self.0.begin().await?;
+
+        let Some((user, Some(profile))) = user::Entity::find()
             .filter(user::Column::Id.eq(id))
             .find_also_related(user_profile::Entity)
-            .one(&self.0)
-            .await?
-            .and_then(|(u, p)| p.map(|pp| (u, pp))))
+            .one(&txn)
+            .await? else {
+            return Ok(None);
+        };
+
+        let Some(keys) = user_keypair::Entity::find()
+            .filter(user_keypair::Column::UserId.eq(id))
+            .one(&txn)
+            .await? else {
+            return Ok(None);
+        };
+
+        Ok(Some((user, profile, keys)))
     }
 
     pub async fn get_user_security_keys_by_id(
@@ -165,20 +177,30 @@ impl CalckeyModel {
             .await?)
     }
 
-    pub async fn get_user_and_profile_by_token(
+    pub async fn get_user_for_cache_by_token(
         &self,
         token: &str,
-    ) -> Result<Option<(user::Model, user_profile::Model)>, CalckeyDbError> {
-        Ok(user::Entity::find()
-            .filter(
-                user::Column::Token
+    ) -> Result<Option<(user::Model, user_profile::Model, user_keypair::Model)>, CalckeyDbError> {
+        let txn = self.0.begin().await?;
+
+        let Some((user, Some(profile))) = user::Entity::find()
+            .filter(user::Column::Token
                 .eq(token)
-                    .and(user::Column::Host.is_null()),
-            )
+                .and(user::Column::Host.is_null()))
             .find_also_related(user_profile::Entity)
-            .one(&self.0)
-            .await?
-            .and_then(|(u, p)| p.map(|pp| (u, pp))))
+            .one(&txn)
+            .await? else {
+            return Ok(None);
+        };
+
+        let Some(keys) = user_keypair::Entity::find()
+            .filter(user_keypair::Column::UserId.eq(&user.id))
+            .one(&txn)
+            .await? else {
+            return Ok(None);
+        };
+
+        Ok(Some((user, profile, keys)))
     }
 
     pub async fn get_user_by_uri(&self, uri: &str) -> Result<Option<user::Model>, CalckeyDbError> {

ext_model/src/note_model/mod.rs

@@ -45,7 +45,8 @@ impl NoteResolveMode {
         match self {
             NoteResolveMode::Single(id) => Ok(id_col.eq(id)),
             NoteResolveMode::Multiple(ids) => Ok(id_col.is_in(ids)),
-            // We add a CTE for pins
+            // We do this in a separate query, because before we used an inner join, and it caused
+            // a massive performance penalty
             NoteResolveMode::PinsFromUserId(user_id) => {
                 let cte_query = user_note_pining::Entity::find()
                     .column(user_note_pining::Column::NoteId)

magnetar_sdk/src/types/mod.rs

@@ -15,7 +15,7 @@ use serde::{Deserialize, Deserializer, Serialize};
 use ts_rs::TS;
 
 pub(crate) mod packed_time {
-    use chrono::{DateTime, NaiveDateTime, Utc};
+    use chrono::{DateTime, Utc};
     use serde::de::Error;
     use serde::{Deserialize, Deserializer, Serializer};
 
@@ -30,15 +30,12 @@ pub(crate) mod packed_time {
     where
         D: Deserializer<'de>,
     {
-        Ok(DateTime::<Utc>::from_naive_utc_and_offset(
-            NaiveDateTime::from_timestamp_millis(
+        DateTime::<Utc>::from_timestamp_millis(
             String::deserialize(deserializer)?
                 .parse::<i64>()
                 .map_err(Error::custom)?,
         )
-            .ok_or_else(|| Error::custom("millisecond value out of range"))?,
-            Utc,
-        ))
+            .ok_or_else(|| Error::custom("millisecond value out of range"))
     }
 }
 

src/service/gen_id.rs

@@ -0,0 +1,29 @@
+use std::{sync::Arc, time::SystemTime};
+
+use super::MagnetarService;
+
+pub struct GenIdService;
+
+impl GenIdService {
+    pub fn new_id(&self) -> ulid::Ulid {
+        ulid::Ulid::new()
+    }
+
+    pub fn new_id_str(&self) -> String {
+        self.new_id().to_string()
+    }
+
+    pub fn new_for_time(&self, time: impl Into<SystemTime>) -> ulid::Ulid {
+        ulid::Ulid::from_datetime(time.into())
+    }
+
+    pub fn new_str_for_time(&self, time: impl Into<SystemTime>) -> String {
+        self.new_for_time(time).to_string()
+    }
+}
+
+impl AsRef<GenIdService> for Arc<MagnetarService> {
+    fn as_ref(&self) -> &GenIdService {
+        &self.gen_id
+    }
+}

src/service/local_user_cache.rs

@@ -7,34 +7,44 @@ use thiserror::Error;
 use tokio::sync::Mutex;
 use tracing::error;
 
+use crate::web::ApiError;
 use magnetar_common::config::MagnetarConfig;
+use magnetar_federation::crypto::{ApHttpPrivateKey, ApHttpPrivateKeyParseError, ApHttpPublicKey, ApHttpPublicKeyParseError};
 use magnetar_model::{
     ck, CalckeyCache, CalckeyCacheError, CalckeyDbError, CalckeyModel, CalckeySub,
     InternalStreamMessage, SubMessage,
 };
 
-use crate::web::ApiError;
-
 #[derive(Debug, Error, VariantNames)]
 pub enum UserCacheError {
     #[error("Database error: {0}")]
     DbError(#[from] CalckeyDbError),
     #[error("Redis error: {0}")]
     RedisError(#[from] CalckeyCacheError),
+    #[error("Private key parse error: {0}")]
+    PrivateKeyParseError(#[from] ApHttpPrivateKeyParseError),
+    #[error("Public key parse error: {0}")]
+    PublicKeyParseError(#[from] ApHttpPublicKeyParseError),
 }
 
 #[derive(Debug, Clone)]
 pub struct CachedLocalUser {
     pub user: Arc<ck::user::Model>,
     pub profile: Arc<ck::user_profile::Model>,
+    pub private_key: Arc<ApHttpPrivateKey<'static>>,
+    pub public_key: Arc<ApHttpPublicKey<'static>>,
 }
 
-impl From<(ck::user::Model, ck::user_profile::Model)> for CachedLocalUser {
-    fn from((user, profile): (ck::user::Model, ck::user_profile::Model)) -> Self {
-        CachedLocalUser {
+impl TryFrom<(ck::user::Model, ck::user_profile::Model, ck::user_keypair::Model)> for CachedLocalUser {
+    type Error = UserCacheError;
+
+    fn try_from((user, profile, key_pair): (ck::user::Model, ck::user_profile::Model, ck::user_keypair::Model)) -> Result<Self, Self::Error> {
+        Ok(CachedLocalUser {
             user: Arc::new(user),
             profile: Arc::new(profile),
-        }
+            private_key: Arc::new(key_pair.private_key.parse()?),
+            public_key: Arc::new(key_pair.public_key.parse()?),
+        })
     }
 }
 
@@ -43,6 +53,8 @@ impl From<UserCacheError> for ApiError {
         let mut api_error: ApiError = match err {
             UserCacheError::DbError(err) => err.into(),
             UserCacheError::RedisError(err) => err.into(),
+            UserCacheError::PublicKeyParseError(err) => err.into(),
+            UserCacheError::PrivateKeyParseError(err) => err.into()
         };
 
         api_error.message = format!("Local user cache error: {}", api_error.message);
@@ -156,7 +168,7 @@ impl LocalUserCacheService {
                         | InternalStreamMessage::UserChangeSuspendedState { id, .. }
                         | InternalStreamMessage::RemoteUserUpdated { id }
                         | InternalStreamMessage::UserTokenRegenerated { id, .. } => {
-                            let user_profile = match db.get_user_and_profile_by_id(&id).await {
+                            let user_profile = match db.get_user_for_cache_by_id(&id).await {
                                 Ok(Some(m)) => m,
                                 Ok(None) => return,
                                 Err(e) => {
@@ -165,7 +177,15 @@ impl LocalUserCacheService {
                                 }
                             };
 
-                            cache.lock().await.refresh(&CachedLocalUser::from(user_profile));
+                            let cached: CachedLocalUser = match user_profile.try_into() {
+                                Ok(c) => c,
+                                Err(e) => {
+                                    error!("Error parsing user from database: {}", e);
+                                    return;
+                                }
+                            };
+
+                            cache.lock().await.refresh(&cached);
                         }
                         _ => {}
                     };
@@ -202,7 +222,7 @@ impl LocalUserCacheService {
             return Ok(Some(user));
         }
 
-        self.map_cache_user(self.db.get_user_and_profile_by_token(token).await?.map(CachedLocalUser::from))
+        self.map_cache_user(self.db.get_user_for_cache_by_token(token).await?.map(CachedLocalUser::try_from).transpose()?)
             .await
     }
 
@@ -216,6 +236,6 @@ impl LocalUserCacheService {
             return Ok(Some(user));
         }
 
-        self.map_cache_user(self.db.get_user_and_profile_by_id(id).await?.map(CachedLocalUser::from)).await
+        self.map_cache_user(self.db.get_user_for_cache_by_id(id).await?.map(CachedLocalUser::try_from).transpose()?).await
     }
 }

src/service/mod.rs

@@ -1,3 +1,4 @@
+use gen_id::GenIdService;
 use magnetar_common::config::MagnetarConfig;
 use magnetar_model::{ck, CalckeyCache, CalckeyModel};
 use std::fmt::{Debug, Formatter};
@@ -5,15 +6,16 @@ use std::time::Duration;
 use thiserror::Error;
 
 pub mod emoji_cache;
+pub mod gen_id;
 pub mod generic_id_cache;
 pub mod instance_cache;
 pub mod instance_meta_cache;
 pub mod local_user_cache;
 
 #[non_exhaustive]
-
 pub struct MagnetarService {
     pub db: CalckeyModel,
+    pub gen_id: GenIdService,
     pub cache: CalckeyCache,
     pub config: &'static MagnetarConfig,
     pub local_user_cache: local_user_cache::LocalUserCacheService,
@@ -67,6 +69,14 @@ impl MagnetarService {
             remote_instance_cache,
             emoji_cache,
             drive_file_cache,
+            gen_id: GenIdService,
         })
     }
+
+    pub fn service<T>(&self) -> &T
+    where
+        Self: AsRef<T>,
+    {
+        self.as_ref()
+    }
 }

src/web/mod.rs

@@ -3,6 +3,7 @@ use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
 use axum::Json;
 use magnetar_common::util::FediverseTagParseError;
+use magnetar_federation::crypto::{ApHttpPrivateKeyParseError, ApHttpPublicKeyParseError};
 use magnetar_model::{CalckeyCacheError, CalckeyDbError};
 use serde::Serialize;
 use serde_json::json;
@@ -190,3 +191,31 @@ impl From<ArgumentOutOfRange> for ApiError {
         }
     }
 }
+
+impl From<ApHttpPublicKeyParseError> for ApiError {
+    fn from(err: ApHttpPublicKeyParseError) -> Self {
+        Self {
+            status: StatusCode::INTERNAL_SERVER_ERROR,
+            code: "ApHttpPublicKeyParseError".error_code(),
+            message: if cfg!(debug_assertions) {
+                format!("User public key parse error: {}", err)
+            } else {
+                "User public key parse error".to_string()
+            },
+        }
+    }
+}
+
+impl From<ApHttpPrivateKeyParseError> for ApiError {
+    fn from(err: ApHttpPrivateKeyParseError) -> Self {
+        Self {
+            status: StatusCode::INTERNAL_SERVER_ERROR,
+            code: "ApHttpPrivateKeyParseError".error_code(),
+            message: if cfg!(debug_assertions) {
+                format!("User private key parse error: {}", err)
+            } else {
+                "User private key parse error".to_string()
+            },
+        }
+    }
+}