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

use crate::{
    crypto::EphemeralSecret,
    header::{delays::Delay, SphinxHeader},
    payload::Payload,
    route::{Destination, Node},
    Result, SphinxPacket,
};

pub const DEFAULT_PAYLOAD_SIZE: usize = 1024;

pub struct SphinxPacketBuilder<'a> {
    payload_size: usize,
    initial_secret: Option<&'a EphemeralSecret>,
}

impl<'a> SphinxPacketBuilder<'a> {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_payload_size(mut self, payload_size: usize) -> Self {
        self.payload_size = payload_size;
        self
    }

    pub fn with_initial_secret(mut self, initial_secret: &'a EphemeralSecret) -> Self {
        self.initial_secret = Some(initial_secret);
        self
    }

    pub fn build_packet<M: AsRef<[u8]>>(
        &self,
        message: M,
        route: &[Node],
        destination: &Destination,
        delays: &[Delay],
    ) -> Result<SphinxPacket> {
        let (header, payload_keys) = match self.initial_secret.as_ref() {
            Some(initial_secret) => SphinxHeader::new(initial_secret, route, delays, destination),
            None => SphinxHeader::new(&EphemeralSecret::new(), route, delays, destination),
        };

        // no need to check if plaintext has correct length as this check is already performed in payload encapsulation
        let payload =
            Payload::encapsulate_message(message.as_ref(), &payload_keys, self.payload_size)?;
        Ok(SphinxPacket { header, payload })
    }
}

impl<'a> Default for SphinxPacketBuilder<'a> {
    fn default() -> Self {
        SphinxPacketBuilder {
            payload_size: DEFAULT_PAYLOAD_SIZE,
            initial_secret: None,
        }
    }
}