mail_auth/
heap.rs

1/*
2 * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
3 *
4 * SPDX-License-Identifier: Apache-2.0 OR MIT
5 */
6
7use crate::{
8    DnsError, Error,
9    crypto::{CryptoError, SigningKey},
10    dkim::{AtpsRecord, DkimReportRecord, DkimSigner, DomainKey, Signature},
11    dmarc::{DmarcRecord, Uri},
12    dns::TxtRecord,
13    mta_sts::{MtaStsRecord, ReportUri, TlsRptRecord},
14    spf::{Directive, Macro, Mechanism, SpfRecord},
15};
16use std::{mem::size_of, sync::Arc};
17
18#[cfg(any(feature = "aws-lc-rs", feature = "rust-crypto"))]
19use crate::dkim2::{Dkim2Hash, Dkim2Signer, Flag, KeyEntry};
20
21pub(crate) trait HeapSize {
22    fn heap_size(&self) -> usize;
23}
24
25impl HeapSize for String {
26    fn heap_size(&self) -> usize {
27        self.capacity()
28    }
29}
30
31impl<T: HeapSize> HeapSize for Vec<T> {
32    fn heap_size(&self) -> usize {
33        self.capacity() * size_of::<T>() + self.iter().map(HeapSize::heap_size).sum::<usize>()
34    }
35}
36
37impl<T: HeapSize> HeapSize for Box<[T]> {
38    fn heap_size(&self) -> usize {
39        self.len() * size_of::<T>() + self.iter().map(HeapSize::heap_size).sum::<usize>()
40    }
41}
42
43impl<T: HeapSize> HeapSize for Option<T> {
44    fn heap_size(&self) -> usize {
45        self.as_ref().map_or(0, HeapSize::heap_size)
46    }
47}
48
49impl<T: HeapSize> HeapSize for Arc<T> {
50    fn heap_size(&self) -> usize {
51        2 * size_of::<usize>() + size_of::<T>() + self.as_ref().heap_size()
52    }
53}
54
55impl HeapSize for u8 {
56    fn heap_size(&self) -> usize {
57        0
58    }
59}
60
61impl HeapSize for Macro {
62    fn heap_size(&self) -> usize {
63        match self {
64            Macro::Literal(bytes) => bytes.len(),
65            Macro::List(items) => items.heap_size(),
66            Macro::Variable { .. } | Macro::None => 0,
67        }
68    }
69}
70
71impl HeapSize for Mechanism {
72    fn heap_size(&self) -> usize {
73        match self {
74            Mechanism::Include { macro_string }
75            | Mechanism::A { macro_string, .. }
76            | Mechanism::Mx { macro_string, .. }
77            | Mechanism::Ptr { macro_string }
78            | Mechanism::Exists { macro_string } => macro_string.heap_size(),
79            Mechanism::All | Mechanism::Ip4 { .. } | Mechanism::Ip6 { .. } => 0,
80        }
81    }
82}
83
84impl HeapSize for Directive {
85    fn heap_size(&self) -> usize {
86        self.mechanism.heap_size()
87    }
88}
89
90impl HeapSize for SpfRecord {
91    fn heap_size(&self) -> usize {
92        self.directives.heap_size()
93            + self.exp.heap_size()
94            + self.redirect.heap_size()
95            + self.ra.as_ref().map_or(0, |ra| ra.len())
96    }
97}
98
99impl HeapSize for Uri {
100    fn heap_size(&self) -> usize {
101        self.uri.heap_size()
102    }
103}
104
105impl HeapSize for DmarcRecord {
106    fn heap_size(&self) -> usize {
107        self.rua.heap_size() + self.ruf.heap_size()
108    }
109}
110
111impl HeapSize for DomainKey {
112    fn heap_size(&self) -> usize {
113        size_of_val(self.p.as_ref()) + self.p.heap_size()
114    }
115}
116
117impl HeapSize for DkimReportRecord {
118    fn heap_size(&self) -> usize {
119        self.ra.heap_size() + self.rs.heap_size()
120    }
121}
122
123impl HeapSize for AtpsRecord {
124    fn heap_size(&self) -> usize {
125        self.d.heap_size()
126    }
127}
128
129impl HeapSize for MtaStsRecord {
130    fn heap_size(&self) -> usize {
131        self.id.heap_size()
132    }
133}
134
135impl HeapSize for ReportUri {
136    fn heap_size(&self) -> usize {
137        match self {
138            ReportUri::Mail(uri) | ReportUri::Http(uri) => uri.heap_size(),
139        }
140    }
141}
142
143impl HeapSize for TlsRptRecord {
144    fn heap_size(&self) -> usize {
145        self.rua.heap_size()
146    }
147}
148
149impl HeapSize for Error {
150    fn heap_size(&self) -> usize {
151        match self {
152            Error::Io(text)
153            | Error::Dns(DnsError::Resolver(text))
154            | Error::Crypto(CryptoError::Library(text)) => text.heap_size(),
155            _ => 0,
156        }
157    }
158}
159
160impl HeapSize for Signature {
161    fn heap_size(&self) -> usize {
162        self.d.heap_size()
163            + self.s.heap_size()
164            + self.b.capacity()
165            + self.bh.capacity()
166            + self.h.heap_size()
167            + self.z.heap_size()
168            + self.i.heap_size()
169            + self.atps.heap_size()
170    }
171}
172
173impl TxtRecord {
174    /// Approximate heap memory owned by this record, in bytes, including the
175    /// shared allocation of the parsed record. Walks the record once, so cache
176    /// implementations should compute it when the entry is inserted.
177    pub fn heap_size(&self) -> usize {
178        match self {
179            TxtRecord::Spf(record) => record.heap_size(),
180            TxtRecord::SpfExplanation(record) => record.heap_size(),
181            TxtRecord::DomainKey(record) => record.heap_size(),
182            TxtRecord::DkimReport(record) => record.heap_size(),
183            TxtRecord::Dmarc(record) => record.heap_size(),
184            TxtRecord::Atps(record) => record.heap_size(),
185            TxtRecord::MtaSts(record) => record.heap_size(),
186            TxtRecord::TlsRpt(record) => record.heap_size(),
187            TxtRecord::Error(err) => err.heap_size(),
188        }
189    }
190}
191
192impl<T: SigningKey, State> DkimSigner<T, State> {
193    /// Approximate heap memory owned by this signer, in bytes: the signature
194    /// template and the private key material.
195    pub fn heap_size(&self) -> usize {
196        self.key.heap_size() + self.template.heap_size()
197    }
198}
199
200#[cfg(any(feature = "aws-lc-rs", feature = "rust-crypto"))]
201impl<State> Dkim2Signer<State> {
202    /// Approximate heap memory owned by this signer, in bytes: the domain,
203    /// selectors, flags, nonce and private key material.
204    pub fn heap_size(&self) -> usize {
205        self.keys.capacity() * size_of::<KeyEntry>()
206            + self
207                .keys
208                .iter()
209                .map(|entry| entry.key.heap_size() + entry.selector.heap_size())
210                .sum::<usize>()
211            + self.domain.heap_size()
212            + self.hashes.capacity() * size_of::<Dkim2Hash>()
213            + self.flags.capacity() * size_of::<Flag>()
214            + self.nonce.heap_size()
215    }
216}
217
218#[cfg(test)]
219#[cfg(any(feature = "aws-lc-rs", feature = "rust-crypto"))]
220mod test {
221    use crate::{
222        Error,
223        crypto::{Ed25519Key, RsaKey, Sha256, SigningKey},
224        dkim::{DkimSigner, DomainKey},
225        dmarc::DmarcRecord,
226        dns::{TxtRecord, TxtRecordParser},
227        spf::SpfRecord,
228    };
229    use encodify::base64;
230    use rustls_pki_types::{PrivateKeyDer, PrivatePkcs1KeyDer, pem::PemObject};
231    use std::{mem::size_of, sync::Arc};
232
233    const ARC_HEADER: usize = 2 * size_of::<usize>();
234
235    fn spf(record: &str) -> TxtRecord {
236        TxtRecord::Spf(Arc::new(SpfRecord::parse(record.as_bytes()).unwrap()))
237    }
238
239    #[test]
240    fn txt_record_heap_size() {
241        // Every parsed record counts its shared allocation and grows with its content
242        let small = spf("v=spf1 include:_spf.google.com ~all").heap_size();
243        let large = spf(concat!(
244            "v=spf1 ip4:192.0.2.0/24 include:_spf.google.com ",
245            "include:spf.protection.outlook.com include:sendgrid.net a mx -all"
246        ))
247        .heap_size();
248        assert!(small > ARC_HEADER + size_of::<SpfRecord>() + "_spf.google.com".len());
249        assert!(large > small + "spf.protection.outlook.com".len() + "sendgrid.net".len());
250
251        let dmarc = TxtRecord::Dmarc(Arc::new(
252            DmarcRecord::parse(b"v=DMARC1; p=reject; rua=mailto:dmarc@example.org").unwrap(),
253        ))
254        .heap_size();
255        assert!(dmarc >= ARC_HEADER + size_of::<DmarcRecord>() + "dmarc@example.org".len());
256
257        let key = TxtRecord::DomainKey(Arc::new(
258            DomainKey::parse(b"v=DKIM1; k=ed25519; p=11qYAYKxCrfVS/7TyWQHOg7hcvPapiMlrwIaaPcHURo=")
259                .unwrap(),
260        ))
261        .heap_size();
262        assert!(key >= ARC_HEADER + size_of::<DomainKey>() + 32);
263
264        assert_eq!(TxtRecord::Error(Error::Parse).heap_size(), 0);
265        assert_eq!(
266            TxtRecord::Error(Error::Io("timeout".to_string())).heap_size(),
267            "timeout".len()
268        );
269    }
270
271    #[test]
272    fn signer_heap_size() {
273        // A signer counts its template strings and its private key material
274        let rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
275            PrivatePkcs1KeyDer::from_pem_slice(
276                include_str!("../resources/rsa-private.pem").as_bytes(),
277            )
278            .unwrap(),
279        ))
280        .unwrap();
281        let rsa_key = rsa.heap_size();
282        assert!(rsa_key >= 2048 / 8 * 4);
283        let signer = DkimSigner::from_key(rsa)
284            .domain("example.org")
285            .selector("default")
286            .headers(["From", "To", "Subject"]);
287        assert!(signer.heap_size() >= rsa_key + "example.org".len() + "default".len());
288
289        let ed = Ed25519Key::from_seed_and_public_key(
290            &base64::LENIENT
291                .decode("nWGxne/9WmC6hEr0kuwsxERJxWl7MmkZcDusAxyuf2A=")
292                .unwrap(),
293            &base64::LENIENT
294                .decode("11qYAYKxCrfVS/7TyWQHOg7hcvPapiMlrwIaaPcHURo=")
295                .unwrap(),
296        )
297        .unwrap();
298        assert!(ed.heap_size() < rsa_key);
299    }
300}
mail_auth/heap.rs

mail_auth/
heap.rs
