Skip to content

Commit

Permalink
user profile, idiomatic improvements
Browse files Browse the repository at this point in the history
  • Loading branch information
@aprxi
aprxi committed Aug 12, 2024
1 parent f061e5b commit 1bd1414
Show file tree
Hide file tree
Showing 10 changed files with 207 additions and 106 deletions.
14 changes: 14 additions & 0 deletions lumni/src/apps/api/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,10 @@ pub enum EncryptionError {
Utf8Error(std::string::FromUtf8Error),
SpkiError(rsa::pkcs8::spki::Error),
Pkcs8Error(rsa::pkcs8::Error),
KeyGenerationFailed(String),
EncryptionFailed(String),
DecryptionFailed(String),
InvalidKey(String),
Other(Box<dyn Error + Send + Sync>),
}

Expand All @@ -235,6 +239,16 @@ impl fmt::Display for EncryptionError {
}
EncryptionError::SpkiError(e) => write!(f, "SPKI error: {}", e),
EncryptionError::Pkcs8Error(e) => write!(f, "PKCS8 error: {}", e),
EncryptionError::KeyGenerationFailed(e) => {
write!(f, "Key generation failed: {}", e)
}
EncryptionError::EncryptionFailed(e) => {
write!(f, "Encryption failed: {}", e)
}
EncryptionError::DecryptionFailed(e) => {
write!(f, "Decryption failed: {}", e)
}
EncryptionError::InvalidKey(e) => write!(f, "Invalid key: {}", e),
EncryptionError::Other(e) => write!(f, "Other error: {}", e),
}
}
Expand Down
2 changes: 1 addition & 1 deletion lumni/src/apps/builtin/llm/prompt/src/chat/db/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ pub use lumni::Timestamp;
pub use model::{ModelIdentifier, ModelSpec};
use serde::{Deserialize, Serialize};
pub use store::ConversationDatabase;
pub use user_profile::UserProfileDbHandler;
pub use user_profile::{EncryptionMode, MaskMode, UserProfileDbHandler};

pub use super::ConversationCache;
use super::PromptRole;
Expand Down
48 changes: 26 additions & 22 deletions lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profile/content_operations.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
use lumni::api::error::ApplicationError;
use serde_json::{json, Map, Value as JsonValue};

use super::{DatabaseOperationError, UserProfileDbHandler};
use super::{
DatabaseOperationError, EncryptionMode, MaskMode, UserProfileDbHandler,
};
use crate::external as lumni;

