MindWorkAI/AI-Studio
MindWorkAI/AI-Studio
1
0
Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2025-11-04 03:20:21 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Refactored encryption-related code

Browse Source
This commit is contained in:
Thorsten Sommer 2024-11-04 20:42:12 +01:00
parent f1104c5e09
commit 74522dc22a
Signed by: tsommer
GPG Key ID: 371BBA77A02C0108
3 changed files with 170 additions and 157 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

165
runtime/src/encryption.rs Normal file
View File

@ -0,0 +1,165 @@
use std::fmt;
use std::time::Instant;
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use hmac::Hmac;
use log::info;
use once_cell::sync::Lazy;
use pbkdf2::pbkdf2;
use rand::{RngCore, SeedableRng};
use rocket::{data, Data, Request};
use rocket::data::ToByteUnit;
use rocket::http::Status;
use rocket::serde::{Deserialize, Serialize};
use sha2::Sha512;
use tokio::io::AsyncReadExt;
type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;
type DataOutcome<'r, T> = data::Outcome<'r, T>;
pub static ENCRYPTION: Lazy<Encryption> = Lazy::new(|| {
//
// Generate a secret key & salt for the AES encryption for the IPC channel:
//
let mut secret_key = [0u8; 512]; // 512 bytes = 4096 bits
let mut secret_key_salt = [0u8; 16]; // 16 bytes = 128 bits
// We use a cryptographically secure pseudo-random number generator
// to generate the secret password & salt. ChaCha20Rng is the algorithm
// of our choice:
let mut rng = rand_chacha::ChaChaRng::from_entropy();
// Fill the secret key & salt with random bytes:
rng.fill_bytes(&mut secret_key);
rng.fill_bytes(&mut secret_key_salt);
Encryption::new(&secret_key, &secret_key_salt).unwrap()
});
pub struct Encryption {
key: [u8; 32],
iv: [u8; 16],
pub secret_password: [u8; 512],
pub secret_key_salt: [u8; 16],
}
impl Encryption {
// The number of iterations to derive the key and IV from the password. For a password
// manager where the user has to enter their primary password, 100 iterations would be
// too few and insecure. Here, the use case is different: We generate a 512-byte long
// and cryptographically secure password at every start. This password already contains
// enough entropy. In our case, we need key and IV primarily because AES, with the
// algorithms we chose, requires a fixed key length, and our password is too long.
const ITERATIONS: u32 = 100;
pub fn new(secret_password: &[u8], secret_key_salt: &[u8]) -> Result<Self, String> {
if secret_password.len() != 512 {
return Err("The secret password must be 512 bytes long.".to_string());
}
if secret_key_salt.len() != 16 {
return Err("The salt must be 16 bytes long.".to_string());
}
info!(Source = "Encryption"; "Initializing encryption...");
let mut encryption = Encryption {
key: [0u8; 32],
iv: [0u8; 16],
secret_password: [0u8; 512],
secret_key_salt: [0u8; 16],
};
encryption.secret_password.copy_from_slice(secret_password);
encryption.secret_key_salt.copy_from_slice(secret_key_salt);
let start = Instant::now();
let mut key_iv = [0u8; 48];
pbkdf2::<Hmac<Sha512>>(secret_password, secret_key_salt, Self::ITERATIONS, &mut key_iv).map_err(|e| format!("Error while generating key and IV: {e}"))?;
encryption.key.copy_from_slice(&key_iv[0..32]);
encryption.iv.copy_from_slice(&key_iv[32..48]);
let duration = start.elapsed();
let duration = duration.as_millis();
info!(Source = "Encryption"; "Encryption initialized in {duration} milliseconds.", );
Ok(encryption)
}
pub fn encrypt(&self, data: &str) -> Result<EncryptedText, String> {
let cipher = Aes256CbcEnc::new(&self.key.into(), &self.iv.into());
let encrypted = cipher.encrypt_padded_vec_mut::<Pkcs7>(data.as_bytes());
let mut result = BASE64_STANDARD.encode(self.secret_key_salt);
result.push_str(&BASE64_STANDARD.encode(&encrypted));
Ok(EncryptedText::new(result))
}
pub fn decrypt(&self, encrypted_data: &EncryptedText) -> Result<String, String> {
let decoded = BASE64_STANDARD.decode(encrypted_data.get_encrypted()).map_err(|e| format!("Error decoding base64: {e}"))?;
if decoded.len() < 16 {
return Err("Encrypted data is too short.".to_string());
}
let (salt, encrypted) = decoded.split_at(16);
if salt != self.secret_key_salt {
return Err("The salt bytes do not match. The data is corrupted or tampered.".to_string());
}
let cipher = Aes256CbcDec::new(&self.key.into(), &self.iv.into());
let decrypted = cipher.decrypt_padded_vec_mut::<Pkcs7>(encrypted).map_err(|e| format!("Error decrypting data: {e}"))?;
String::from_utf8(decrypted).map_err(|e| format!("Error converting decrypted data to string: {}", e))
}
}
#[derive(Clone, Serialize, Deserialize)]
pub struct EncryptedText(String);
impl EncryptedText {
pub fn new(encrypted_data: String) -> Self {
EncryptedText(encrypted_data)
}
pub fn get_encrypted(&self) -> &str {
&self.0
}
}
impl fmt::Debug for EncryptedText {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "EncryptedText(**********)")
}
}
impl fmt::Display for EncryptedText {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "**********")
}
}
// Use Case: When we receive encrypted text from the client as body (e.g., in a POST request).
// We must interpret the body as EncryptedText.
#[rocket::async_trait]
impl<'r> data::FromData<'r> for EncryptedText {
type Error = String;
async fn from_data(req: &'r Request<'_>, data: Data<'r>) -> DataOutcome<'r, Self> {
let content_type = req.content_type();
if content_type.map_or(true, |ct| !ct.is_text()) {
return DataOutcome::Forward((data, Status::Ok));
}
let mut stream = data.open(2.mebibytes());
let mut body = String::new();
if let Err(e) = stream.read_to_string(&mut body).await {
return DataOutcome::Error((Status::InternalServerError, format!("Failed to read data: {}", e)));
}
DataOutcome::Success(EncryptedText(body))
}
}

1
runtime/src/lib.rs Normal file
View File

@ -0,0 +1 @@
pub mod encryption;

161
runtime/src/main.rs
View File

@ -5,16 +5,14 @@ extern crate rocket;
extern crate core;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fmt;
use std::net::TcpListener;
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
use std::time::Duration;
use once_cell::sync::Lazy;
use arboard::Clipboard;
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use keyring::Entry;
use serde::{Deserialize, Serialize};
use tauri::{Manager, Url, Window};
@ -22,29 +20,21 @@ use tauri::api::process::{Command, CommandChild, CommandEvent};
use tokio::time;
use flexi_logger::{DeferredNow, Duplicate, FileSpec, Logger};
use flexi_logger::writers::FileLogWriter;
use hmac::Hmac;
use keyring::error::Error::NoEntry;
use log::{debug, error, info, kv, warn};
use log::kv::{Key, Value, VisitSource};
use pbkdf2::pbkdf2;
use rand::{RngCore, SeedableRng};
use rcgen::generate_simple_self_signed;
use rocket::figment::Figment;
use rocket::{data, get, post, routes, Data, Request};
use rocket::{get, post, routes, Request};
use rocket::config::{Shutdown};
use rocket::data::{ToByteUnit};
use rocket::http::Status;
use rocket::request::{FromRequest};
use rocket::serde::json::Json;
use sha2::{Sha256, Sha512, Digest};
use sha2::{Sha256, Digest};
use tauri::updater::UpdateResponse;
use tokio::io::AsyncReadExt;
type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;
type DataOutcome<'r, T> = data::Outcome<'r, T>;
use mindwork_ai_studio::encryption::{EncryptedText, ENCRYPTION};
type RequestOutcome<R, T> = rocket::request::Outcome<R, T>;
@ -75,25 +65,6 @@ static MAIN_WINDOW: Lazy<Mutex<Option<Window>>> = Lazy::new(|| Mutex::new(None))
// The update response coming from the Tauri updater.
static CHECK_UPDATE_RESPONSE: Lazy<Mutex<Option<UpdateResponse<tauri::Wry>>>> = Lazy::new(|| Mutex::new(None));
static ENCRYPTION: Lazy<Encryption> = Lazy::new(|| {
//
// Generate a secret key & salt for the AES encryption for the IPC channel:
//
let mut secret_key = [0u8; 512]; // 512 bytes = 4096 bits
let mut secret_key_salt = [0u8; 16]; // 16 bytes = 128 bits
// We use a cryptographically secure pseudo-random number generator
// to generate the secret password & salt. ChaCha20Rng is the algorithm
// of our choice:
let mut rng = rand_chacha::ChaChaRng::from_entropy();
// Fill the secret key & salt with random bytes:
rng.fill_bytes(&mut secret_key);
rng.fill_bytes(&mut secret_key_salt);
Encryption::new(&secret_key, &secret_key_salt).unwrap()
});
static API_TOKEN: Lazy<APIToken> = Lazy::new(|| {
let mut token = [0u8; 32];
let mut rng = rand_chacha::ChaChaRng::from_entropy();
@ -590,130 +561,6 @@ pub fn file_logger_format(
write!(w, "{}", &record.args())
}
pub struct Encryption {
key: [u8; 32],
iv: [u8; 16],
secret_password: [u8; 512],
secret_key_salt: [u8; 16],
}
impl Encryption {
// The number of iterations to derive the key and IV from the password. For a password
// manager where the user has to enter their primary password, 100 iterations would be
// too few and insecure. Here, the use case is different: We generate a 512-byte long
// and cryptographically secure password at every start. This password already contains
// enough entropy. In our case, we need key and IV primarily because AES, with the
// algorithms we chose, requires a fixed key length, and our password is too long.
const ITERATIONS: u32 = 100;
pub fn new(secret_password: &[u8], secret_key_salt: &[u8]) -> Result<Self, String> {
if secret_password.len() != 512 {
return Err("The secret password must be 512 bytes long.".to_string());
}
if secret_key_salt.len() != 16 {
return Err("The salt must be 16 bytes long.".to_string());
}
info!(Source = "Encryption"; "Initializing encryption...");
let mut encryption = Encryption {
key: [0u8; 32],
iv: [0u8; 16],
secret_password: [0u8; 512],
secret_key_salt: [0u8; 16],
};
encryption.secret_password.copy_from_slice(secret_password);
encryption.secret_key_salt.copy_from_slice(secret_key_salt);
let start = Instant::now();
let mut key_iv = [0u8; 48];
pbkdf2::<Hmac<Sha512>>(secret_password, secret_key_salt, Self::ITERATIONS, &mut key_iv).map_err(|e| format!("Error while generating key and IV: {e}"))?;
encryption.key.copy_from_slice(&key_iv[0..32]);
encryption.iv.copy_from_slice(&key_iv[32..48]);
let duration = start.elapsed();
let duration = duration.as_millis();
info!(Source = "Encryption"; "Encryption initialized in {duration} milliseconds.", );
Ok(encryption)
}
pub fn encrypt(&self, data: &str) -> Result<EncryptedText, String> {
let cipher = Aes256CbcEnc::new(&self.key.into(), &self.iv.into());
let encrypted = cipher.encrypt_padded_vec_mut::<Pkcs7>(data.as_bytes());
let mut result = BASE64_STANDARD.encode(self.secret_key_salt);
result.push_str(&BASE64_STANDARD.encode(&encrypted));
Ok(EncryptedText::new(result))
}
pub fn decrypt(&self, encrypted_data: &EncryptedText) -> Result<String, String> {
let decoded = BASE64_STANDARD.decode(encrypted_data.get_encrypted()).map_err(|e| format!("Error decoding base64: {e}"))?;
if decoded.len() < 16 {
return Err("Encrypted data is too short.".to_string());
}
let (salt, encrypted) = decoded.split_at(16);
if salt != self.secret_key_salt {
return Err("The salt bytes do not match. The data is corrupted or tampered.".to_string());
}
let cipher = Aes256CbcDec::new(&self.key.into(), &self.iv.into());
let decrypted = cipher.decrypt_padded_vec_mut::<Pkcs7>(encrypted).map_err(|e| format!("Error decrypting data: {e}"))?;
String::from_utf8(decrypted).map_err(|e| format!("Error converting decrypted data to string: {}", e))
}
}
#[derive(Clone, Serialize, Deserialize)]
pub struct EncryptedText(String);
impl EncryptedText {
pub fn new(encrypted_data: String) -> Self {
EncryptedText(encrypted_data)
}
pub fn get_encrypted(&self) -> &str {
&self.0
}
}
impl fmt::Debug for EncryptedText {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "EncryptedText(**********)")
}
}
impl fmt::Display for EncryptedText {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "**********")
}
}
// Use Case: When we receive encrypted text from the client as body (e.g., in a POST request).
// We must interpret the body as EncryptedText.
#[rocket::async_trait]
impl<'r> data::FromData<'r> for EncryptedText {
type Error = String;
async fn from_data(req: &'r Request<'_>, data: Data<'r>) -> DataOutcome<'r, Self> {
let content_type = req.content_type();
if content_type.map_or(true, |ct| !ct.is_text()) {
return DataOutcome::Forward((data, Status::Ok));
}
let mut stream = data.open(2.mebibytes());
let mut body = String::new();
if let Err(e) = stream.read_to_string(&mut body).await {
return DataOutcome::Error((Status::InternalServerError, format!("Failed to read data: {}", e)));
}
DataOutcome::Success(EncryptedText(body))
}
}
#[get("/system/dotnet/port")]
fn dotnet_port(_token: APIToken) -> String {
let dotnet_server_port = *DOTNET_SERVER_PORT;