Ironforce/src/message.rs

use crate::crypto::{Keys, PublicKey};
use crate::res::IFResult;
use crate::tunnel::TunnelPublic;
use alloc::string::String;
use alloc::vec::Vec;
use serde::{Deserialize, Serialize};
use sha2::Digest;

/// A serialized message
pub(crate) type MessageBytes = Vec<u8>;

/// Signature of the message: optional and optionally encrypted sender's key and signed hash
#[derive(Serialize, Deserialize, Clone)]
pub enum Signature {
    /// The message is signed. Author is unknown
    NotSigned,
    /// The message is signed with the sender's key visible to everyone
    Signed {
        sender: PublicKey,
        signature: Vec<u8>,
    },
    /// Sender's key is encrypted for the recipient
    SignedPrivately {
        sender_encrypted: Vec<u8>,
        signature: Vec<u8>,
    },
}

impl Signature {
    /// Get sender's key or its encrypted version for hashing
    pub(crate) fn sender_or_encrypted_sender(&self) -> Option<Vec<u8>> {
        match &self {
            Signature::NotSigned => None,
            Signature::Signed { sender, .. } => Some(sender.to_vec()),
            Signature::SignedPrivately {
                sender_encrypted, ..
            } => Some(sender_encrypted.clone()),
        }
    }
}

/// Network name and version
#[derive(Serialize, Deserialize, Clone)]
pub struct NetworkInfo {
    network_name: String,
    version: String,
}

impl Default for NetworkInfo {
    fn default() -> Self {
        Self {
            version: String::from("0.1.0"),
            network_name: String::from("test"),
        }
    }
}

#[derive(Serialize, Deserialize, Clone)]
pub enum MessageType {
    SingleCast,
    Broadcast,
    Service(ServiceMessageType),
}

impl MessageType {
    fn hash(&self) -> Vec<u8> {
        match self {
            MessageType::SingleCast => Vec::from([0]),
            MessageType::Broadcast => Vec::from([1]),
            MessageType::Service(ServiceMessageType::TunnelBuilding(tunnel)) => {
                [2, 0].iter().chain(tunnel.hash().iter()).copied().collect()
            }
        }
    }
}

#[derive(Serialize, Deserialize, Clone)]
pub enum ServiceMessageType {
    TunnelBuilding(TunnelPublic),
}

#[derive(Serialize, Deserialize, Clone)]
pub enum MessageContent {
    /// Just plaintext message content
    Plain(Vec<u8>),
    /// Message content bytes encrypted for the recipient
    Encrypted(Vec<u8>),
    /// There is no content
    None,
}

impl MessageContent {
    pub fn hash(&self) -> Vec<u8> {
        match self {
            MessageContent::Plain(v) => sha2::Sha512::new()
                .chain(&[0u8; 1])
                .chain(v.as_slice())
                .result()
                .to_vec(),
            MessageContent::Encrypted(v) => sha2::Sha512::new()
                .chain(&[1; 1])
                .chain(v.as_slice())
                .result()
                .to_vec(),
            MessageContent::None => Vec::new(),
        }
    }
}

/// The struct for messages that are sent in the network
#[derive(Serialize, Deserialize, Clone)]
pub struct Message {
    /// Content of the message (not to be confused with the bytes that we are sending through interfaces)
    ///
    /// AKA useful payload
    pub content: MessageContent,
    /// The type of this message
    pub message_type: MessageType,
    /// Sender's signature
    pub signature: Signature,
    /// A random number that is used in hash together with the content
    salt: u64,
    /// Hash of message content and the salt
    hash: Vec<u8>,
    /// Optional: hash of the message encrypted for the recipient, so that the recipient can know that this message is for them, but nobody else
    recipient_verification: Option<Vec<u8>>,
    /// ID of the tunnel that is used
    tunnel_id: u64,
    /// Network info
    network_info: NetworkInfo,
}

