1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110

use std::io::{Read, Write};
use std::net::TcpStream;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use byteorder::{BigEndian, ByteOrder};
use crypto_box::aead::{Aead, Payload};
use parking_lot::Mutex;
use sha2::digest::generic_array::GenericArray;
use sha2::Digest;
use typenum::U24;

use crate::event::Event;
use crate::result::VentedResult;

pub use crypto_box::PublicKey;
pub use crypto_box::SecretKey;

/// A cryptographical stream object that handles encryption and decryption of streams
#[derive(Clone)]
pub struct CryptoStream {
    send_stream: Arc<Mutex<TcpStream>>,
    recv_stream: Arc<Mutex<TcpStream>>,
    sent_count: Arc<AtomicUsize>,
    recv_count: Arc<AtomicUsize>,
    secret_box: Arc<Mutex<crypto_box::ChaChaBox>>,
}

impl CryptoStream {
    /// Creates a new crypto stream from a given Tcp Stream and with a given secret
    pub fn new(inner: TcpStream, secret_box: crypto_box::ChaChaBox) -> VentedResult<Self> {
        inner.set_nonblocking(false)?;
        let send_stream = Arc::new(Mutex::new(inner.try_clone()?));
        let recv_stream = Arc::new(Mutex::new(inner));

        Ok(Self {
            send_stream,
            recv_stream,
            sent_count: Arc::new(AtomicUsize::new(0)),
            recv_count: Arc::new(AtomicUsize::new(0)),
            secret_box: Arc::new(Mutex::new(secret_box)),
        })
    }

    /// Sends a new event encrypted
    /// format:
    /// length: u64
    /// data: length
    pub fn send(&self, mut event: Event) -> VentedResult<()> {
        let number = self.sent_count.fetch_add(1, Ordering::SeqCst);
        let nonce = generate_nonce(number);

        let ciphertext = self.secret_box.lock().encrypt(
            &nonce,
            Payload {
                msg: &event.as_bytes(),
                aad: &[],
            },
        )?;
        let mut stream = self.send_stream.lock();
        let mut length_raw = [0u8; 8];
        BigEndian::write_u64(&mut length_raw, ciphertext.len() as u64);

        log::trace!("Encoded event '{}' to raw message", event.name);

        stream.write(&length_raw)?;
        stream.write(&ciphertext)?;
        stream.flush()?;

        log::trace!("Event sent");

        Ok(())
    }

    /// Reads an event from the stream. Blocks until data is received
    pub fn read(&self) -> VentedResult<Event> {
        let mut stream = self.recv_stream.lock();
        let mut length_raw = [0u8; 8];
        stream.read_exact(&mut length_raw)?;

        let length = BigEndian::read_u64(&length_raw);
        let mut ciphertext = vec![0u8; length as usize];
        stream.read(&mut ciphertext)?;
        log::trace!("Received raw message");

        let number = self.recv_count.fetch_add(1, Ordering::SeqCst);
        let nonce = generate_nonce(number);
        let plaintext = self.secret_box.lock().decrypt(
            &nonce,
            Payload {
                msg: &ciphertext,
                aad: &[],
            },
        )?;

        let event = Event::from_bytes(&mut &plaintext[..])?;
        log::trace!("Decoded message to event '{}'", event.name);

        Ok(event)
    }
}

/// Generates a nonce by hashing the input number which is the message counter
fn generate_nonce(number: usize) -> GenericArray<u8, U24> {
    let result = sha2::Sha256::digest(&number.to_be_bytes()).to_vec();
    let mut nonce = [0u8; 24];
    nonce.copy_from_slice(&result[0..24]);

    nonce.into()
}