dope-tls 0.1.0

TCP TLS streams for dope runtime, shin-backed
Documentation

use std::collections::VecDeque;

use shin::record::{
    ContentType, HEADER_LEN, MAX_CIPHERTEXT_BODY, Opener, PlaintextRecord, RecordError, Sealer,
};
use shin::{Epoch as ShinEpoch, Event, KeyDirection};

const REC_CCS: u8 = 20;
const REC_ALERT: u8 = 21;
const REC_HS_PLAIN: u8 = 22;
const REC_AEAD: u8 = 23;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Phase {
    Handshaking,
    Established,
    Closed,
}

#[derive(Debug)]
pub enum TlsError {
    Other,
    Detail(shin::Error),
    Record(RecordError),
    UnexpectedRecord,
    NotEstablished,
    Io(std::io::Error),
    PlaintextAlert { level: u8, desc: u8 },
}

impl From<RecordError> for TlsError {
    fn from(e: RecordError) -> Self {
        Self::Record(e)
    }
}

impl PartialEq for TlsError {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Other, Self::Other)
            | (Self::UnexpectedRecord, Self::UnexpectedRecord)
            | (Self::NotEstablished, Self::NotEstablished) => true,
            (Self::Record(a), Self::Record(b)) => a == b,
            (Self::Io(a), Self::Io(b)) => a.kind() == b.kind(),
            (
                Self::PlaintextAlert { level: la, desc: da },
                Self::PlaintextAlert { level: lb, desc: db },
            ) => la == lb && da == db,
            _ => false,
        }
    }
}

enum Side {
    Client(shin::client::Client),
    Server(shin::server::Server),
}

pub struct TlsState {
    side: Side,
    phase: Phase,

    handshake_sealer: Option<Sealer>,
    handshake_opener: Option<Opener>,
    app_sealer: Option<Sealer>,
    app_opener: Option<Opener>,

    recv_buf: Vec<u8>,
    pending_send: Vec<u8>,
    incoming_app: VecDeque<Vec<u8>>,

    handshake_keys_ready: bool,
}

impl TlsState {
    pub fn new_client(config: shin::client::Config) -> Result<Self, TlsError> {
        let mut state = Self::empty(Side::Client(shin::client::Client::new(config)));
        let evs = match &mut state.side {
            Side::Client(c) => c.start().map_err(|_| TlsError::Other)?,
            _ => unreachable!(),
        };
        state.absorb_events(evs)?;
        Ok(state)
    }

    pub fn new_server(config: shin::server::Config) -> Self {
        Self::empty(Side::Server(shin::server::Server::new(config)))
    }

    fn seal_app(&mut self, ct: ContentType, data: &[u8]) -> Result<(), TlsError> {
        let sealer = self.app_sealer.as_mut().ok_or(TlsError::NotEstablished)?;
        let wire = sealer.seal(ct, data);
        self.pending_send.extend_from_slice(&wire);
        Ok(())
    }

    fn seal_handshake(&mut self, data: &[u8]) -> Result<(), TlsError> {
        let sealer = self
            .handshake_sealer
            .as_mut()
            .ok_or(TlsError::UnexpectedRecord)?;
        let wire = sealer.seal(ContentType::Handshake, data);
        self.pending_send.extend_from_slice(&wire);
        Ok(())
    }

    fn empty(side: Side) -> Self {
        Self {
            side,
            phase: Phase::Handshaking,
            handshake_sealer: None,
            handshake_opener: None,
            app_sealer: None,
            app_opener: None,
            recv_buf: Vec::new(),
            pending_send: Vec::new(),
            incoming_app: VecDeque::new(),
            handshake_keys_ready: false,
        }
    }

    pub fn read_tcp(&mut self, bytes: &[u8]) -> Result<(), TlsError> {
        self.recv_buf.extend_from_slice(bytes);
        while self.consume_one_record()? {}
        Ok(())
    }

    pub fn try_read_tcp(&mut self, bytes: &[u8]) -> bool {
        bytes.is_empty() || self.read_tcp(bytes).is_ok()
    }

    pub fn pull_send(&mut self) -> Vec<u8> {
        std::mem::take(&mut self.pending_send)
    }

    pub fn pull_send_into(&mut self, sink: &mut Vec<u8>) {
        sink.append(&mut self.pending_send);
    }

    pub fn write_app(&mut self, plaintext: &[u8]) -> Result<(), TlsError> {
        if !self.is_established() {
            return Err(TlsError::NotEstablished);
        }
        self.seal_app(ContentType::ApplicationData, plaintext)
    }

    pub fn send_key_update(&mut self, request_update: bool) -> Result<(), TlsError> {
        if !self.is_established() {
            return Err(TlsError::NotEstablished);
        }
        let evs = match &mut self.side {
            Side::Client(c) => c.send_key_update(request_update).map_err(TlsError::Detail)?,
            Side::Server(s) => s.send_key_update(request_update).map_err(TlsError::Detail)?,
        };
        self.absorb_events(evs)
    }

    pub fn send_close_notify(&mut self) -> Result<(), TlsError> {
        if self.is_closed() {
            return Ok(());
        }
        if !self.is_established() {
            return Err(TlsError::NotEstablished);
        }
        self.seal_app(ContentType::Alert, &[1, 0])?;
        self.phase = Phase::Closed;
        Ok(())
    }

    pub fn pull_app(&mut self) -> Option<Vec<u8>> {
        self.incoming_app.pop_front()
    }

    pub fn is_handshaking(&self) -> bool {
        self.phase == Phase::Handshaking
    }

