volt_ws_protocol/

rpc_manager.rs

use crate::{
    binary_message::BinaryMessage,
    crypto_utils::{get_identity_token, signing_key_from_pem, IdentityToken, SigningKeyPair},
    log,
    rpc::Rpc,
    volt_client_configuration::VoltClientConfiguration,
    volt_configuration::RelayConfiguration,
};
use base64ct::{Base64UrlUnpadded, Encoding};
use ed25519_dalek::Signer;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

pub struct RpcManager {
    pub client_config: VoltClientConfiguration,
    pub connect_url: String,
    pub key_pair: SigningKeyPair,
    pub relaying: bool,
    pub next_id: u64,
    pub rpcs: Arc<Mutex<HashMap<u64, Rpc>>>,
}

impl RpcManager {
    pub fn new() -> RpcManager {
        log("creating new VoltClient");

        RpcManager {
            client_config: VoltClientConfiguration::default(),
            connect_url: String::new(),
            key_pair: SigningKeyPair {
                public_key_pem: None,
                private_key: None,
            },
            relaying: false,
            next_id: 0,
            rpcs: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /**
     * Initialise the client with the provided configuration.
     * @param config_json The JSON configuration for the client, in stringified form.
     * @returns The URL to connect to the Volt websocket server.
     */
    pub fn set_configuration(&mut self, config_json: &str) -> Result<String, String> {
        log(&format!("initialising client with config: {}", config_json));

        self.rpcs.lock().unwrap().clear();

        // Parse the client configuration from the provided JSON.
        self.client_config = match serde_json::from_str(config_json) {
            Ok(config) => {
                log("successfully parsed client config");
                config
            }
            Err(e) => {
                log(&format!("failed to parse client config: {}", e));
                return Err("Failed to parse client config".to_string());
            }
        };

        self.relaying = self.client_config.volt.relay.is_some();
        log(&format!("relaying connection {}", self.relaying));

        // Make sure a public key is provided in the Volt configuration.
        // We could extract it from the certificate, but that would involve
        // adding dependencies, which we want to keep to a minimum at the moment.
        let volt_config = &self.client_config.volt;
        match volt_config.public_key.len() {
            0 => {
                log("no public key provided");
                return Err("No public key provided".to_string());
            }
            _ => {
                log(&format!("volt public key {}", volt_config.public_key));
            }
        }

        // Parse and load the private key from the PEM-encoded credential string.
        let credential = &self.client_config.credential;
        self.key_pair = match signing_key_from_pem(&credential.key) {
            Ok(key_pair) => {
                log("successfully parsed private key");
                key_pair
            }
            Err(e) => {
                log(&format!("failed to parse private key: {}", e));
                return Err("Failed to parse private key".to_string());
            }
        };

        let http_addresss = match &self.client_config.volt.relay {
            // Relay can be either a string or an object.
            Some(relay) => match &relay {
                RelayConfiguration::String(url) => url.clone(),
                RelayConfiguration::Object(relay) => relay.http_address.clone(),
            },
            None => {
                // There is no relay configuration, so use the P2P address.
                self.client_config.volt.http_address.clone()
            }
        };

        // Determine the websocket protocol to use.
        self.connect_url = http_addresss.replace("http", "ws");
        log(&format!(
            "initialised, connect URL is: {}",
            self.connect_url
        ));

        Ok(self.connect_url.clone())
    }

    /**
     * Create a new RPC method.
     * @param method_name The name of the method to create.
     * @param service_json The JSON string representing the service, or an empty string if this is a built-in service.
     * @returns The ID of the new RPC method.
     */
    #[cfg(feature = "chrono")]
    pub fn create_rpc(&mut self, method_name: &str, service_json: &str) -> Result<u64, String> {
        let token_time = chrono::Utc::now().timestamp() as u64;
        self.create_rpc_with_time(method_name, token_time, service_json)
    }

    /**
     * Create a new RPC method.
     * @param method_name The name of the method to create.
     * @param token_time The time in seconds since the Unix epoch for the rpc token.
     * @param service_json The JSON string representing the service, or an empty string if this is a built-in service.
     * @returns The ID of the new RPC method.
     */
    pub fn create_rpc_with_time(
        &mut self,
        method_name: &str,
        token_time: u64,
        service_json: &str,
    ) -> Result<u64, String> {
        let mut service_id = String::new();
        let mut service_client_id = String::new();
        let mut service_client_key = String::new();
        let mut service_relayed = false;

        // Parse the service JSON if it's provided.
        if service_json.len() > 0 {
            log(&format!("service JSON: {}", service_json));

            let service: Value = match serde_json::from_str(&service_json) {
                Ok(service) => service,
                Err(e) => {
                    return Err("Failed to parse service JSON ".to_string() + &e.to_string());
                }
            };

            // Extract the `service_id` property from the service JSON, which is a string.
            service_id = match service["id"].as_str() {
                Some(id) => id.to_string(),
                _ => {
                    return Err("Service ID is not a string".to_string());
                }
            };

            // Extract the `service_description` property from the service JSON, which is a JSON object.
            let service_description = match service["service_description"].as_object() {
                Some(description) => description,
                None => {
                    return Err("Service description is not an object".to_string());
                }
            };

            // Extract the host type from the service description, which is a string.
            let host_type = match service_description["host_type"].as_str() {
                Some(host_type) => host_type,
                None => {
                    return Err("Host type is not a string".to_string());
                }
            };

            service_relayed = match host_type {
                "SERVICE_HOST_TYPE_RELAYED" => true,
                _ => false,
            };

            // Extract the `client_id` property from the service description, which is a string.
            service_client_id = match service_description["host_client_id"].as_str() {
                Some(client_id) => client_id.to_string(),
                None => {
                    return Err("Client ID is not a string".to_string());
                }
            };

            let has_pk = service_description.contains_key("host_public_key");
            if has_pk {
                // Extract the `client_key` property from the service description, which is a string.
                service_client_key = match service_description["host_public_key"].as_str() {
                    Some(client_key) => client_key.to_string(),
                    None => {
                        return Err("Client key is not a string".to_string());
                    }
                };
            } else {
                // Default to the Volt public key.
                service_client_key = self.client_config.volt.public_key.clone();
            }
        }

        if service_client_id.is_empty() || service_client_id == self.client_config.volt.id {
            // If no service is provided or this is a built-in service, we use the Volt public key.
            service_client_key = self.client_config.volt.public_key.clone();
            service_relayed = false;
            service_id = "".to_string();
            service_client_id = "".to_string();
        }

        let mut target_did = Vec::<String>::new();
        if self.relaying {
            target_did.push(self.client_config.volt.id.clone());
        }

        if service_relayed
            && !service_client_id.is_empty()
            && &service_client_id != &self.client_config.volt.id
        {
            target_did.push(service_client_id.clone());
        }

        log(&format!(
            "creating call metadata with {} and {}",
            service_client_id, service_client_key
        ));

        let metadata = match self.get_call_metadata(&service_client_id, token_time) {
            Ok(metadata) => metadata,
            Err(e) => {
                return Err("Failed to create call metadata: ".to_string() + &e);
            }
        };

        log("created call metadata");

        self.next_id += 1;

        let new_rpc = Rpc::new(
            &self.next_id,
            method_name,
            target_did,
            metadata,
            service_client_key,
            service_id,
        );

        let mut rpcs = self.rpcs.lock().unwrap();
        rpcs.insert(self.next_id, new_rpc);

        log(&format!("allocated new ID: {}", self.next_id));

        Ok(self.next_id)
    }

    /**
     * Remove an RPC method with the given method ID.
     * @param method_id The ID of the RPC method to remove.
     */
    pub fn remove_rpc(&mut self, method_id: &u64) {
        log(&format!("removing RPC with method ID: {}", method_id));
        let mut rpcs = self.rpcs.lock().unwrap();
        match rpcs.remove(method_id) {
            Some(rpc) => {
                drop(rpc);
            }
            None => {
                log(&format!("No RPC found for method ID {}", method_id));
            }
        }
    }

    /**
     * Create an identity token for the given target audience, using the client's private key.
     * @param target_host_audience The target audience for the identity token.
     * @param token_time The time in seconds since the Unix epoch for the token.
     * @returns The identity token.
     */
    fn get_call_metadata(
        &self,
        target_host_audience: &String,
        token_time: u64,
    ) -> Result<IdentityToken, String> {
        // Default the audience to the Volt id.
        let mut target_audience = self.client_config.volt.id.clone();
        if !target_host_audience.is_empty() {
            target_audience = target_host_audience.clone();
        }

        let session_id = match &self.client_config.credential.session_id {
            Some(session_id) => session_id,
            None => {
                log("no session ID found");
                return Err("No session ID found".to_string());
            }
        };

        log("creating identity token");

        // Check if we have a session and we're using long-lived tokens.
        let token_info = get_identity_token(
            &session_id,
            &self.key_pair,
            &target_audience,
            true,
            token_time,
            60,
        );

        Ok(token_info)
    }

    /**
     * Encodes and encrypts a JSON payload for the given method ID.
     * @param method_id The method ID to encode the payload for.
     * @param payload_json The JSON payload to encode.
     * @returns The encoded and encrypted payload as a binary buffer.
     */
    pub fn encode_payload(&self, method_id: &u64, payload_json: &str) -> Result<Vec<u8>, String> {
        log(&format!("preparing payload for method ID: {}", method_id));

        // See if we have an RPC for this method ID.
        let mut rpcs = self.rpcs.lock().unwrap();
        let lookup = rpcs.get_mut(method_id);

        match lookup {
            Some(rpc) => {
                // We have an RPC for this method ID.
                log("RPC found for method ID");

                match payload_json.is_empty() {
                    true => {
                        // We have an empty payload, so just end the RPC.
                        Ok(rpc.end())
                    }
                    false => {
                        // We have a payload, so send it.
                        Ok(rpc.send(payload_json.to_string()))
                    }
                }
            }
            None => {
                // We don't have an RPC for this method ID.
                Err(format!("No RPC found for method ID {}", method_id))
            }
        }
    }

    /**
     * Decode and decrypts a binary payload.
     * @param payload The binary payload to decode.
     * @returns The decoded payload as a JSON string.
     */
    pub fn decode_payload(&self, payload: Vec<u8>) -> Result<String, String> {
        let message = BinaryMessage::from_bytes(&payload);

        let mut rpcs = self.rpcs.lock().unwrap();
        let rpc = rpcs.get_mut(&message.method_id);

        match rpc {
            Some(rpc) => rpc.receive(message),
            None => {
                log(&format!("no RPC found for method ID {}", message.method_id));
                Err("No RPC found for method ID".to_string())
            }
        }
    }

    /**
     * Returns any pending payload for the given method ID.
     * @param method_id The method ID to check for pending payloads.
     * @returns The pending payload, or empty array if none.
     */
    pub fn pending_payload(&self, method_id: &u64) -> Result<Vec<u8>, String> {
        let mut rpcs = self.rpcs.lock().unwrap();
        let rpc = rpcs.get_mut(method_id);

        match rpc {
            Some(rpc) => Ok(rpc.pending_payload()),
            None => {
                log("No RPC found for method ID");
                Err("No RPC found for method ID".to_string())
            }
        }
    }

    /**
     * Sign the provided data using the client's private key.
     * @param data The data to sign.
     * @returns The signature of the data, encoded as a base64Url string.
     */
    pub fn sign_base64(&self, data: &str) -> Result<String, String> {
        // Use the client private key to sign the data.
        let key = match &self.key_pair.private_key {
            Some(key) => key,
            None => {
                log("no private key found");
                return Err("No private key found".to_string());
            }
        };

        // Sign the payload using the private key
        let signature = match key.try_sign(data.as_bytes()) {
            Ok(signature) => signature,
            Err(e) => {
                log(&format!("failed to sign data: {}", e));
                return Err("Failed to sign data".to_string());
            }
        };

        let signature_bytes = signature.to_bytes();

        let signature_base64 = Base64UrlUnpadded::encode_string(&signature_bytes);

        Ok(signature_base64)
    }

    /**
     * Get the client's public key in PEM format.
     * @returns The client's public key in PEM format.
     */
    pub fn public_key_pem(&self) -> String {
        match &self.key_pair.public_key_pem {
            Some(key) => key.clone(),
            None => "".to_string(),
        }
    }

    pub fn drop(self) {
        // Instance is dropped.
    }
}

impl Drop for RpcManager {
    fn drop(&mut self) {
        log(&format!(
            "dropping RpcManager: {}",
            self.client_config.volt.id
        ));
    }
}