impl Message {
    /// Verify message's hash
    pub fn verify_hash(&self) -> bool {
        self.hash
            == Self::calculate_hash(
                &self.content,
                self.message_type.clone(),
                self.signature.sender_or_encrypted_sender(),
                &self.network_info,
            )
    }

    /// Verify sender's signature
    pub fn verify_signature(&self, recipient_keys: Keys) -> bool {
        match &self.signature {
            Signature::NotSigned => true,
            Signature::Signed { signature, sender } => {
                sender.verify_sign(self.hash.as_slice(), signature.as_slice())
            }
            Signature::SignedPrivately { signature, .. } => {
                if let Some(sender) = self.get_sender(&recipient_keys) {
                    sender.verify_sign(
                        self.hash.as_slice(),
                        &match recipient_keys.decrypt_data(signature.as_slice()) {
                            Ok(r) => r,
                            Err(_e) => return false,
                        },
                    )
                } else {
                    false
                }
            }
        }
    }

    /// Check if this message is for this set of keys
    pub fn check_recipient(&self, keys: Keys) -> bool {
        keys.decrypt_data(&self.recipient_verification.clone().unwrap_or_default())
            .is_ok()
    }

    /// Get decrypted content of the message
    pub fn get_decrypted(&self, keys: Keys) -> IFResult<Vec<u8>> {
        Ok(match &self.content {
            MessageContent::Plain(c) => c.clone(),
            MessageContent::Encrypted(encrypted_content) => {
                keys.decrypt_data(encrypted_content.as_slice())?
            }
            MessageContent::None => Vec::new(),
        })
    }

    pub fn calculate_hash(
        content: &MessageContent,
        message_type: MessageType,
        sender_or_encrypted_sender: Option<Vec<u8>>,
        network_info: &NetworkInfo,
    ) -> Vec<u8> {
        sha2::Sha512::new()
            .chain(content.hash().as_slice())
            .chain(message_type.hash().as_slice())
            .chain(sender_or_encrypted_sender.unwrap_or_default().as_slice())
            .chain(network_info.network_name.as_bytes())
            .chain(network_info.version.as_bytes())
            .result()
            .to_vec()
    }

    /// Encrypt hash of the message for the recipient
    pub fn generate_recipient_verification(
        hash: Vec<u8>,
        recipient: PublicKey,
    ) -> rsa::errors::Result<Vec<u8>> {
        recipient.encrypt_data(&hash)
    }

    /// Try to get sender from the signature
    fn get_sender(&self, keys: &Keys) -> Option<PublicKey> {
        match &self.signature {
            Signature::NotSigned => None,
            Signature::Signed { sender, .. } => Some(sender.clone()),
            Signature::SignedPrivately {
                sender_encrypted, ..
            } => {
                if let Some(Some(res)) = keys
                    .decrypt_data(sender_encrypted.as_slice())
                    .ok()
                    .map(|k| PublicKey::from_vec(k).ok())
                {
                    Some(res)
                } else {
                    None
                }
            }
        }
    }

    /// Create new MessageBuilder
    pub fn build() -> MessageBuilder {
        MessageBuilder::new()
    }
}

/// Message builder to create a new message step-by-step, like `Message::build().message_type(...).sign(...)`
pub struct MessageBuilder {
    content: MessageContent,
    /// The type of the message to be built
    message_type: Option<MessageType>,
    /// Sender's keys
    sender: Option<Keys>,
    /// Recipient's public key (if present, the content will be encrypted and recipient verification field will be set)
    recipient: Option<PublicKey>,
    /// ID of the tunnel that is used
    tunnel_id: u64,
}

impl MessageBuilder {
    /// Create a new `MessageBuilder` with default parameters
    pub fn new() -> Self {
        Self {
            content: MessageContent::None,
            message_type: None,
            sender: None,
            recipient: None,
            tunnel_id: 0,
        }
    }

    pub fn content(mut self, cont: Vec<u8>) -> Self {
        self.content = MessageContent::Plain(cont);
        self
    }

