1use std::collections::BTreeSet;
17
18use serde::{Deserialize, Serialize};
19
20#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
22pub struct TaintSource(String);
23
24impl TaintSource {
25 pub fn new(label: impl Into<String>) -> Option<TaintSource> {
26 let label = label.into();
27 if label.is_empty() {
28 None
29 } else {
30 Some(TaintSource(label))
31 }
32 }
33
34 pub fn as_str(&self) -> &str {
35 &self.0
36 }
37}
38
39#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
43pub struct TaintSet(BTreeSet<TaintSource>);
44
45impl TaintSet {
46 pub fn empty() -> TaintSet {
47 TaintSet(BTreeSet::new())
48 }
49
50 pub fn of(source: TaintSource) -> TaintSet {
51 let mut set = BTreeSet::new();
52 set.insert(source);
53 TaintSet(set)
54 }
55
56 pub fn union(&self, other: &TaintSet) -> TaintSet {
57 TaintSet(self.0.union(&other.0).cloned().collect())
58 }
59
60 pub fn is_empty(&self) -> bool {
61 self.0.is_empty()
62 }
63
64 pub fn contains(&self, source: &TaintSource) -> bool {
65 self.0.contains(source)
66 }
67
68 pub fn len(&self) -> usize {
69 self.0.len()
70 }
71
72 pub fn sources(&self) -> impl Iterator<Item = &TaintSource> {
73 self.0.iter()
74 }
75}
76
77impl std::fmt::Display for TaintSet {
80 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81 self.0.iter().enumerate().try_for_each(|(at, source)| {
82 let comma = if at == 0 { "" } else { ", " };
83 write!(f, "{comma}{}", source.as_str())
84 })
85 }
86}
87
88#[derive(Debug, Clone, PartialEq, Eq)]
92pub struct Tainted<T> {
93 value: T,
94 taint: TaintSet,
95}
96
97impl<T> Tainted<T> {
98 pub fn new(value: T, source: TaintSource) -> Tainted<T> {
102 Tainted {
103 value,
104 taint: TaintSet::of(source),
105 }
106 }
107
108 pub fn peek(&self) -> &T {
110 &self.value
111 }
112
113 pub fn taint(&self) -> &TaintSet {
114 &self.taint
115 }
116
117 pub fn map<U>(self, f: impl FnOnce(&T) -> U) -> Tainted<U> {
120 Tainted {
121 value: f(&self.value),
122 taint: self.taint,
123 }
124 }
125
126 pub fn join<U, V>(self, other: Tainted<U>, f: impl FnOnce(&T, &U) -> V) -> Tainted<V> {
128 Tainted {
129 value: f(&self.value, &other.value),
130 taint: self.taint.union(&other.taint),
131 }
132 }
133}
134
135#[cfg(test)]
141#[allow(
142 clippy::unwrap_used,
143 clippy::expect_used,
144 clippy::panic,
145 clippy::indexing_slicing,
146 reason = "test code"
147)]
148mod tests {
149 use super::*;
150 use proptest::prelude::*;
151
152 fn source(label: &str) -> TaintSource {
153 TaintSource::new(label).unwrap()
154 }
155
156 #[test]
157 fn empty_label_is_not_a_source() {
158 assert!(TaintSource::new("").is_none());
159 }
160
161 #[test]
162 fn map_keeps_the_full_set_and_join_unions() {
163 let web = Tainted::new("body".to_owned(), source("web:example.com"));
164 let derived = web.map(|s| s.len());
165 assert_eq!(derived.taint().len(), 1);
166 let upload = Tainted::new(3usize, source("file:upload"));
167 let joined = derived.join(upload, |a, b| a.saturating_add(*b));
168 assert_eq!(*joined.peek(), 7);
169 assert!(joined.taint().contains(&source("web:example.com")));
170 assert!(joined.taint().contains(&source("file:upload")));
171 }
172
173 #[test]
174 fn peek_lends_without_moving() {
175 let tainted = Tainted::new(vec![1u8, 2], source("mail:inbound"));
176 assert_eq!(tainted.peek().len(), 2);
177 assert_eq!(tainted.taint().len(), 1);
179 }
180
181 proptest! {
182 #[test]
187 fn a_set_holds_and_names_exactly_its_sources(
188 labels in proptest::collection::btree_set("[a-z]{1,6}", 0..5),
189 stranger in "[A-Z]{1,6}",
190 ) {
191 let set = labels
192 .iter()
193 .map(|label| TaintSet::of(source(label)))
194 .fold(TaintSet::empty(), |all, one| all.union(&one));
195 let named: Vec<&str> = set.sources().map(TaintSource::as_str).collect();
196 let expected: Vec<&str> = labels.iter().map(String::as_str).collect();
197 prop_assert_eq!(&named, &expected);
198 prop_assert_eq!((set.len(), set.is_empty()), (labels.len(), labels.is_empty()));
199 prop_assert!(!set.contains(&source(&stranger)));
200 prop_assert_eq!(set.to_string(), expected.join(", "));
201 }
202
203 #[test]
205 fn join_output_contains_both_inputs(labels_a in proptest::collection::vec("[a-z]{1,8}", 1..4),
206 labels_b in proptest::collection::vec("[a-z]{1,8}", 1..4)) {
207 let mut a = Tainted::new(0u32, source(&format!("a:{}", labels_a[0])));
208 for label in &labels_a[1..] {
209 let extra = Tainted::new(1u32, source(&format!("a:{label}")));
210 a = a.join(extra, |x, y| x.saturating_add(*y));
211 }
212 let mut b = Tainted::new(0u32, source(&format!("b:{}", labels_b[0])));
213 for label in &labels_b[1..] {
214 let extra = Tainted::new(1u32, source(&format!("b:{label}")));
215 b = b.join(extra, |x, y| x.saturating_add(*y));
216 }
217 let a_set = a.taint().clone();
218 let b_set = b.taint().clone();
219 let joined = a.join(b, |x, y| x.saturating_add(*y));
220 for s in a_set.sources() {
221 prop_assert!(joined.taint().contains(s));
222 }
223 for s in b_set.sources() {
224 prop_assert!(joined.taint().contains(s));
225 }
226 prop_assert_eq!(joined.taint().len(), a_set.union(&b_set).len());
227 }
228 }
229}