volt_ws_protocol/

rpc.rs

use crate::{
    binary_message::BinaryMessage,
    constant::{PAYLOAD_TYPE_KEY_EXCHANGE, PAYLOAD_TYPE_PAYLOAD},
    crypto_utils::{
        self, aes_decrypt, aes_encrypt, derive_shared_key, verify_signature, IdentityToken,
    },
    log,
};
use base64ct::{Base64, Base64UrlUnpadded, Encoding};
use serde_json::{json, Value};
use std::{collections::VecDeque, vec};

pub struct Rpc {
    method_id: u64,
    method_name: String,
    metadata: IdentityToken,
    token: String,
    response_count: u32,
    client_end_pending: bool,
    client_ended: bool,
    server_ended: bool,
    target_did: Vec<String>,
    target_public_key: String,
    service_id: String,
    call_started: bool,
    key_exchange_pending: bool,
    encryption_key: Option<[u8; 32]>,
    pending_payloads: VecDeque<String>,
}

impl Rpc {
    pub fn new(
        id: &u64,
        method_name: &str,
        target_did: Vec<String>,
        metadata: IdentityToken,
        target_public_key: String,
        service_id: String,
    ) -> Rpc {
        Rpc {
            method_id: id.clone(),
            method_name: method_name.to_string(),
            token: metadata.token.clone(),
            metadata: metadata,
            response_count: 0,
            client_end_pending: false,
            client_ended: false,
            server_ended: false,
            target_did: target_did,
            target_public_key: target_public_key,
            service_id: service_id,
            call_started: false,
            key_exchange_pending: false,
            encryption_key: None,
            pending_payloads: VecDeque::new(),
        }
    }

    pub fn encode_payload(
        &self,
        payload_type: &u8,
        payload_json: &str,
        meta_json: Option<&str>,
    ) -> Vec<u8> {
        // Create a random IV.
        let iv_generate = crypto_utils::random_bytes(16);
        let iv: [u8; 16] = iv_generate.try_into().unwrap();

        // Encrypt the call payload.
        let encrypted_payload = aes_encrypt(&payload_json, &self.encryption_key.unwrap(), &iv);

        // Create a binary message with the encrypted payload.
        let bin_message = BinaryMessage::new_payload(
            payload_type,
            self.method_id,
            meta_json.unwrap_or("").as_bytes().to_vec(),
            iv.to_vec(),
            encrypted_payload,
        );

        bin_message.to_bytes()
    }

    pub fn send(&mut self, payload: String) -> Vec<u8> {
        if self.server_ended {
            // The server has closed their side of the stream.
            log("server has ended the stream, not sending payload");

            // This means we can't send any more data.
            self.client_ended = true;

            // Don't send any payload.
            return vec![];
        }

        if self.call_started {
            // The call has already started, so we're either sending more data or ending the call.
            if !payload.is_empty() {
                log("sending payload");

                // Send the payload to the server
                let payload_json = json!({
                    "method_payload": {
                      "json_payload": payload
                    }
                })
                .to_string();

                self.encode_payload(&PAYLOAD_TYPE_PAYLOAD, &payload_json, None)
            } else {
                log("ending call");

                // Tell the server we've finished sending data (the server may continue to send us data).
                self.client_ended = true;

                let payload_json = json!({
                    "method_end": {
                      "ended": true
                  },
                })
                .to_string();

                let meta_json = json!({
                  "client_end": true
                })
                .to_string();

                self.encode_payload(&PAYLOAD_TYPE_PAYLOAD, &payload_json, Some(&meta_json))
            }
        } else if let Some(_key) = &self.encryption_key {
            // The call hasn't started yet but we have an encryption key, indicating that
            // the key exchange has completed. We can now start the call.
            log("encryption key is present");

            self.call_started = true;

            let mut method_invoke = json!({
                "method_name": self.method_name,
                "json_request": payload
            });

            if !self.service_id.is_empty() {
                method_invoke["service_id"] = Value::String(self.service_id.clone());
            }

            let payload_json = json!({
                "method_invoke": method_invoke
            })
            .to_string();

            self.encode_payload(&PAYLOAD_TYPE_PAYLOAD, &payload_json, None)
        } else {
            // The call hasn't started yet and we don't have an encryption key, indicating that
            // we need to initiate a key exchange.
            log("encryption key is not present");

            // Key exchange is about to begin or pending, so queue this payload.
            self.pending_payloads.push_back(payload);

            // Only send the key exchange payload if it's not already pending.
            if !self.key_exchange_pending {
                log("sending key exchange payload");

                self.key_exchange_pending = true;

                let mut key_exchange_json = json!({
                  "i": self.method_id,
                  "d": self.target_did,
                });

                if !self.service_id.is_empty() {
                    key_exchange_json["s"] = Value::String(self.service_id.clone());
                }

                log(&format!("key exchange json: {}", key_exchange_json));

                // Extract the signature and payload from the token.
                let token_parts: Vec<&str> = self.token.split('.').collect();
                let token_payload = Base64UrlUnpadded::decode_vec(token_parts[1]);
                let token_signature = Base64UrlUnpadded::decode_vec(token_parts[2]);

                let bin_message = BinaryMessage::new_key_exchange(
                    self.method_id,
                    key_exchange_json.to_string().as_bytes().to_vec(),
                    token_signature.unwrap(),
                    token_payload.unwrap(),
                );

                bin_message.to_bytes()
            } else {
                // We're already waiting for a key exchange to complete, so don't send the payload yet.
                vec![]
            }
        }
    }