impl UserProfileDbHandler {
Expand Down Expand Up @@ -88,24 +90,26 @@ impl UserProfileDbHandler {
pub fn process_settings(
&self,
value: &JsonValue,
encrypt: bool,
mask_encrypted: bool,
encryption_mode: EncryptionMode,
mask_mode: MaskMode,
) -> Result<JsonValue, ApplicationError> {
match value {
JsonValue::Object(obj) => {
let mut new_obj = Map::new();
for (k, v) in obj {
new_obj.insert(
k.clone(),
self.process_value(v, encrypt, mask_encrypted)?,
self.process_value(v, encryption_mode, mask_mode)?,
);
}
Ok(JsonValue::Object(new_obj))
}
JsonValue::Array(arr) => {
let new_arr: Result<Vec<JsonValue>, _> = arr
.iter()
.map(|v| self.process_settings(v, encrypt, mask_encrypted))
.map(|v| {
self.process_settings(v, encryption_mode, mask_mode)
})
.collect();
Ok(JsonValue::Array(new_arr?))
}
Expand All @@ -116,13 +120,13 @@ impl UserProfileDbHandler {
fn process_value(
&self,
value: &JsonValue,
encrypt: bool,
mask_encrypted: bool,
encryption_mode: EncryptionMode,
mask_mode: MaskMode,
) -> Result<JsonValue, ApplicationError> {
if encrypt {
if encryption_mode == EncryptionMode::Encrypt {
self.handle_encryption(value)
} else {
self.handle_decryption(value, mask_encrypted)
self.handle_decryption(value, mask_mode)
}
}

Expand All @@ -146,11 +150,11 @@ impl UserProfileDbHandler {
fn handle_decryption(
&self,
value: &JsonValue,
mask_encrypted: bool,
mask_mode: MaskMode,
) -> Result<JsonValue, ApplicationError> {
if Self::is_encrypted_value(value) {
if self.encryption_handler.is_some() {
if mask_encrypted {
if mask_mode == MaskMode::Mask {
Ok(JsonValue::String("*****".to_string()))
} else {
self.decrypt_value(value)
Expand All @@ -166,17 +170,17 @@ impl UserProfileDbHandler {
pub fn process_settings_with_metadata(
&self,
value: &JsonValue,
encrypt: bool,
mask_encrypted: bool,
encryption_mode: EncryptionMode,
mask_mode: MaskMode,
) -> Result<JsonValue, ApplicationError> {
match value {
JsonValue::Object(obj) => {
let mut new_obj = Map::new();
for (k, v) in obj {
let processed = self.process_value_with_metadata(
v,
encrypt,
mask_encrypted,
encryption_mode,
mask_mode,
)?;
new_obj.insert(k.clone(), processed);
}
Expand All @@ -188,8 +192,8 @@ impl UserProfileDbHandler {
.map(|v| {
self.process_settings_with_metadata(
v,
encrypt,
mask_encrypted,
encryption_mode,
mask_mode,
)
})
.collect();
Expand All @@ -202,13 +206,13 @@ impl UserProfileDbHandler {
fn process_value_with_metadata(
&self,
value: &JsonValue,
encrypt: bool,
mask_encrypted: bool,
encryption_mode: EncryptionMode,
mask_mode: MaskMode,
) -> Result<JsonValue, ApplicationError> {
if encrypt {
if encryption_mode == EncryptionMode::Encrypt {
self.handle_encryption(value)
} else {
let decrypted = self.handle_decryption(value, mask_encrypted)?;
let decrypted = self.handle_decryption(value, mask_mode)?;
if Self::is_encrypted_value(value) {
Ok(json!({
"value": decrypted,
Expand Down Expand Up @@ -237,4 +241,4 @@ impl UserProfileDbHandler {
false
}
}
}
}
53 changes: 33 additions & 20 deletions lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profile/database_operations.rs
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,16 @@
use std::path::PathBuf;

use lumni::api::error::{ApplicationError, EncryptionError};
use rusqlite::{params, OptionalExtension};

use super::{DatabaseOperationError, EncryptionHandler, UserProfileDbHandler};
use crate::external as lumni;

use super::{
DatabaseOperationError, EncryptionHandler, UserProfileDbHandler,
};


impl UserProfileDbHandler {
pub async fn profile_exists(&self, profile_name: &str) -> Result<bool, ApplicationError> {
pub async fn profile_exists(
&self,
profile_name: &str,
) -> Result<bool, ApplicationError> {
let mut db = self.db.lock().await;
db.process_queue_with_result(|tx| {
let count: i64 = tx
Expand All @@ -25,7 +25,10 @@ impl UserProfileDbHandler {
.map_err(ApplicationError::from)
}

pub async fn delete_profile(&self, profile_name: &str) -> Result<(), ApplicationError> {
pub async fn delete_profile(
&self,
profile_name: &str,
) -> Result<(), ApplicationError> {
let mut db = self.db.lock().await;

db.process_queue_with_result(|tx| {
Expand Down Expand Up @@ -56,7 +59,9 @@ impl UserProfileDbHandler {
.map_err(ApplicationError::from)
}

pub async fn get_default_profile(&self) -> Result<Option<String>, ApplicationError> {
pub async fn get_default_profile(
&self,
) -> Result<Option<String>, ApplicationError> {
let mut db = self.db.lock().await;
db.process_queue_with_result(|tx| {
tx.query_row(
Expand All @@ -70,7 +75,10 @@ impl UserProfileDbHandler {
.map_err(ApplicationError::from)
}

pub async fn set_default_profile(&self, profile_name: &str) -> Result<(), ApplicationError> {
pub async fn set_default_profile(
&self,
profile_name: &str,
) -> Result<(), ApplicationError> {
let mut db = self.db.lock().await;

db.process_queue_with_result(|tx| {
Expand All @@ -89,7 +97,10 @@ impl UserProfileDbHandler {
.map_err(ApplicationError::from)
}

pub async fn load_encryption_handler(&self, encryption_key_id: i64) -> Result<EncryptionHandler, ApplicationError> {
pub async fn load_encryption_handler(
&self,
encryption_key_id: i64,
) -> Result<EncryptionHandler, ApplicationError> {
let mut db = self.db.lock().await;
let key_path: String = db.process_queue_with_result(|tx| {
tx.query_row(
Expand All @@ -100,16 +111,18 @@ impl UserProfileDbHandler {
.map_err(DatabaseOperationError::SqliteError)
})?;

EncryptionHandler::new_from_path(&PathBuf::from(key_path))?.ok_or_else(
|| {
ApplicationError::EncryptionError(EncryptionError::Other(
Box::new(std::io::Error::new(
std::io::ErrorKind::Other,
"Failed to create encryption handler from key file",
)),
EncryptionHandler::new_from_path(&PathBuf::from(key_path))
.map_err(|e| {
ApplicationError::EncryptionError(EncryptionError::InvalidKey(
e.to_string(),
))
})?
.ok_or_else(|| {
ApplicationError::EncryptionError(EncryptionError::InvalidKey(
"Failed to create encryption handler from key file"
.to_string(),
))
},
)
})
}

pub async fn register_encryption_key(
Expand Down Expand Up @@ -198,4 +211,4 @@ impl UserProfileDbHandler {
})
.map_err(ApplicationError::from)
}
}
}
Loading

0 comments on commit 1bd1414

Please sign in to comment.