brass_aphid_wire_decryption/decryption/
key_manager.rs

1use brass_aphid_wire_messages::iana;
2
3use crate::{
4    decryption::{key_space::KeySpace, Mode},
5    key_log::NssLog,
6};
7use std::{
8    collections::HashMap,
9    ffi::c_void,
10    fmt::Debug,
11    pin::Pin,
12    sync::{Arc, Mutex},
13};
14
15#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
16struct TlsKeys {
17    /// TLS 1.3
18    pub client_handshake_traffic_secret: Option<Vec<u8>>,
19    pub server_handshake_traffic_secret: Option<Vec<u8>>,
20    pub client_application_traffic_secret: Option<Vec<u8>>,
21    pub server_application_traffic_secret: Option<Vec<u8>>,
22    /// TLS 1.2-ish
23    pub master_secret: Option<Vec<u8>>,
24}
25
26impl TlsKeys {
27    fn update_from_nss_log(&mut self, keys: NssLog) -> anyhow::Result<()> {
28        match keys.label.as_str() {
29            "CLIENT_HANDSHAKE_TRAFFIC_SECRET" => {
30                self.client_handshake_traffic_secret = Some(keys.secret);
31            }
32            "SERVER_HANDSHAKE_TRAFFIC_SECRET" => {
33                self.server_handshake_traffic_secret = Some(keys.secret);
34            }
35            "CLIENT_TRAFFIC_SECRET_0" => {
36                self.client_application_traffic_secret = Some(keys.secret);
37            }
38            "SERVER_TRAFFIC_SECRET_0" => {
39                self.server_application_traffic_secret = Some(keys.secret);
40            }
41            "CLIENT_RANDOM" => {
42                self.master_secret = Some(keys.secret);
43            }
44            _ => {}
45        }
46        Ok(())
47    }
48}
49
50/// The KeyManager holds all of the keys that are generated by the key logging callbacks.
51/// This has been primarily designed to work with s2n-tls's key logging, but should
52/// work with other implementations as well.
53///
54/// The KeyManager is internally reference counter and can be safely shared across
55/// connections.
56///
57/// Notes that a single KeyManager is generally set on a config but decrypting is
58/// done per-connection, so that is a one-to-many relationship.
59///
60/// The `Pin` is necessary for safe use with s2n-tls, which will case the reference
61/// to a c_void pointer.
62#[derive(Debug, Clone)]
63pub struct KeyManager(Pin<Arc<Mutex<HashMap<Vec<u8>, TlsKeys>>>>);
64
65impl Default for KeyManager {
66    fn default() -> Self {
67        Self::new()
68    }
69}
70
71impl KeyManager {
72    pub fn new() -> Self {
73        let value = Arc::pin(Mutex::new(HashMap::new()));
74        Self(value)
75    }
76
77    pub fn enable_s2n_logging(&self, config: &mut s2n_tls::config::Builder) {
78        unsafe {
79            config
80                .set_key_log_callback(
81                    Some(Self::s2n_tls_key_log_cb),
82                    self as *const KeyManager as *mut c_void,
83                )
84                .unwrap();
85        }
86    }
87
88    unsafe extern "C" fn s2n_tls_key_log_cb(
89        ctx: *mut std::ffi::c_void,
90        _conn: *mut s2n_tls_sys::s2n_connection,
91        logline: *mut u8,
92        len: usize,
93    ) -> std::ffi::c_int {
94        let handle = &mut *(ctx as *mut Self);
95        let logline = core::slice::from_raw_parts(logline, len);
96
97        // ignore any errors
98        handle.parse_key_log_line(logline);
99
100        0
101    }
102
103    /// This signature is required by OpenSSL
104    pub fn parse_key_log_line(&self, line: &[u8]) {
105        let line = String::from_utf8(line.to_vec()).unwrap();
106        let key = NssLog::from_log_line(&line).unwrap();
107        self.register_key(key);
108    }
109
110    pub fn register_key(&self, key: NssLog) {
111        tracing::debug!("received key {key:?}");
112        self.0
113            .lock()
114            .unwrap()
115            .entry(key.client_random.clone())
116            .or_default()
117            .update_from_nss_log(key)
118            .unwrap();
119    }
120
121    /// the key used for initial handshake messages in TLS 1.3
122    pub fn handshake_space(
123        &self,
124        mode: Mode,
125        client_random: &[u8],
126        cipher: iana::Cipher,
127    ) -> Option<KeySpace> {
128        let key = self.0.lock().unwrap().get(client_random)?.clone();
129        let secret = match mode {
130            Mode::Client => key.client_handshake_traffic_secret?,
131            Mode::Server => key.server_handshake_traffic_secret?,
132        };
133        let space = KeySpace::handshake_traffic_secret(secret, cipher);
134        Some(space)
135    }
136
137    /// Retrieve the KeySpace for a particular application space.
138    ///
139    /// Index will be `0` for the initial set of traffic keys, but higher indicies
140    /// will be used in the event of key updates.
141    pub fn first_application_space(
142        &self,
143        mode: Mode,
144        client_random: &[u8],
145        cipher: iana::Cipher,
146    ) -> Option<KeySpace> {
147        let key = self.0.lock().unwrap().get(client_random)?.clone();
148        let secret = match mode {
149            Mode::Client => key.client_application_traffic_secret?,
150            Mode::Server => key.server_application_traffic_secret?,
151        };
152        let space = KeySpace::first_traffic_secret(secret, cipher);
153        Some(space)
154    }
155}
156
157impl rustls::KeyLog for KeyManager {
158    fn log(&self, label: &str, client_random: &[u8], secret: &[u8]) {
159        let key = NssLog {
160            label: label.to_string(),
161            client_random: client_random.to_vec(),
162            secret: secret.to_vec(),
163        };
164        self.register_key(key)
165    }
166}
brass_aphid_wire_decryption/decryption/key_manager.rs

brass_aphid_wire_decryption/decryption/
key_manager.rs
