Canonical Transport Layer for Private Notes: APIs for the `RustClient` (integrated)

private_note_transport.rs

mod miden_note_transport_client {
    pub trait NoteTransportClient {
        /// Send a serialized, encrypted note to the receiver `AccountId`
        fn send_note(&self, note: Note, account_id: AccountId) -> Result<(), Error>;
        /// Convenience method that constructs the tag from the account ID and fetches the notes by this tag.
        /// Fetches all the notes addressed to a given account, decrypt them and import them into the client
        /// Fails if the decryption key associated with `account_id` is missing from client's encryption store.
        fn fetch_notes(&self, client: &mut Client, account_id: AccountId) -> Result<Vec<Note>, Error>;
        /// Fetch all the notes with a given tag. Allows more customization than `fetch_notes`.
        fn fetch_notes_by_tag(&self, client: &mut Client, tag: NoteTag, account_id: AccountId) -> Result<Vec<Note>, Error>;
    }
}

mod canonical_transport_layer {
    use miden_note_transport_client::NoteTransportClient;
    use miden_client::Client;

    pub struct CanonicalTransportClient {
        rpc_api: ApiClient,
    }

    impl CanonicalTransportClient {
        pub fn new(endpoint: String) -> Self {
            Self { rpc_api: ApiClient::new(endpoint) }
        }
    }

    impl NoteTransportClient for CanonicalTransportClient {
        fn send_note(&self, note: Note, account_id: AccountId) -> Result<(), Error> {
            // only support v1 accounts which contain the public encryption key
            if account_id.version() != AccountIdVersion::V1 {
                return Err(Error::InvalidAccountId);
            }

            // encrypt the note with the receiver public key
            // serialize the note and send it to the RPC API
            let note_data = note.to_bytes()?;
            let encrypted_note = encrypt(note_data, account_id.enc_pub_key());
            // send the encrypted note to the RPC API
            // The server will store the encrypted note with the tag as the key
            self.rpc_api.send_note(encrypted_note, note.metadata().tag());
            Ok(())
        }

        fn fetch_notes(&self, client: &mut Client, account_id: AccountId) -> Result<Vec<Note>, Error> {
            let tag = NoteTag::from_account_id(account_id);
            self.fetch_notes_by_tag(client, tag, account_id)
        }

        fn fetch_notes_by_tag(&self, client: &mut Client, tag: NoteTag, account_id: AccountId) -> Result<Vec<Note>, Error> {
            let encrypted_notes: Vec<EncryptedNote> = self.rpc_api.fetch_notes(tag)?;
            let mut decrypted_notes = Vec::new();

            for encrypted_note in encrypted_notes {
                // To decrypt, we need to know which of our keys to use.
                // Assuming one encryption key per account, the client can figure this out.
                // The public key used for encryption should be part of the encrypted envelope
                // or derivable.
                let decrypted_data = client.decrypt(account_id.enc_pub_key(), encrypted_note)?;
                let note = Note::read_from_bytes(&decrypted_data)?;
                decrypted_notes.push(note);
            }
            Ok(decrypted_notes)
        }
    }
}

/// Like `TransactionAuthenticator` but for encryption.
pub trait EncryptionStore {
    /// Attempt to decrypt a message with the given public encryption key
    fn decrypt(&self, pub_enc_key: Word, msg: &[u8]) -> Result<Vec<u8>, Error>;
}

/// A filesystem-based encryption store that stores keys in separate files and provides encryption functionality.
/// Similar to `FilesystemKeyStore` but for encryption keys, not signing.
/// 
/// TODO maybe re-use `FilesystemKeyStore`? 
/// (Rather not, since we already keep a reference to the `FilesystemKeyStore` in the `Client` via the tx authenticator)
pub struct FilesystemEncryptionStore {
    ...
}

impl FilesystemEncryptionStore {
    fn add_key(&self, key: &EncryptionKey) -> Result<(), Error> {
        ...
    }

    fn get_key(&self, pub_enc_key: Word) -> Result<Option<EncryptionKey>, Error> {
        ...
    }
}

impl EncryptionStore for FilesystemEncryptionStore {
    fn decrypt(&self, pub_enc_key: Word, msg: &[u8]) -> Result<Vec<u8>, Error> {
        let key = self.get_key(pub_enc_key)?;
        let decrypted = key.decrypt(msg)?;
        Ok(decrypted)
    }
}


pub struct Client {
    ...
    // the client needs a way to decrypt messages (notes)
    encryption_store: Arc<dyn EncryptionStore>,
}

impl Client {
    ...

    // TODO: Do we want the client to decrypt generic messages, or only notes?
    // Not sure on this one, so far there hasn't been any other use case for decrypting generic messages
    pub fn decrypt(&self, pub_enc_key: Word, msg: &[u8]) -> Result<Vec<u8>, Error> {
        self.encryption_store.decrypt(pub_enc_key, msg)
    }

}

// We further require that Alice's pub encryption key is part of the `AccountId`
pub struct AccountIdV1 {
    prefix: Felt,
    suffix: Felt,
    enc_pub_key: Word,
}


/// ---------------------------------------------------------------------------
/// Usage example
/// ---------------------------------------------------------------------------

use miden_note_transport_client::NoteTransportClient;


// Register an external transport accessible via RPC
let endpoint = ...
let encryption_keypair = ...; 
// Similar to the `FilesystemKeyStore` but for the encryption keypair. Maybe can re-use `FilesystemKeyStore`?
let encryption_store_alice = FilesystemEncryptionStore::new("keystore/alice".into()).unwrap();
encryption_store_alice.add_key(&encryption_keypair.secret_key());

let authenticator_alice = ...

let mut alice_client: Client = setup_client(authenticator_alice, encryption_store_alice, ...);

let transport_client_alice = CanonicalTransportClient::new(endpoint);
// In a real-world scenario, Alice and Bob would instantiate these on their own machines
let transport_client_bob = CanonicalTransportClient::new(endpoint);

let (alice, alice_seed) = create_basic_wallet(
    ...,
    // AccountIdV1 contains the public encryption key
    encryption_keypair.public_key(),
);
let (bob, bob_seed) = create_basic_wallet(...);

// Assume that Alice has shared her v1 `AccountId` with Bob.

// Now Bob prepares a note for Alice and sends it over the transport layer
let note = Note::new(
    assets,
    NoteMetadata::new(
        bob.id(),
        NoteType::Private,
        // here Bob is revealing 14 bits about the receiver's account ID
        NoteTag::from_account_id(alice.id()).unwrap(),
        NoteExecutionHint::Always,
        aux,
    )
    .unwrap(),
    build_p2id_recipient(alice.id(), serial_num).unwrap(),
);
transport_client_bob.send_note(note, alice.id());

// Alice fetches the notes
let notes = transport_client_alice.fetch_notes(&mut alice_client, alice.id())?;
alice_client.import_notes(notes);


// Alice consumes the notes
let consume_notes_request = TransactionRequestBuilder::new()
    .build_consume_notes(notes.iter().map(|note| note.id()).collect());

let tx_result = alice_client.new_transaction(alice_client.id(), consume_notes_request);
alice_client.submit_transaction(tx_result);