frxd/tests/encryption.rs

mod common;

use std::fs;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};

use common::{
    ask, client, collection, config_for, poll_messages, publish, query_envelope, register,
    spawn_relay,
};
use frxd::config::Config;
use frxd::crypto::{self, Keypair, now_ts};
use frxd::index::LocalIndex;
use frxd::message::{EXPOSURE_FULL, TYPE_RESPONSE};
use frxd::node::Node;
use frxd::registry::{self, KeyEntry, RegistryDoc, RegistryMember};
use serde_json::Value;

fn member_with_enc(id: &str, ed: &Keypair, enc_public: &str) -> RegistryMember {
    RegistryMember {
        id: id.to_string(),
        keys: vec![KeyEntry {
            key: ed.public_hex(),
            not_before: 0,
            not_after: None,
        }],
        enc_key: Some(enc_public.to_string()),
    }
}

fn save_registry(dir: &Path, ma: &Keypair, members: Vec<RegistryMember>) -> PathBuf {
    let path = dir.join("registry.json");
    registry::save_registry(
        &path,
        &registry::sign_registry(
            RegistryDoc {
                version: 1,
                issued_at: now_ts(),
                ma_key: String::new(),
                zone: "frx.invalid".to_string(),
                members,
                relays: Vec::new(),
            },
            ma,
        ),
    )
    .unwrap();
    path
}

fn node_config(
    dir: &Path,
    name: &str,
    id: &str,
    relay_url: &str,
    registry_path: &Path,
    ma_key: &str,
    enc_secret: &str,
) -> Config {
    let mut config = config_for(dir, name, relay_url);
    config.node.id = Some(id.to_string());
    config.node.registry = Some(registry_path.display().to_string());
    config.node.ma_key = Some(ma_key.to_string());
    config.node.dev_bootstrap = false;
    config.save_enc_key(enc_secret).unwrap();
    config
}

async fn response_envelopes(
    http: &reqwest::Client,
    relay_url: &str,
    asker: &Keypair,
    text: &str,
    timeout_ms: u64,
) -> Vec<Value> {
    register(http, relay_url, asker).await;
    let envelope = query_envelope(asker, "alice.frx.example", text, 5);
    assert!(
        publish(http, relay_url, &envelope)
            .await
            .status()
            .is_success()
    );
    let start = Instant::now();
    let deadline = Duration::from_millis(timeout_ms);
    let mut responses = Vec::new();
    while start.elapsed() < deadline {
        let messages = poll_messages(http, relay_url, asker, 200).await;
        responses.extend(
            messages
                .iter()
                .filter(|m| m.get("type").and_then(Value::as_str) == Some(TYPE_RESPONSE))
                .cloned(),
        );
        if !responses.is_empty() {
            break;
        }
    }
    responses
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn relay_sees_ciphertext_and_recipient_decrypts() {
    let root = tempfile::tempdir().unwrap();
    let relay_url = spawn_relay().await;
    let ma = Keypair::generate();

    let alice = Keypair::generate();
    let (alice_enc_secret, alice_enc_public) = crypto::generate_enc_keypair();
    let (bob_enc_secret, bob_enc_public) = {
        let pair = crypto::generate_enc_keypair();
        (pair.0, pair.1)
    };

    let bob_docs = root.path().join("bob-docs");
    fs::create_dir_all(&bob_docs).unwrap();
    fs::write(bob_docs.join("doc.txt"), "confidential rust document").unwrap();
    let bob_config = node_config(
        &root.path().join("bob"),
        "bob",
        "bob.frx.example",
        &relay_url,
        &root.path().join("registry.json"),
        &ma.public_hex(),
        &bob_enc_secret,
    );
    let registry_path = save_registry(
        root.path(),
        &ma,
        vec![
            member_with_enc("alice.frx.example", &alice, &alice_enc_public),
            member_with_enc(
                "bob.frx.example",
                &bob_config.load_key().unwrap(),
                &bob_enc_public,
            ),
        ],
    );
    {
        let index = LocalIndex::open(&bob_config.index_dir()).unwrap();
        index
            .add_collection(&collection("docs", &bob_docs, true, EXPOSURE_FULL))
            .unwrap();
    }
    let _bob = Node::start(bob_config).await.unwrap();
    let _ = &registry_path;

    let http = client();
    let responses = response_envelopes(&http, &relay_url, &alice, "rust", 900).await;
    assert_eq!(responses.len(), 1, "expected one encrypted response");
    let response = &responses[0];
    let raw = serde_json::to_string(response).unwrap();
    assert!(
        !raw.contains("\"results\""),
        "relay-visible payload leaked plaintext: {raw}"
    );
    let enc = response.pointer("/body/enc").unwrap();
    assert_eq!(
        enc.get("alg").and_then(Value::as_str),
        Some(crypto::ENC_ALG)
    );

    let context = crypto::unicast_context("bob.frx.example", &alice_enc_public);
    let plaintext = crypto::decrypt_unicast(&alice_enc_secret, &context, enc).unwrap();
    let body: Value = serde_json::from_slice(&plaintext).unwrap();
    assert_eq!(
        body.get("results").and_then(Value::as_array).map(Vec::len),
        Some(1)
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn node_decrypts_encrypted_responses_end_to_end() {
    let root = tempfile::tempdir().unwrap();
    let relay_url = spawn_relay().await;
    let ma = Keypair::generate();

    let (alice_enc_secret, alice_enc_public) = crypto::generate_enc_keypair();
    let (bob_enc_secret, bob_enc_public) = crypto::generate_enc_keypair();

    let bob_docs = root.path().join("bob-docs");
    fs::create_dir_all(&bob_docs).unwrap();
    fs::write(bob_docs.join("doc.txt"), "confidential rust document").unwrap();
    let alice_config = node_config(
        &root.path().join("alice"),
        "alice",
        "alice.frx.example",
        &relay_url,
        &root.path().join("registry.json"),
        &ma.public_hex(),
        &alice_enc_secret,
    );
    let bob_config = node_config(
        &root.path().join("bob"),
        "bob",
        "bob.frx.example",
        &relay_url,
        &root.path().join("registry.json"),
        &ma.public_hex(),
        &bob_enc_secret,
    );
    save_registry(
        root.path(),
        &ma,
        vec![
            member_with_enc(
                "alice.frx.example",
                &alice_config.load_key().unwrap(),
                &alice_enc_public,
            ),
            member_with_enc(
                "bob.frx.example",
                &bob_config.load_key().unwrap(),
                &bob_enc_public,
            ),
        ],
    );
    {
        let index = LocalIndex::open(&bob_config.index_dir()).unwrap();
        index
            .add_collection(&collection("docs", &bob_docs, true, EXPOSURE_FULL))
            .unwrap();
    }
    let _bob = Node::start(bob_config).await.unwrap();
    let alice = Node::start(alice_config).await.unwrap();

    let (_status, _raw, value) = ask(&client(), &alice.addr.to_string(), "rust", 5).await;
    let responses = value.get("responses").and_then(Value::as_array).unwrap();
    assert_eq!(responses.len(), 1);
    assert_eq!(
        responses[0]
            .get("results")
            .and_then(Value::as_array)
            .map(Vec::len),
        Some(1),
        "alice failed to decrypt the response: {value}"
    );
}