    pub fn is_established(&self) -> bool {
        self.phase == Phase::Established
    }

    pub fn is_closed(&self) -> bool {
        self.phase == Phase::Closed
    }

    pub fn phase(&self) -> Phase {
        self.phase
    }

    pub fn selected_alpn(&self) -> Option<&[u8]> {
        match &self.side {
            Side::Client(c) => c.selected_alpn(),
            Side::Server(s) => s.selected_alpn(),
        }
    }

    fn consume_one_record(&mut self) -> Result<bool, TlsError> {
        if self.recv_buf.len() < HEADER_LEN {
            return Ok(false);
        }
        let outer = self.recv_buf[0];
        let body_len = u16::from_be_bytes([self.recv_buf[3], self.recv_buf[4]]) as usize;
        if body_len > MAX_CIPHERTEXT_BODY {
            return Err(TlsError::Record(RecordError::BodyTooLarge));
        }
        let total = HEADER_LEN + body_len;
        if self.recv_buf.len() < total {
            return Ok(false);
        }

        match outer {
            REC_CCS => self.handle_ccs(total),
            REC_ALERT => self.handle_alert(total),
            REC_HS_PLAIN => self.handle_handshake_plaintext(total),
            REC_AEAD => self.handle_aead(total),
            _ => Err(TlsError::UnexpectedRecord),
        }
    }

    fn handle_ccs(&mut self, total: usize) -> Result<bool, TlsError> {
        self.recv_buf.drain(0..total);
        Ok(true)
    }

    fn handle_alert(&mut self, total: usize) -> Result<bool, TlsError> {
        let level = self.recv_buf.get(HEADER_LEN).copied().unwrap_or(0);
        let desc = self.recv_buf.get(HEADER_LEN + 1).copied().unwrap_or(0);
        self.recv_buf.drain(0..total);
        self.phase = Phase::Closed;
        Err(TlsError::PlaintextAlert { level, desc })
    }

    fn handle_handshake_plaintext(&mut self, total: usize) -> Result<bool, TlsError> {
        let (rec, consumed) = PlaintextRecord::parse(&self.recv_buf[..total])?
            .ok_or(TlsError::UnexpectedRecord)?;
        let inner = rec.body.to_vec();
        self.recv_buf.drain(0..consumed);
        self.feed_shin(ShinEpoch::Plaintext, &inner)?;
        Ok(true)
    }

    fn handle_aead(&mut self, total: usize) -> Result<bool, TlsError> {
        let opener = if self.is_handshaking() {
            self.handshake_opener.as_mut()
        } else {
            self.app_opener
                .as_mut()
                .or_else(|| self.handshake_opener.as_mut())
        }
        .ok_or(TlsError::UnexpectedRecord)?;

        let (inner_type, range, consumed) = opener
            .open(&mut self.recv_buf[..total])?
            .ok_or(TlsError::UnexpectedRecord)?;
        let inner = self.recv_buf[range].to_vec();
        self.recv_buf.drain(0..consumed);

        match inner_type {
            ContentType::Handshake => {
                let epoch = if self.is_handshaking() && self.handshake_keys_ready {
                    ShinEpoch::Handshake
                } else {
                    ShinEpoch::Application
                };
                self.feed_shin(epoch, &inner)?;
            }
            ContentType::ApplicationData => {
                self.incoming_app.push_back(inner);
            }
            ContentType::Alert => {
                self.phase = Phase::Closed;
            }
            ContentType::ChangeCipherSpec => {}
        }
        Ok(true)
    }

    fn feed_shin(&mut self, epoch: ShinEpoch, data: &[u8]) -> Result<(), TlsError> {
        let evs = match &mut self.side {
            Side::Client(c) => c.read(epoch, data).map_err(TlsError::Detail)?,
            Side::Server(s) => s.read(epoch, data).map_err(TlsError::Detail)?,
        };
        self.absorb_events(evs)
    }

    fn absorb_events(&mut self, events: Vec<Event>) -> Result<(), TlsError> {
        for e in events {
            match e {
                Event::Send { epoch, data } => match epoch {
                    ShinEpoch::Plaintext => {
                        PlaintextRecord::encode(ContentType::Handshake, &data, &mut self.pending_send);
                    }
                    ShinEpoch::Handshake => self.seal_handshake(&data)?,
                    ShinEpoch::Application => {
                        self.seal_app(ContentType::Handshake, &data)?
                    }
                },
                Event::KeysReady {
                    epoch,
                    read_secret,
                    write_secret,
                } => match epoch {
                    ShinEpoch::Handshake => {
                        self.handshake_opener = Some(Opener::from_secret(&read_secret));
                        self.handshake_sealer = Some(Sealer::from_secret(&write_secret));
                        self.handshake_keys_ready = true;
                    }
                    ShinEpoch::Application => {
                        self.app_opener = Some(Opener::from_secret(&read_secret));
                        self.app_sealer = Some(Sealer::from_secret(&write_secret));
                    }
                    ShinEpoch::Plaintext => {}
                },
                Event::KeyUpdate { direction, secret } => match direction {
                    KeyDirection::Read => {
                        self.app_opener = Some(Opener::from_secret(&secret));
                    }
                    KeyDirection::Write => {
                        self.app_sealer = Some(Sealer::from_secret(&secret));
                    }
                },
                Event::PeerExtension { .. } => {}
                Event::NewSessionTicket { .. } | Event::ResumptionSecret { .. } => {}
                Event::ZeroRttKeysReady { .. } => {}
                Event::Done => {
                    self.phase = Phase::Established;
                }
            }
        }
        Ok(())
    }
}