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

// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    connection,
    event::{api::SocketAddress, IntoEvent},
    inet, random,
};

#[non_exhaustive]
pub struct Context<'a> {
    pub remote_address: SocketAddress<'a>,
    pub peer_connection_id: &'a [u8],
    pub random: &'a mut dyn random::Generator,
}

impl<'a> Context<'a> {
    #[inline]
    #[doc(hidden)]
    pub fn new(
        remote_address: &'a inet::SocketAddress,
        peer_connection_id: &'a connection::PeerId,
        random: &'a mut dyn random::Generator,
    ) -> Self {
        Self {
            remote_address: remote_address.into_event(),
            peer_connection_id: peer_connection_id.as_bytes(),
            random,
        }
    }
}

pub trait Format: 'static + Send {
    const TOKEN_LEN: usize;

    /// Generate a signed token to be delivered in a NEW_TOKEN frame.
    /// This function will only be called if the provider support NEW_TOKEN frames.
    fn generate_new_token(
        &mut self,
        context: &mut Context<'_>,
        source_connection_id: &connection::LocalId,
        output_buffer: &mut [u8],
    ) -> Option<()>;

    /// Generate a signed token to be delivered in a Retry Packet
    fn generate_retry_token(
        &mut self,
        context: &mut Context<'_>,
        original_destination_connection_id: &connection::InitialId,
        output_buffer: &mut [u8],
    ) -> Option<()>;

    /// Return the original destination connection id of a valid token.
    /// If the token is invalid, return None.
    /// Callers should detect duplicate tokens and treat them as invalid.
    fn validate_token(
        &mut self,
        context: &mut Context<'_>,
        token: &[u8],
    ) -> Option<connection::InitialId>;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Source {
    RetryPacket,
    NewTokenFrame,
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use super::*;
    use crate::crypto::retry;

    #[derive(Debug, Default)]
    pub struct Format(());

    impl super::Format for Format {
        const TOKEN_LEN: usize = retry::example::TOKEN_LEN;

        fn generate_new_token(
            &mut self,
            _context: &mut Context<'_>,
            _source_connection_id: &connection::LocalId,
            _output_buffer: &mut [u8],
        ) -> Option<()> {
            // TODO implement one for testing
            None
        }

        fn generate_retry_token(
            &mut self,
            _context: &mut Context<'_>,
            _original_destination_connection_id: &connection::InitialId,
            output_buffer: &mut [u8],
        ) -> Option<()> {
            output_buffer.copy_from_slice(&retry::example::TOKEN);
            Some(())
        }

        fn validate_token(
            &mut self,
            _context: &mut Context<'_>,
            token: &[u8],
        ) -> Option<connection::InitialId> {
            if token == retry::example::TOKEN {
                return Some(connection::InitialId::TEST_ID);
            }

            None
        }
    }

    impl Format {
        pub fn new() -> Self {
            Self(())
        }
    }
}