Skip to main content

ethrex_p2p/discv5/
session.rs

1use ethrex_common::H256;
2use hkdf::Hkdf;
3use secp256k1::{
4    Message as SecpMessage, PublicKey, SECP256K1, SecretKey, ecdh::shared_secret_point,
5    ecdsa::Signature,
6};
7use sha2::{Digest, Sha256};
8
9/// A discv5 session
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct Session {
12    pub outbound_key: [u8; 16],
13    pub inbound_key: [u8; 16],
14}
15
16/// Builds the challenge-data from a WHOAREYOU packet
17pub fn build_challenge_data(masking_iv: &[u8], static_header: &[u8], authdata: &[u8]) -> Vec<u8> {
18    let mut data = Vec::with_capacity(masking_iv.len() + static_header.len() + authdata.len());
19    data.extend_from_slice(masking_iv);
20    data.extend_from_slice(static_header);
21    data.extend_from_slice(authdata);
22    data
23}
24
25/// Derives session keys from the handshake.
26/// - `secret_key`: The secret key for ECDH (ephemeral for initiator, static for recipient)
27/// - `public_key`: The public key for ECDH (dest static for initiator, ephemeral for recipient)
28/// - `node_id_a`: The initiator's node ID
29/// - `node_id_b`: The recipient's node ID
30/// - `challenge_data`: The challenge data from WHOAREYOU
31/// - `is_initiator`: True if we are the initiator (node A), false if recipient (node B)
32pub fn derive_session_keys(
33    secret_key: &SecretKey,
34    public_key: &PublicKey,
35    node_id_a: &H256,
36    node_id_b: &H256,
37    challenge_data: &[u8],
38    is_initiator: bool,
39) -> Session {
40    let shared_secret = compressed_shared_secret(public_key, secret_key);
41    let hkdf = Hkdf::<Sha256>::new(Some(challenge_data), &shared_secret);
42
43    let mut kdf_info = b"discovery v5 key agreement".to_vec();
44    kdf_info.extend_from_slice(node_id_a.as_bytes());
45    kdf_info.extend_from_slice(node_id_b.as_bytes());
46
47    let mut key_data = [0u8; 32];
48    hkdf.expand(&kdf_info, &mut key_data)
49        .expect("key_data is 32 bytes long, it can never fail");
50
51    // First 16 bytes are initiator's outbound key, second 16 are recipient's outbound key
52    let mut initiator_key = [0u8; 16];
53    let mut recipient_key = [0u8; 16];
54    initiator_key.copy_from_slice(&key_data[..16]);
55    recipient_key.copy_from_slice(&key_data[16..]);
56
57    let (outbound_key, inbound_key) = if is_initiator {
58        (initiator_key, recipient_key)
59    } else {
60        (recipient_key, initiator_key)
61    };
62
63    Session {
64        outbound_key,
65        inbound_key,
66    }
67}
68
69/// Signs the id-signature input used in the handshake
70pub fn create_id_signature(
71    static_key: &SecretKey,
72    challenge_data: &[u8],
73    ephemeral_pubkey: &[u8],
74    node_id_b: &H256,
75) -> Signature {
76    /*
77    *  id-signature-text  = "discovery v5 identity proof"
78       id-signature-input = id-signature-text || challenge-data || ephemeral-pubkey || node-id-B
79       id-signature       = id_sign(sha256(id-signature-input))
80    */
81    let mut id_signature_input = b"discovery v5 identity proof".to_vec();
82    id_signature_input.extend_from_slice(challenge_data);
83    id_signature_input.extend_from_slice(ephemeral_pubkey);
84    id_signature_input.extend_from_slice(node_id_b.as_bytes());
85
86    let digest = Sha256::digest(&id_signature_input);
87    let message = SecpMessage::from_digest_slice(&digest).expect("32 byte digest");
88    SECP256K1.sign_ecdsa(&message, static_key)
89}
90
91/// Verifies the id-signature from the handshake
92pub fn verify_id_signature(
93    src_pubkey: &PublicKey,
94    challenge_data: &[u8],
95    ephemeral_pubkey: &[u8],
96    node_id_b: &H256,
97    signature: &Signature,
98) -> bool {
99    let mut id_signature_input = b"discovery v5 identity proof".to_vec();
100    id_signature_input.extend_from_slice(challenge_data);
101    id_signature_input.extend_from_slice(ephemeral_pubkey);
102    id_signature_input.extend_from_slice(node_id_b.as_bytes());
103
104    let digest = Sha256::digest(&id_signature_input);
105    let Ok(message) = SecpMessage::from_digest_slice(&digest) else {
106        return false;
107    };
108    SECP256K1
109        .verify_ecdsa(&message, signature, src_pubkey)
110        .is_ok()
111}
112
113/// Creates a secret through elliptic-curve Diffie-Hellman key agreement
114///
115/// ecdh(pubkey, privkey) from the spec
116///
117/// https://github.com/ethereum/devp2p/blob/master/discv5/discv5-theory.md#identity-specific-cryptography-in-the-handshake
118fn compressed_shared_secret(dest_pubkey: &PublicKey, ephemeral_key: &SecretKey) -> [u8; 33] {
119    let xy_point = shared_secret_point(dest_pubkey, ephemeral_key);
120    let mut compressed = [0u8; 33];
121    let y = &xy_point[32..];
122    compressed[0] = if y[31] & 1 == 0 { 0x02 } else { 0x03 };
123    compressed[1..].copy_from_slice(&xy_point[..32]);
124    compressed
125}