Skip to content

Commit

Permalink
Add support for generic public key caching
Browse files Browse the repository at this point in the history
The cache must be provided by the user of this crate by implementing the
`PublicKeyCache` trait.
  • Loading branch information
dbrgn committed Jul 21, 2024
1 parent f126b1e commit 96d3451
Show file tree
Hide file tree
Showing 5 changed files with 104 additions and 10 deletions.
47 changes: 40 additions & 7 deletions examples/lookup_pubkey.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
use docopt::Docopt;
use threema_gateway::ApiBuilder;
use threema_gateway::{ApiBuilder, PublicKeyCache};

const USAGE: &str = "
Usage: lookup_pubkey [options] <our_id> <secret> <their_id>
Usage: lookup_pubkey [--with-cache] <our_id> <secret> <their_id>
Options:
-h, --help Show this help
--with-cache Simulate a cache
";

#[tokio::main(flavor = "current_thread")]
Expand All @@ -18,14 +18,47 @@ async fn main() {
let our_id = args.get_str("<our_id>");
let their_id = args.get_str("<their_id>");
let secret = args.get_str("<secret>");
let simulate_cache = args.get_bool("--with-cache");

// Fetch recipient public key
let api = ApiBuilder::new(our_id, secret).into_simple();
let recipient_key = api.lookup_pubkey(their_id).await;
let pubkey = if simulate_cache {
let cache = SimulatedCache;
api.lookup_pubkey_with_cache(their_id, &cache)
.await
.unwrap_or_else(|e| {
println!("Could not fetch public key: {}", e);
std::process::exit(1);
})
} else {
api.lookup_pubkey(their_id).await.unwrap_or_else(|e| {
println!("Could not fetch and cache public key: {}", e);
std::process::exit(1);
})
};

// Show result
match recipient_key {
Ok(key) => println!("Public key for {} is {}.", their_id, key.to_hex_string()),
Err(e) => println!("Could not fetch public key: {}", e),
println!("Public key for {} is {}.", their_id, pubkey.to_hex_string());
}

struct SimulatedCache;

impl PublicKeyCache for SimulatedCache {
type Error = std::io::Error;

async fn store(
&self,
identity: &str,
_key: &threema_gateway::RecipientKey,
) -> Result<(), Self::Error> {
println!("[cache] Storing public key for identity {identity}");
Ok(())
}

async fn load(
&self,
_identity: &str,
) -> Result<Option<threema_gateway::RecipientKey>, Self::Error> {
unimplemented!("Not implemented in this example")
}
}
35 changes: 32 additions & 3 deletions src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,12 @@ use data_encoding::HEXLOWER_PERMISSIVE;
use reqwest::Client;

use crate::{
cache::PublicKeyCache,
connection::{blob_download, blob_upload, send_e2e, send_simple, Recipient},
crypto::{
encrypt, encrypt_file_msg, encrypt_image_msg, encrypt_raw, EncryptedMessage, RecipientKey,
},
errors::{ApiBuilderError, ApiError, CryptoError},
errors::{ApiBuilderError, ApiError, ApiOrCacheError, CryptoError},
lookup::{
lookup_capabilities, lookup_credits, lookup_id, lookup_pubkey, Capabilities,
LookupCriterion,
Expand Down Expand Up @@ -42,8 +43,9 @@ macro_rules! impl_common_functionality {
/// and therefore you can also look up the key associated with a given ID from
/// the server.
///
/// It is strongly recommended that you cache the public keys to avoid querying
/// the API for each message.
/// *Note:* It is strongly recommended that you cache the public keys to avoid
/// querying the API for each message. To simplify this, the
/// `lookup_pubkey_with_cache` method can be used instead.
pub async fn lookup_pubkey(&self, id: &str) -> Result<RecipientKey, ApiError> {
lookup_pubkey(
&self.client,
Expand All @@ -55,6 +57,33 @@ macro_rules! impl_common_functionality {
.await
}

/// Fetch the recipient public key for the specified Threema ID and store it
/// in the [`PublicKeyCache`].
///
/// For the end-to-end encrypted mode, you need the public key of the recipient
/// in order to encrypt a message. While it's best to obtain this directly from
/// the recipient (extract it from the QR code), this may not be convenient,
/// and therefore you can also look up the key associated with a given ID from
/// the server.
pub async fn lookup_pubkey_with_cache<C>(
&self,
id: &str,
public_key_cache: &C,
) -> Result<RecipientKey, ApiOrCacheError<C::Error>>
where
C: PublicKeyCache,
{
let pubkey = self
.lookup_pubkey(id)
.await
.map_err(ApiOrCacheError::ApiError)?;
public_key_cache
.store(id, &pubkey)
.await
.map_err(ApiOrCacheError::CacheError)?;
Ok(pubkey)
}

/// Look up a Threema ID in the directory.
///
/// An ID can be looked up either by a phone number or an e-mail
Expand Down
22 changes: 22 additions & 0 deletions src/cache.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
use std::future::Future;

use crate::crypto::RecipientKey;

/// A cache for Threema public keys
pub trait PublicKeyCache {
/// Error returned if cache operations fail
type Error: std::error::Error;

/// Store a public key for `identity` in the cache
fn store(
&self,
identity: &str,
key: &RecipientKey,
) -> impl Future<Output = Result<(), Self::Error>>;

/// Retrieve a public key for `identity` from the cache
fn load(
&self,
identity: &str,
) -> impl Future<Output = Result<Option<RecipientKey>, Self::Error>>;
}
8 changes: 8 additions & 0 deletions src/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,14 @@ impl From<ReqwestError> for ApiError {
}
}

#[derive(Debug, Error)]
pub enum ApiOrCacheError<C: std::error::Error> {
#[error("api error: {0}")]
ApiError(ApiError),
#[error("cache error: {0}")]
CacheError(C),
}

/// Crypto related errors.
#[derive(Debug, PartialEq, Clone, Error)]
pub enum CryptoError {
Expand Down
2 changes: 2 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@
extern crate log;

mod api;
mod cache;
mod connection;
mod crypto;
pub mod errors;
Expand All @@ -88,6 +89,7 @@ pub use crypto_secretbox::Nonce;

pub use crate::{
api::{ApiBuilder, E2eApi, SimpleApi},
cache::PublicKeyCache,
connection::Recipient,
crypto::{
decrypt_file_data, encrypt, encrypt_file_data, encrypt_raw, EncryptedFileData,
Expand Down

0 comments on commit 96d3451

Please sign in to comment.