    /// Sign the message
    pub fn sign(mut self, keys: &Keys) -> Self {
        self.sender = Some(keys.clone());
        self
    }

    /// Set message's recipient (and therefore set recipient verification and encrypt the content)
    pub fn recipient(mut self, recipient: PublicKey) -> Self {
        self.recipient = Some(recipient);
        self
    }

    /// Set tunnel id
    pub fn tunnel(mut self, tunnel_id: u64) -> Self {
        self.tunnel_id = tunnel_id;
        self
    }

    /// Set message's type
    pub fn message_type(mut self, message_type: MessageType) -> Self {
        self.message_type = Some(message_type);
        self
    }

    /// Get the resulting message
    pub fn build(self) -> IFResult<Message> {
        let salt = rand::random();
        let sender_encrypted = if let (Some(sender_keys), Some(recipient)) =
            (self.sender.as_ref(), self.recipient.as_ref())
        {
            Some(recipient.encrypt_data(&sender_keys.get_public().to_vec())?)
        } else {
            None
        };
        let network_info = NetworkInfo::default();
        let hash = Message::calculate_hash(
            &self.content,
            self.message_type.clone().unwrap(),
            sender_encrypted.clone().or_else(|| {
                self.sender
                    .as_ref()
                    .map(|sender_keys| sender_keys.get_public().to_vec())
            }),
            &network_info,
        );
        let recipient_verification = self
            .recipient
            .as_ref()
            .map(|rec| rec.encrypt_data(&hash).unwrap());
        let signature = match (self.sender, self.recipient) {
            (Some(sender_keys), Some(recipient_key)) => Signature::SignedPrivately {
                sender_encrypted: sender_encrypted.unwrap(),
                signature: recipient_key.encrypt_data(&sender_keys.sign(&hash)?)?,
            },
            (Some(sender_keys), None) => Signature::Signed {
                sender: sender_keys.get_public(),
                signature: sender_keys.sign(&hash)?,
            },
            (None, _) => Signature::NotSigned,
        };
        Ok(Message {
            content: self.content,
            message_type: self.message_type.unwrap(),
            signature,
            salt,
            hash,
            recipient_verification,
            tunnel_id: self.tunnel_id,
            network_info,
        })
    }
}

#[cfg(test)]
use alloc::vec;

#[test]
fn test_hashing_message_type() {
    let msg_type_1 = MessageType::Broadcast;
    let msg_type_2 = MessageType::Service(ServiceMessageType::TunnelBuilding(
        TunnelPublic::new_for_test(),
    ));
    assert_eq!(msg_type_1.hash(), msg_type_1.hash());
    assert_eq!(msg_type_2.hash(), msg_type_2.hash());
    assert_ne!(msg_type_1.hash(), msg_type_2.hash())
}

#[test]
fn test_hash_message_content() {
    let content_1 = MessageContent::Plain(vec![1, 2, 4, 5]);
    let content_2 = MessageContent::Encrypted(vec![1, 2, 4, 5]);
    let content_3 = MessageContent::Plain(vec![1, 3, 4, 5]);
    assert_eq!(content_1.hash(), content_1.hash());
    assert_ne!(content_1.hash(), MessageContent::None.hash());
    assert_ne!(content_1.hash(), content_2.hash());
    assert_ne!(content_1.hash(), content_3.hash());
    assert_ne!(content_3.hash(), content_2.hash());
}

#[test]
fn test_building_message() -> IFResult<()> {
    let keys_1 = Keys::generate();
    let keys_2 = Keys::generate();
    let msg = Message::build()
        .content(b"hello".to_vec())
        .sign(&keys_1)
        .recipient(keys_2.get_public())
        .tunnel(1)
        .message_type(MessageType::SingleCast)
        .build()?;
    assert!(msg.verify_hash());
    assert!(msg.verify_signature(keys_2));
    Ok(())
}