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
111
112
113
114
115
116
117
118
119
use std::net::Ipv4Addr;
use p256::ecdh::EphemeralSecret;
use crate::session::clients::Client;
use crate::crypto::{asym_to_sym_key, get_random_asym_pair, get_shared_asym_secret, trim_public_key};
use crate::message::{ARdata, DNSMessage, QClass, QType, ResourceRecord};
use crate::message::record::CnameRdata;
use crate::string;
use crate::string::{append_base_domain_to_key, encode_domain_name};

/// Result of a client's handshake request including server key pair and prepared response
pub struct KeySwapContext {
    /// New client structure to track derived shared secret and last seen time
    pub new_client: Client,
    /// Response message to send to the client with the server's public key
    pub response: DNSMessage,
    /// Public key of the server's key pair
    pub server_public: String,
    /// Public key extracted from the client's request
    pub client_public: String
}

/// Generate a random asymmetric key pair, append the dnstp base domain to the secret
pub fn get_key_request_with_base_domain(base_domain: String) -> (EphemeralSecret, String)
{
    let (private, public) = get_random_asym_pair();

    (private, append_base_domain_to_key(trim_public_key(&public), &base_domain))
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Debug, Copy, Clone)]
pub enum DecodeKeyRequestError {
    QuestionCount(usize),
    FirstQuestionNotA(QType),
    SecondQuestionNotA(QType),
    SharedSecretDerivation,
}

/// Take a client's handshake request, process the crypto and prepare a response
///
/// Includes generating a server key pair, using the public key in the response, deriving the shared secret.
pub fn decode_key_request(message: &DNSMessage) -> Result<KeySwapContext, DecodeKeyRequestError>
{
    if message.questions.len() == 2 {

        if message.questions[0].qtype != QType::A
        {
            return Err(DecodeKeyRequestError::FirstQuestionNotA(message.questions[0].qtype));
        }

        let key_question = &message.questions[1];

        if key_question.qtype != QType::A
        {
            return Err(DecodeKeyRequestError::SecondQuestionNotA(key_question.qtype));
        }

        // key is transmitted wihout --- BEGIN KEY -- header and trailer bits and with '.' instead of new lines
        let (fattened_public_key, base_domain) = string::get_fattened_public_key(&key_question.qname);
        // generate the servers public/private key pair
        let (server_private, server_public) = get_random_asym_pair();

        // do the Diffie-Hellman shared secret derivation for the server side symmetric encryption
        match get_shared_asym_secret(&server_private, &fattened_public_key) {
            Ok(secret) => {

                let sym_key = asym_to_sym_key(&secret);
                let new_client = Client::new(sym_key);
                let mut response = message.empty_resp_from_request();

                // return an empty null response for the key hostname (static.BLANK.TLD) question, not that important
                let first_record = ResourceRecord {
                    name_offset: 12,
                    answer_type: QType::A,
                    class: QClass::Internet,
                    ttl: 0,
                    rd_length: 4,
                    r_data: Box::new(ARdata::from(Ipv4Addr::from([127,0,0,1])))
                };

                let server_public_domain = append_base_domain_to_key(
                    trim_public_key(&server_public),
                    &base_domain
                );
                // for the public key question that the client sent, respond with the server's public key
                let second_record = ResourceRecord {
                    name_offset: 12 + (&message.questions[0]).to_bytes().len() as u16,
                    answer_type: QType::CNAME,
                    class: QClass::Internet,
                    ttl: 0,
                    rd_length: encode_domain_name(&server_public_domain).len() as u16,
                    r_data: Box::new(
                        CnameRdata::from(
                            server_public_domain
                        )
                    )
                };

                response.header.answer_record_count = 2;
                response.answer_records = vec![
                    first_record, second_record
                ];

                return Ok(KeySwapContext {
                    new_client,
                    response,
                    server_public,
                    client_public: key_question.qname.to_string()
                });
            }
            Err(_) => {
                return Err(DecodeKeyRequestError::SharedSecretDerivation);
            }
        }
    }
    else
    {
        return Err(DecodeKeyRequestError::QuestionCount(message.questions.len()));
    }
}