    pub fn receive(&mut self, bin_message: BinaryMessage) -> Result<String, String> {
        self.response_count += 1;

        if &bin_message.payload_type == PAYLOAD_TYPE_KEY_EXCHANGE {
            log("received key exchange message");

            // The signed message is a concatenation of the public key and nonce.
            let mut verify_message = bin_message.public_key.clone();
            verify_message.extend_from_slice(&bin_message.nonce);

            // Verify the signature using the target service public key (or the Volt
            let verified = verify_signature(
                &self.target_public_key,
                &verify_message,
                &bin_message.public_key_signature,
            );

            if !verified {
                log("signature verification failed");
                return Err("signature verification failed".to_string());
            }

            // Add PEM headers and base64 encode the public key.
            let public_key_pem = format!(
                "-----BEGIN PUBLIC KEY-----\n{}\n-----END PUBLIC KEY-----",
                Base64::encode_string(&bin_message.public_key)
            );

            let derived_key = derive_shared_key(
                &self.metadata.encryption_key.as_ref().unwrap().0,
                &public_key_pem,
            );

            match derived_key {
                Ok(k) => {
                    self.encryption_key = Some(k);

                    Ok(json!({
                      "method_id": self.method_id,
                      "key_exchanged": true
                    })
                    .to_string())
                }
                Err(e) => {
                    log(&format!("error deriving key: {}", e));
                    self.encryption_key = None;

                    Err(e)
                }
            }
        } else if &bin_message.payload_type == PAYLOAD_TYPE_PAYLOAD {
            log("received payload message");

            let decrypted_json: String;

            if bin_message.iv.len() > 0 {
                let iv: [u8; 16] = bin_message.iv[..16].try_into().unwrap();
                let mut encrypted_payload = bin_message.payload.clone();

                decrypted_json = aes_decrypt(
                    encrypted_payload.as_mut_slice(),
                    &self.encryption_key.unwrap(),
                    &iv,
                );
            } else {
                decrypted_json = "{}".to_string();
            }

            let decrypted_payload: Value = serde_json::from_str(&decrypted_json).unwrap();

            let meta_json = match bin_message.meta.len() {
                0 => "{}".to_string(),
                _ => String::from_utf8_lossy(&bin_message.meta).to_string(),
            };

            let mut response: Value = serde_json::from_str(&meta_json).unwrap();
            response["method_id"] = self.method_id.into();
            response["payload"] = decrypted_payload;

            let response_json = response.to_string();

            Ok(response_json)
        } else {
            log("received non-key exchange message");
            Err("unsupported message type".to_string())
        }
    }

    pub fn end(&mut self) -> Vec<u8> {
        match self.client_end_pending {
            true => {
                // The client has already ended the stream, so do nothing.
                vec![]
            }
            false => {
                // The client hasn't ended the stream yet, so do so now.
                log(&format!("client side ended: {}", self.method_id));
                self.client_end_pending = true;
                self.send(String::new())
            }
        }
    }

    /**
     * Get the next pending payload to send - there may be none.
     * This is used to send payloads that were queued while waiting for the key exchange to complete.
     * @returns The next pending payload to send, or an empty array if there are no pending payloads.
     */
    pub fn pending_payload(&mut self) -> Vec<u8> {
        match self.pending_payloads.pop_front() {
            Some(p) => self.send(p),
            None => vec![],
        }
    }

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

impl Drop for Rpc {
    fn drop(&mut self) {
        log(&format!("dropping RPC: {}", self.method_id));
    }
}