238 lines
10 KiB
Rust
238 lines
10 KiB
Rust
|
use std::{io::{Read, Write, ErrorKind}, net::{TcpStream}, option::Option, result::Result, time::Duration};
|
||
|
use ed25519_dalek;
|
||
|
use ed25519_dalek::{ed25519::signature::Signature, Verifier, PUBLIC_KEY_LENGTH, SIGNATURE_LENGTH};
|
||
|
use x25519_dalek;
|
||
|
use rand_7::{RngCore, rngs::OsRng};
|
||
|
use sha2::{Sha384, Digest};
|
||
|
use aes_gcm::{Aes128Gcm, aead::Aead, NewAead, aead::Payload, aead::generic_array::GenericArray};
|
||
|
use crate::utils::*;
|
||
|
use crate::crypto::*;
|
||
|
use crate::identity::Identity;
|
||
|
use crate::session_manager::SessionError;
|
||
|
use crate::print_error;
|
||
|
|
||
|
const RANDOM_LEN: usize = 64;
|
||
|
const MESSAGE_LEN_LEN: usize = 4;
|
||
|
type MessageLenType = u32;
|
||
|
|
||
|
pub struct Session {
|
||
|
stream: TcpStream,
|
||
|
handshake_sent_buff: Vec<u8>,
|
||
|
handshake_recv_buff: Vec<u8>,
|
||
|
local_cipher: Option<Aes128Gcm>,
|
||
|
local_iv: Option<[u8; IV_LEN]>,
|
||
|
local_counter: usize,
|
||
|
peer_cipher: Option<Aes128Gcm>,
|
||
|
peer_iv: Option<[u8; IV_LEN]>,
|
||
|
peer_counter: usize,
|
||
|
pub peer_public_key: Option<[u8; PUBLIC_KEY_LENGTH]>,
|
||
|
//pub name: Option<String>,
|
||
|
//pub message_queue: Vec<Vec<u8>>,
|
||
|
}
|
||
|
|
||
|
impl Session {
|
||
|
pub fn new(stream: TcpStream) -> Session {
|
||
|
stream.set_read_timeout(Some(Duration::from_millis(1000))).unwrap();
|
||
|
Session {
|
||
|
stream: stream,
|
||
|
handshake_sent_buff: Vec::new(),
|
||
|
handshake_recv_buff: Vec::new(),
|
||
|
local_cipher: None,
|
||
|
local_iv: None,
|
||
|
local_counter: 0,
|
||
|
peer_cipher: None,
|
||
|
peer_iv: None,
|
||
|
peer_counter: 0,
|
||
|
peer_public_key: None,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
pub fn get_ip(&self) -> String {
|
||
|
self.stream.peer_addr().unwrap().ip().to_string()
|
||
|
}
|
||
|
|
||
|
fn socket_read(&mut self, buff: &mut [u8]) -> Result<usize, SessionError> {
|
||
|
match self.stream.read(buff) {
|
||
|
Ok(read) => {
|
||
|
if read > 0 {
|
||
|
return Ok(read);
|
||
|
} else if read == 0 {
|
||
|
return Err(SessionError::BrokenPipe);
|
||
|
}
|
||
|
}
|
||
|
Err(e) => {
|
||
|
match e.kind() {
|
||
|
ErrorKind::WouldBlock => return Err(SessionError::SocketTimeout),
|
||
|
ErrorKind::ConnectionReset => return Err(SessionError::ConnectionReset),
|
||
|
_ => print_error(format!("Receive error: {}", e))
|
||
|
}
|
||
|
}
|
||
|
};
|
||
|
Err(SessionError::Unknown)
|
||
|
}
|
||
|
|
||
|
fn socket_write(&mut self, buff: &[u8]) -> Result<(), SessionError> {
|
||
|
match self.stream.write(buff) {
|
||
|
Ok(sent) => {
|
||
|
if buff.len() == sent {
|
||
|
return Ok(());
|
||
|
} else if sent == 0 {
|
||
|
return Err(SessionError::BrokenPipe);
|
||
|
}
|
||
|
}
|
||
|
Err(e) => print_error(format!("Send error: {}", e))
|
||
|
};
|
||
|
Err(SessionError::Unknown)
|
||
|
}
|
||
|
|
||
|
fn handshake_read(&mut self, buff: &mut [u8]) -> Result<(), SessionError> {
|
||
|
match self.socket_read(buff) {
|
||
|
Ok(_) => {
|
||
|
self.handshake_recv_buff.extend(buff.as_ref());
|
||
|
Ok(())
|
||
|
}
|
||
|
Err(e) => Err(e)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
fn handshake_write(&mut self, buff: &[u8]) -> Result<(), SessionError> {
|
||
|
match self.socket_write(buff) {
|
||
|
Ok(_) => {
|
||
|
self.handshake_sent_buff.extend(buff);
|
||
|
Ok(())
|
||
|
}
|
||
|
Err(e) => Err(e)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
fn hash_handshake(&self, i_am_bob: bool) -> [u8; 48] {
|
||
|
let handshake_bytes = if i_am_bob {
|
||
|
[self.handshake_sent_buff.as_slice(), self.handshake_recv_buff.as_slice()].concat()
|
||
|
} else {
|
||
|
[self.handshake_recv_buff.as_slice(), self.handshake_sent_buff.as_slice()].concat()
|
||
|
};
|
||
|
let mut hasher = Sha384::new();
|
||
|
hasher.update(handshake_bytes);
|
||
|
let handshake_hash = hasher.finalize();
|
||
|
to_array_48(handshake_hash.as_slice())
|
||
|
}
|
||
|
|
||
|
fn on_handshake_successful(&mut self, application_keys: ApplicationKeys){
|
||
|
self.local_cipher = Some(Aes128Gcm::new_varkey(&application_keys.local_key).unwrap());
|
||
|
self.local_iv = Some(application_keys.local_iv);
|
||
|
self.peer_cipher = Some(Aes128Gcm::new_varkey(&application_keys.peer_key).unwrap());
|
||
|
self.peer_iv = Some(application_keys.peer_iv);
|
||
|
self.handshake_sent_buff.clear();
|
||
|
self.handshake_sent_buff.shrink_to_fit();
|
||
|
self.handshake_recv_buff.clear();
|
||
|
self.handshake_recv_buff.shrink_to_fit();
|
||
|
}
|
||
|
|
||
|
pub fn do_handshake(&mut self, identity: &Identity) -> Result<(), SessionError> {
|
||
|
//ECDHE initial exchange
|
||
|
//generate random bytes
|
||
|
let mut handshake_buffer = [0; RANDOM_LEN+PUBLIC_KEY_LENGTH];
|
||
|
OsRng.fill_bytes(&mut handshake_buffer[..RANDOM_LEN]);
|
||
|
//generate ephemeral x25519 keys
|
||
|
let ephemeral_secret = x25519_dalek::EphemeralSecret::new(OsRng);
|
||
|
let ephemeral_public_key = x25519_dalek::PublicKey::from(&ephemeral_secret);
|
||
|
handshake_buffer[RANDOM_LEN..].copy_from_slice(&ephemeral_public_key.to_bytes());
|
||
|
self.handshake_write(&handshake_buffer)?;
|
||
|
self.handshake_read(&mut handshake_buffer)?;
|
||
|
let peer_ephemeral_public_key = x25519_dalek::PublicKey::from(to_array_32(&handshake_buffer[RANDOM_LEN..]));
|
||
|
//calc handshake keys
|
||
|
let i_am_bob = self.handshake_sent_buff < self.handshake_recv_buff; //mutual consensus for keys attribution
|
||
|
let handshake_hash = self.hash_handshake(i_am_bob);
|
||
|
let shared_secret = ephemeral_secret.diffie_hellman(&peer_ephemeral_public_key);
|
||
|
let handshake_keys = HandshakeKeys::derive_keys(shared_secret.to_bytes(), handshake_hash, i_am_bob);
|
||
|
|
||
|
|
||
|
//encrypted handshake
|
||
|
//generate random bytes
|
||
|
let mut random_bytes = [0; RANDOM_LEN];
|
||
|
OsRng.fill_bytes(&mut random_bytes);
|
||
|
self.handshake_write(&random_bytes)?;
|
||
|
//receive peer random bytes
|
||
|
let mut peer_random = [0; RANDOM_LEN];
|
||
|
self.handshake_read(&mut peer_random)?;
|
||
|
//get public key & sign our ephemeral public key
|
||
|
let mut auth_msg = [0; PUBLIC_KEY_LENGTH+SIGNATURE_LENGTH];
|
||
|
auth_msg[..PUBLIC_KEY_LENGTH].copy_from_slice(&identity.get_public_key());
|
||
|
auth_msg[PUBLIC_KEY_LENGTH..].copy_from_slice(&identity.sign(ephemeral_public_key.as_bytes()));
|
||
|
//encrypt auth_msg
|
||
|
let local_cipher = Aes128Gcm::new_varkey(&handshake_keys.local_key).unwrap();
|
||
|
let mut local_handshake_counter = 0;
|
||
|
let nonce = iv_to_nonce(&handshake_keys.local_iv, &mut local_handshake_counter);
|
||
|
let encrypted_auth_msg = local_cipher.encrypt(GenericArray::from_slice(&nonce), auth_msg.as_ref()).unwrap();
|
||
|
self.handshake_write(&encrypted_auth_msg)?;
|
||
|
|
||
|
let mut encrypted_peer_auth_msg = [0; PUBLIC_KEY_LENGTH+SIGNATURE_LENGTH+AES_TAG_LEN];
|
||
|
self.handshake_read(&mut encrypted_peer_auth_msg)?;
|
||
|
//decrypt peer_auth_msg
|
||
|
let peer_cipher = Aes128Gcm::new_varkey(&handshake_keys.peer_key).unwrap();
|
||
|
let mut peer_handshake_counter = 0;
|
||
|
let peer_nonce = iv_to_nonce(&handshake_keys.peer_iv, &mut peer_handshake_counter);
|
||
|
match peer_cipher.decrypt(GenericArray::from_slice(&peer_nonce), encrypted_peer_auth_msg.as_ref()) {
|
||
|
Ok(peer_auth_msg) => {
|
||
|
//verify ephemeral public key signature
|
||
|
self.peer_public_key = Some(to_array_32(&peer_auth_msg[..PUBLIC_KEY_LENGTH]));
|
||
|
let peer_public_key = ed25519_dalek::PublicKey::from_bytes(&self.peer_public_key.unwrap()).unwrap();
|
||
|
let peer_signature = Signature::from_bytes(&peer_auth_msg[PUBLIC_KEY_LENGTH..]).unwrap();
|
||
|
if peer_public_key.verify(peer_ephemeral_public_key.as_bytes(), &peer_signature).is_ok() {
|
||
|
let handshake_hash = self.hash_handshake(i_am_bob);
|
||
|
//sending handshake finished
|
||
|
let handshake_finished = compute_handshake_finished(handshake_keys.local_handshake_traffic_secret, handshake_hash);
|
||
|
self.socket_write(&handshake_finished)?;
|
||
|
let mut peer_handshake_finished = [0; HASH_OUTPUT_LEN];
|
||
|
self.socket_read(&mut peer_handshake_finished)?;
|
||
|
if verify_handshake_finished(peer_handshake_finished, handshake_keys.peer_handshake_traffic_secret, handshake_hash) {
|
||
|
//calc application keys
|
||
|
let application_keys = ApplicationKeys::derive_keys(handshake_keys.handshake_secret, handshake_hash, i_am_bob);
|
||
|
self.on_handshake_successful(application_keys);
|
||
|
return Ok(());
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
Err(_) => {}
|
||
|
}
|
||
|
Err(SessionError::TransmissionCorrupted)
|
||
|
}
|
||
|
|
||
|
pub fn encrypt_and_send(&mut self, message: &[u8]) -> Result<(), SessionError> {
|
||
|
let cipher_len = (message.len() as MessageLenType).to_be_bytes();
|
||
|
let payload = Payload {
|
||
|
msg: message,
|
||
|
aad: &cipher_len
|
||
|
};
|
||
|
let nonce = iv_to_nonce(&self.local_iv.unwrap(), &mut self.local_counter);
|
||
|
println!("nonce: {}", hex::encode(&nonce));
|
||
|
let cipher_text = self.local_cipher.as_ref().unwrap().encrypt(GenericArray::from_slice(&nonce), payload).unwrap();
|
||
|
self.socket_write([&cipher_len, cipher_text.as_slice()].concat().as_slice())
|
||
|
}
|
||
|
|
||
|
pub fn receive_and_decrypt(&mut self) -> Result<Vec<u8>, SessionError> {
|
||
|
let mut message_len = [0; MESSAGE_LEN_LEN];
|
||
|
self.socket_read(&mut message_len)?;
|
||
|
let recv_len = MessageLenType::from_be_bytes(message_len) as usize + AES_TAG_LEN;
|
||
|
if recv_len <= 50000000 {
|
||
|
let mut cipher_text = vec![0; recv_len];
|
||
|
let mut read = 0;
|
||
|
while read < recv_len {
|
||
|
read += self.socket_read(&mut cipher_text[read..])?;
|
||
|
}
|
||
|
let peer_nonce = iv_to_nonce(&self.peer_iv.unwrap(), &mut self.peer_counter);
|
||
|
let peer_cipher = self.peer_cipher.as_ref().unwrap();
|
||
|
let payload = Payload {
|
||
|
msg: &cipher_text,
|
||
|
aad: &message_len
|
||
|
};
|
||
|
match peer_cipher.decrypt(GenericArray::from_slice(&peer_nonce), payload) {
|
||
|
Ok(plain_text) => Ok(plain_text),
|
||
|
Err(_) => Err(SessionError::TransmissionCorrupted)
|
||
|
}
|
||
|
} else {
|
||
|
Err(SessionError::BufferTooLarge)
|
||
|
}
|
||
|
}
|
||
|
}
|