tensamin/omikron
Watch
6
0
Fork
You've already forked omikron
 0
Code Issues Pull requests 3 Projects Releases Packages Wiki Activity Actions

Updated Crypto to use MTP-Crypto

Browse source
This commit is contained in:
Alex Emmet 2026-07-03 06:28:08 +02:00
commit 594f13e974
10 changed files with 465 additions and 513 deletions
Download patch file Download diff file Expand all files Collapse all files

37
src/main.rs
View file

@ -11,28 +11,37 @@ use std::path::PathBuf;
use dotenv::dotenv;
use once_cell::sync::Lazy;
pub static WORKING_DIR: Lazy<PathBuf> = Lazy::new(|| {
std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
});
pub static WORKING_DIR: Lazy<PathBuf> =
Lazy::new(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
use rustls::crypto::aws_lc_rs::default_provider;
use base64::engine::general_purpose::STANDARD as BASE64_STD;
use base64::Engine as _;
use mtp::crypto::Keyring;
use crate::{
calls::call_util::garbage_collect_calls,
omega::omega_connection::get_omega_connection,
rho::server::start,
util::{
crypto_helper::{load_public_key, load_secret_key},
logger::startup,
},
util::logger::startup,
};
static PRIVATE_KEY: Lazy<String> = Lazy::new(|| env::var("PRIVATE_KEY").unwrap());
pub fn get_private_key() -> x448::Secret {
load_secret_key(&*PRIVATE_KEY).unwrap()
}
static PUBLIC_KEY: Lazy<String> = Lazy::new(|| env::var("PUBLIC_KEY").unwrap());
pub fn get_public_key() -> x448::PublicKey {
load_public_key(&*PUBLIC_KEY).unwrap()
static KEYRING: Lazy<Keyring> = Lazy::new(|| {
if let Ok(encoded) = env::var("KEYRING") {
let bytes = BASE64_STD.decode(&encoded).expect("Invalid KEYRING base64");
Keyring::from_bytes(&bytes).expect("Invalid KEYRING data")
} else {
let kr = Keyring::generate();
eprintln!(
"Generated KEYRING (save to env): {}",
BASE64_STD.encode(&kr.to_bytes())
);
kr
}
});
pub fn get_keyring() -> &'static Keyring {
&KEYRING
}
#[tokio::main]

27
src/omega/omega_connection.rs
View file

@ -1,9 +1,12 @@
use base64::engine::general_purpose::STANDARD as BASE64_STD;
use base64::Engine as _;
use mtp::crypto::decrypt_with;
use crate::{
data::user::UserStatus,
get_private_key, log, log_cv_in, log_cv_out, log_err, log_in,
get_keyring, log, log_cv_in, log_cv_out, log_err, log_in,
rho::rho_manager::{self, RHO_CONNECTIONS, connection_count},
util::{
crypto_helper::{decrypt_b64, secret_key_to_base64},
file_util::load_file_vec,
logger::PrintType,
},
@ -389,22 +392,18 @@ impl OmegaConnection {
}
async fn handle_challenge(&self, cv: CommunicationValue) -> Result<(), String> {
let challenge = cv
let challenge_b64 = cv
.get_data(DataType::Challenge)
.as_str()
.ok_or("Challenge not found")?;
let server_pub_key = cv
.get_data(DataType::PublicKey)
.as_str()
.ok_or("Public key not found")?;
let decrypted_challenge = decrypt_b64(
&secret_key_to_base64(&get_private_key()),
server_pub_key,
challenge,
)
.map_err(|e| format!("Decryption failed: {:?}", e))?;
let blob = BASE64_STD
.decode(challenge_b64)
.map_err(|e| format!("Base64 decode failed: {}", e))?;
let decrypted = decrypt_with(&blob, get_keyring(), b"challenge")
.map_err(|e| format!("Decryption failed: {:?}", e))?;
let decrypted_challenge =
String::from_utf8(decrypted).map_err(|_| "Decrypted challenge not valid UTF-8")?;
let response_msg = CommunicationValue::new(CommunicationType::ChallengeResponse)
.with_id(cv.get_id())

116
src/rho/connection.rs
View file

@ -5,20 +5,23 @@ use std::{collections::BTreeMap, collections::HashMap, sync::Arc, time::Duration
use tokio::sync::RwLock;
use uuid::Uuid;
use base64::engine::general_purpose::STANDARD as BASE64_STD;
use base64::Engine as _;
use mtp::crypto::{
encrypt_for, EncryptionType, KemPublicKey, PublicKeyBundle, SignaturePqPublicKey,
SignaturePublicKey,
};
use crate::{
anonymous_clients::anonymous_client_connection::AnonymousClientConnection,
calls::call_manager,
get_private_key, get_public_key, log_cv_in, log_cv_out, log_err, log_in, log_out,
get_keyring, log_cv_in, log_cv_out, log_err, log_in, log_out,
omega::omega_connection::get_omega_connection,
rho::{
app_connection::AppConnection, client_connection::ClientConnection,
iota_connection::IotaConnection, rho_connection::RhoConnection, rho_manager,
},
util::{
crypto_helper::{load_public_key, public_key_to_base64},
crypto_util::{DataFormat, SecurePayload},
logger::PrintType,
},
util::logger::PrintType,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@ -146,11 +149,11 @@ impl GeneralConnection {
*self.app_session.write().await = Some(app_session_id);
*self.connection_kind.write().await = Some(ConnectionKind::Phi);
let pub_key = match load_public_key(pub_key_str) {
Some(pk) => pk,
None => return,
let kem_bytes = match BASE64_STD.decode(pub_key_str) {
Ok(b) => b,
Err(_) => return,
};
*self.pub_key.write().await = Some(pub_key.as_bytes().to_vec());
*self.pub_key.write().await = Some(kem_bytes.clone());
let challenge: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
@ -161,18 +164,29 @@ impl GeneralConnection {
*self.challenge.write().await = challenge.clone();
*self.identified.write().await = true;
let encrypted_challenge =
SecurePayload::new(challenge.as_bytes(), DataFormat::Raw, get_private_key())
.unwrap()
.encrypt_x448(pub_key)
.unwrap()
.export(DataFormat::Base64);
let peer_bundle = PublicKeyBundle::new(
KemPublicKey::new(kem_bytes),
SignaturePqPublicKey::new(vec![]),
SignaturePublicKey::new(vec![]),
);
let encrypted_challenge = BASE64_STD.encode(
&encrypt_for(
EncryptionType::MlKemChaCha20Poly1305,
&peer_bundle,
challenge.as_bytes(),
b"challenge",
)
.unwrap(),
);
let our_pk = BASE64_STD.encode(
get_keyring().public_key_bundle().kem_public_key.as_bytes(),
);
let response = CommunicationValue::new(CommunicationType::AppChallenge)
.with_id(cv.get_id())
.add_typed_default(
DataType::PublicKey,
DataValue::Str(public_key_to_base64(&get_public_key())),
DataValue::Str(our_pk),
)
.add_typed_default(DataType::Challenge, DataValue::Str(encrypted_challenge));
@ -214,20 +228,20 @@ impl GeneralConnection {
.as_str()
.unwrap_or("");
let pub_key = match load_public_key(base64_pub) {
Some(pk) => pk,
None => {
let kem_bytes = match BASE64_STD.decode(base64_pub) {
Ok(b) => b,
Err(_) => {
log_err!(
*iota_id as i64,
PrintType::Iota,
"Failed to load public key for iota_id={}",
"Failed to decode public key for iota_id={}",
iota_id
);
return;
}
};
*self.pub_key.write().await = Some(pub_key.as_bytes().to_vec());
*self.pub_key.write().await = Some(kem_bytes.clone());
let challenge: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
@ -238,18 +252,29 @@ impl GeneralConnection {
*self.challenge.write().await = challenge.clone();
*self.identified.write().await = true;
let encrypted_challenge =
SecurePayload::new(challenge.as_bytes(), DataFormat::Raw, get_private_key())
.unwrap()
.encrypt_x448(pub_key)
.unwrap()
.export(DataFormat::Base64);
let peer_bundle = PublicKeyBundle::new(
KemPublicKey::new(kem_bytes),
SignaturePqPublicKey::new(vec![]),
SignaturePublicKey::new(vec![]),
);
let encrypted_challenge = BASE64_STD.encode(
&encrypt_for(
EncryptionType::MlKemChaCha20Poly1305,
&peer_bundle,
challenge.as_bytes(),
b"challenge",
)
.unwrap(),
);
let our_pk = BASE64_STD.encode(
get_keyring().public_key_bundle().kem_public_key.as_bytes(),
);
let response = CommunicationValue::new(CommunicationType::Challenge)
.with_id(cv.get_id())
.add_typed_default(
DataType::PublicKey,
DataValue::Str(public_key_to_base64(&get_public_key())),
DataValue::Str(our_pk),
)
.add_typed_default(DataType::Challenge, DataValue::Str(encrypted_challenge));
@ -306,14 +331,14 @@ impl GeneralConnection {
.to_string();
}
let pub_key = match load_public_key(&base64_pub) {
Some(pk) => pk,
None => {
let kem_bytes = match BASE64_STD.decode(&base64_pub) {
Ok(b) => b,
Err(_) => {
return;
}
};
*self.pub_key.write().await = Some(pub_key.as_bytes().to_vec());
*self.pub_key.write().await = Some(kem_bytes.clone());
let challenge: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
@ -324,12 +349,20 @@ impl GeneralConnection {
*self.challenge.write().await = challenge.clone();
*self.identified.write().await = true;
let encrypted_challenge =
SecurePayload::new(challenge.as_bytes(), DataFormat::Raw, get_private_key())
.unwrap()
.encrypt_x448(pub_key)
.unwrap()
.export(DataFormat::Base64);
let peer_bundle = PublicKeyBundle::new(
KemPublicKey::new(kem_bytes),
SignaturePqPublicKey::new(vec![]),
SignaturePublicKey::new(vec![]),
);
let encrypted_challenge = BASE64_STD.encode(
&encrypt_for(
EncryptionType::MlKemChaCha20Poly1305,
&peer_bundle,
challenge.as_bytes(),
b"challenge",
)
.unwrap(),
);
let challenge_type = if cv.is_type(CommunicationType::AppIdentification) {
CommunicationType::AppChallenge
@ -337,12 +370,15 @@ impl GeneralConnection {
CommunicationType::Challenge
};
let our_pk = BASE64_STD.encode(
get_keyring().public_key_bundle().kem_public_key.as_bytes(),
);
let response = CommunicationValue::new(challenge_type)
.with_id(cv.get_id())
.with_receiver(*self.session_id.read().await)
.add_typed_default(
DataType::PublicKey,
DataValue::Str(public_key_to_base64(&get_public_key())),
DataValue::Str(our_pk),
)
.add_typed_default(DataType::Challenge, DataValue::Str(encrypted_challenge));

11
src/rho/iota_connection.rs
View file

@ -20,7 +20,7 @@ use std::collections::BTreeMap;
use std::{collections::HashMap, sync::Arc, time::Duration};
use tokio::sync::RwLock;
use tokio::sync::mpsc;
use x448::PublicKey;
use mtp::crypto::KemPublicKey;
use super::{rho_connection::RhoConnection, rho_manager};
use crate::omega::omega_connection::OmegaConnection;
@ -79,12 +79,9 @@ impl IotaConnection {
}
#[allow(dead_code)]
pub async fn get_public_key(&self) -> Option<PublicKey> {
if let Some(public_key) = self.pub_key.read().await.clone() {
PublicKey::from_bytes(&public_key)
} else {
None
}
pub async fn get_public_key(&self) -> Option<KemPublicKey> {
let guard = self.pub_key.read().await;
guard.as_ref().map(|bytes| KemPublicKey::new(bytes.clone()))
}
/// Get the user IDs

149
src/util/crypto_helper.rs
View file

@ -1,149 +0,0 @@
use aes_gcm::{
Aes256Gcm, Nonce,
aead::{Aead, KeyInit, OsRng},
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use rand_core::RngCore;
use sha2::{Digest, Sha256};
use x448::{PublicKey, Secret, SharedSecret};
/// Errors for crypto operations
#[allow(dead_code)]
#[derive(Debug)]
pub enum CryptoError {
Base64Decode(base64::DecodeError),
InvalidKey,
AgreementError,
EncryptionError(aes_gcm::Error),
DecryptionError(aes_gcm::Error),
}
impl From<base64::DecodeError> for CryptoError {
fn from(err: base64::DecodeError) -> Self {
CryptoError::Base64Decode(err)
}
}
#[allow(dead_code)]
pub struct KeyPair {
pub secret: Secret,
pub public: PublicKey,
}
#[allow(dead_code)]
pub fn generate_keypair() -> KeyPair {
let mut buf = [0u8; 56];
let mut rng = OsRng;
rng.fill_bytes(&mut buf);
let secret = Secret::from_bytes(&buf).unwrap();
let public = PublicKey::from(&secret);
KeyPair { secret, public }
}
pub fn public_key_to_base64(pubkey: &PublicKey) -> String {
STANDARD.encode(pubkey.as_bytes().as_ref())
}
pub fn secret_key_to_base64(secret: &Secret) -> String {
STANDARD.encode(secret.as_bytes().as_ref())
}
pub fn load_public_key(base64_pub: &str) -> Option<PublicKey> {
let bytes = STANDARD.decode(base64_pub).unwrap();
PublicKey::from_bytes(&bytes)
}
pub fn load_secret_key(base64_secret: &str) -> Option<Secret> {
let bytes = STANDARD.decode(base64_secret).unwrap();
Secret::from_bytes(&bytes)
}
fn derive_aes_key(shared: &SharedSecret) -> [u8; 32] {
let mut hasher = Sha256::new();
hasher.update(shared.as_bytes());
let result = hasher.finalize();
let mut key = [0u8; 32];
key.copy_from_slice(&result[..32]);
key
}
#[allow(dead_code)]
pub fn encrypt_b64(
base64_secret: &str,
base64_peer_pub: &str,
plaintext: &str,
) -> Result<String, CryptoError> {
let secret = load_secret_key(base64_secret).unwrap();
let peer_pub = load_public_key(base64_peer_pub).unwrap();
encrypt(secret, peer_pub, plaintext)
}
pub fn encrypt(
secret: Secret,
peer_pub: PublicKey,
plaintext: &str,
) -> Result<String, CryptoError> {
let shared = secret
.to_diffie_hellman(&peer_pub)
.ok_or(CryptoError::AgreementError)?;
let key_bytes = derive_aes_key(&shared);
let cipher = Aes256Gcm::new_from_slice(&key_bytes).expect("Key length should be correct");
let mut nonce_bytes = [0u8; 12];
OsRng.fill_bytes(&mut nonce_bytes);
let nonce = Nonce::from_slice(&nonce_bytes);
let ciphertext = cipher
.encrypt(nonce, plaintext.as_bytes())
.map_err(CryptoError::EncryptionError)?;
// prefix nonce to ciphertext
let mut out = Vec::with_capacity(nonce_bytes.len() + ciphertext.len());
out.extend_from_slice(&nonce_bytes);
out.extend_from_slice(&ciphertext);
Ok(STANDARD.encode(&out))
}
pub fn decrypt_b64(
base64_secret: &str,
base64_peer_pub: &str,
encrypted_base64: &str,
) -> Result<String, CryptoError> {
let secret = load_secret_key(base64_secret).unwrap();
let peer_pub = load_public_key(base64_peer_pub).unwrap();
decrypt(secret, peer_pub, encrypted_base64)
}
pub fn decrypt(
secret: Secret,
peer_pub: PublicKey,
encrypted_base64: &str,
) -> Result<String, CryptoError> {
let shared = secret
.to_diffie_hellman(&peer_pub)
.ok_or(CryptoError::AgreementError)?;
let key_bytes = derive_aes_key(&shared);
let cipher = Aes256Gcm::new_from_slice(&key_bytes).expect("Key length should be correct");
let encrypted = STANDARD.decode(encrypted_base64)?;
if encrypted.len() < 12 {
return Err(CryptoError::DecryptionError(aes_gcm::Error));
}
let nonce_bytes = &encrypted[..12];
let ciphertext = &encrypted[12..];
let nonce = Nonce::from_slice(nonce_bytes);
let plaintext_bytes = cipher
.decrypt(nonce, ciphertext)
.map_err(CryptoError::DecryptionError)?;
let plaintext = String::from_utf8(plaintext_bytes)
.map_err(|_| CryptoError::DecryptionError(aes_gcm::Error))?;
Ok(plaintext)
}
#[allow(dead_code)]
pub fn hash_it(input: &str) -> Vec<u8> {
let mut hasher = Sha256::new();
hasher.update(input.as_bytes());
hasher.finalize().to_vec()
}
#[allow(dead_code)]
pub fn hex_hash(input: &str) -> String {
let digest = hash_it(input);
digest.iter().map(|b| format!("{:02x}", b)).collect()
}

178
src/util/crypto_util.rs
View file

@ -1,178 +0,0 @@
use aes_gcm::{
Aes256Gcm, Nonce,
aead::{Aead, KeyInit, Payload},
};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STD};
use hkdf::Hkdf;
type HkdfSha256 = sha2::Sha256;
use sha2::{Digest, Sha256 as HashSha256};
use x448::{PublicKey, Secret};
#[derive(Debug)]
#[allow(dead_code)]
pub enum SecurePayloadError {
InvalidBase64,
InvalidHex,
EncryptionError,
DecryptionError,
InvalidKeyLength,
}
#[derive(Clone, Copy, Debug)]
#[allow(dead_code)]
pub enum DataFormat {
Raw,
Base64,
Hex,
}
pub struct SecurePayload {
inner_data: Vec<u8>,
private_key: Secret,
}
impl Clone for SecurePayload {
fn clone(&self) -> Self {
Self {
inner_data: self.inner_data.clone(),
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
}
}
}
#[allow(dead_code)]
impl SecurePayload {
pub fn new<S, T: AsRef<[u8]>>(
data: T,
format: DataFormat,
private_key: S,
) -> Result<Self, SecurePayloadError>
where
S: Into<Secret>,
{
let raw_data = match format {
DataFormat::Raw => data.as_ref().to_vec(),
DataFormat::Base64 => BASE64_STD
.decode(data.as_ref())
.map_err(|_| SecurePayloadError::InvalidBase64)?,
DataFormat::Hex => {
hex::decode(data.as_ref()).map_err(|_| SecurePayloadError::InvalidHex)?
}
};
Ok(Self {
inner_data: raw_data,
private_key: private_key.into(),
})
}
pub fn get_public_key(&self) -> [u8; 56] {
*PublicKey::from(&self.private_key).as_bytes()
}
pub fn export(&self, format: DataFormat) -> String {
match format.into() {
DataFormat::Raw => String::from_utf8_lossy(&self.inner_data).to_string(),
DataFormat::Base64 => BASE64_STD.encode(&self.inner_data),
DataFormat::Hex => hex::encode(&self.inner_data),
}
}
pub fn get_bytes(&self) -> &[u8] {
&self.inner_data
}
pub fn get_hash(&self, format: DataFormat) -> String {
let mut hasher = HashSha256::new();
hasher.update(&self.inner_data);
let result = hasher.finalize();
match format {
DataFormat::Raw => String::from_utf8_lossy(&result).to_string(),
DataFormat::Base64 => BASE64_STD.encode(result),
DataFormat::Hex => hex::encode(result),
}
}
pub fn encrypt_x448<S>(&self, public_key: S) -> Result<SecurePayload, SecurePayloadError>
where
S: Into<PublicKey>,
{
let peer_pub = public_key.into();
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
let hkdf = Hkdf::<HkdfSha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44];
hkdf.expand(b"x448-aes-gcm-no-overhead", &mut okm)
.map_err(|_| SecurePayloadError::EncryptionError)?;
let key = &okm[..32];
let nonce_bytes = &okm[32..];
let cipher = Aes256Gcm::new(key.into());
let nonce = Nonce::from_slice(nonce_bytes);
let ciphertext = cipher
.encrypt(
nonce,
Payload {
msg: &self.inner_data,
aad: &[],
},
)
.map_err(|_| SecurePayloadError::EncryptionError)?;
Ok(SecurePayload {
inner_data: ciphertext,
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
})
}
pub fn decrypt_to_format(
&self,
peer_public_key_bytes: &[u8; 56],
output_format: DataFormat,
) -> Result<String, SecurePayloadError> {
let decrypted_instance =
self.decrypt_x448(PublicKey::from_bytes(peer_public_key_bytes).unwrap())?;
Ok(decrypted_instance.export(output_format))
}
pub fn decrypt_x448<S>(
&self,
peer_public_key_bytes: S,
) -> Result<SecurePayload, SecurePayloadError>
where
S: Into<PublicKey>,
{
let peer_pub = peer_public_key_bytes.into();
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
let hkdf = Hkdf::<HkdfSha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44];
hkdf.expand(b"x448-aes-gcm-no-overhead", &mut okm)
.map_err(|_| SecurePayloadError::DecryptionError)?;
let key = &okm[..32];
let nonce_bytes = &okm[32..];
let cipher = Aes256Gcm::new(key.into());
let nonce = Nonce::from_slice(nonce_bytes);
let plaintext = cipher
.decrypt(
nonce,
Payload {
msg: &self.inner_data,
aad: &[],
},
)
.map_err(|_| SecurePayloadError::DecryptionError)?;
Ok(SecurePayload {
inner_data: plaintext,
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
})
}
}

2
src/util/mod.rs
View file

@ -1,4 +1,2 @@
pub mod crypto_helper;
pub mod crypto_util;
pub mod file_util;
pub mod logger;