1
//! An object mapper for looking up `rpc::Object`s by ID.
2
//!
3
//! This mapper stores strong or weak references, and uses a generational index
4
//! to keep track of names for them.
5
//!
6
//! TODO RPC: Add an object diagram here once the implementation settles down.
7

            
8
use std::sync::{Arc, Weak};
9

            
10
use slotmap_careful::{Key as _, KeyData, SlotMap};
11
use tor_rpcbase as rpc;
12

            
13
pub(crate) mod methods;
14

            
15
slotmap_careful::new_key_type! {
16
    pub(crate) struct GenIdx;
17

            
18
}
19

            
20
/// A weak or a strong reference to an RPC object.
21
//
22
// Note: This type does not pack very efficiently, due to Rust's current lack
23
// of alignment-based niche optimization.
24
// If this ever matters, we can either use two slotmaps, or we can implement
25
// some kind of kludgey hack on our own.
26
#[derive(Clone, derive_more::From)]
27
enum ObjectRef {
28
    /// A strong reference.
29
    Strong(Arc<dyn rpc::Object>),
30
    /// A weak reference reference.
31
    Weak(Weak<dyn rpc::Object>),
32
}
33

            
34
impl ObjectRef {
35
    /// Return this reference as an Arc, if it is present.
36
16
    fn get(&self) -> Option<Arc<dyn rpc::Object>> {
37
16
        match self {
38
10
            ObjectRef::Strong(s) => Some(Arc::clone(s)),
39
6
            ObjectRef::Weak(w) => w.upgrade(),
40
        }
41
16
    }
42
}
43

            
44
/// A mechanism to look up RPC `Objects` by their `ObjectId`.
45
#[derive(Default)]
46
pub(crate) struct ObjMap {
47
    /// Generationally indexed arena of strong object references.
48
    arena: SlotMap<GenIdx, ObjectRef>,
49
}
50

            
51
/// Encoding functions for GenIdx.
52
///
53
/// The encoding is deliberately nondeterministic: we want to avoid situations
54
/// where applications depend on the details of our ObjectIds, or hardcode the
55
/// ObjectIds they expect, or rely on the same  weak generational index getting
56
/// encoded the same way every time they see it.
57
///
58
/// The encoding is deliberately non-cryptographic: we do not want to imply
59
/// that this gives any security. It is just a mild deterrent to misuse.
60
///
61
/// If you find yourself wanting to reverse-engineer this code so that you can
62
/// analyze these object IDs, please contact the Arti developers instead and let
63
/// us give you a better way to do whatever you want.
64
impl GenIdx {
65
    /// The length of a byte-encoded (but not base-64 encoded) GenIdx.
66
    pub(crate) const BYTE_LEN: usize = 16;
67

            
68
    /// Encode `self` into an rpc::ObjectId that we can give to a client.
69
    pub(crate) fn encode(self) -> rpc::ObjectId {
70
        self.encode_with_rng(&mut rand::rng())
71
    }
72

            
73
    /// As `encode`, but take a Rng as an argument. For testing.
74
1040
    fn encode_with_rng<R: rand::RngCore>(self, rng: &mut R) -> rpc::ObjectId {
75
        use base64ct::Encoding;
76
1040
        let bytes = self.to_bytes(rng);
77
1040
        rpc::ObjectId::from(base64ct::Base64UrlUnpadded::encode_string(&bytes[..]))
78
1040
    }
79

            
80
    /// As `encode_with_rng`, but return an array of bytes.
81
1046
    pub(crate) fn to_bytes<R: rand::RngCore>(self, rng: &mut R) -> [u8; Self::BYTE_LEN] {
82
        use rand::Rng;
83
        use tor_bytes::Writer;
84
1046
        let ffi_idx = self.data().as_ffi();
85
1046
        let x = rng.random::<u64>();
86
1046
        let mut bytes = Vec::with_capacity(Self::BYTE_LEN);
87
1046
        bytes.write_u64(x);
88
1046
        bytes.write_u64(ffi_idx.wrapping_add(x));
89

            
90
1046
        bytes.try_into().expect("Length was wrong!")
91
1046
    }
92

            
93
    /// Attempt to decode `id` into a `GenIdx` than an ObjMap can use.
94
1040
    pub(crate) fn try_decode(id: &rpc::ObjectId) -> Result<Self, rpc::LookupError> {
95
        use base64ct::Encoding;
96

            
97
1040
        let bytes = base64ct::Base64UrlUnpadded::decode_vec(id.as_ref())
98
1040
            .map_err(|_| rpc::LookupError::NoObject(id.clone()))?;
99
1040
        Self::from_bytes(&bytes).ok_or_else(|| rpc::LookupError::NoObject(id.clone()))
100
1040
    }
101

            
102
    /// As `try_decode`, but take a slice of bytes.
103
1042
    pub(crate) fn from_bytes(bytes: &[u8]) -> Option<Self> {
104
        use tor_bytes::Reader;
105
1042
        let mut r = Reader::from_slice(bytes);
106
1042
        let x = r.take_u64().ok()?;
107
1042
        let ffi_idx = r.take_u64().ok()?;
108
1042
        r.should_be_exhausted().ok()?;
109

            
110
1042
        let ffi_idx = ffi_idx.wrapping_sub(x);
111
1042
        Some(GenIdx::from(KeyData::from_ffi(ffi_idx)))
112
1042
    }
113
}
114

            
115
impl ObjMap {
116
    /// Create a new empty ObjMap.
117
8
    pub(crate) fn new() -> Self {
118
8
        Self::default()
119
8
    }
120

            
121
    /// Unconditionally insert a strong entry for `value` in self, and return its index.
122
14
    pub(crate) fn insert_strong(&mut self, value: Arc<dyn rpc::Object>) -> GenIdx {
123
14
        self.arena.insert(ObjectRef::Strong(value))
124
14
    }
125

            
126
    /// Unconditionally insert a weak entry for `value` in self, and return its index.
127
10
    pub(crate) fn insert_weak(&mut self, value: &Arc<dyn rpc::Object>) -> GenIdx {
128
10
        self.arena.insert(ObjectRef::Weak(Arc::downgrade(value)))
129
10
    }
130

            
131
    /// Return the entry from this ObjMap for `idx`.
132
24
    pub(crate) fn lookup(&self, idx: GenIdx) -> Result<Arc<dyn rpc::Object>, LookupError> {
133
24
        self.arena
134
24
            .get(idx)
135
24
            .ok_or(LookupError::NoObject)?
136
16
            .get()
137
16
            .ok_or(LookupError::Expired)
138
24
    }
139

            
140
    /// Remove the entry at `idx`.
141
    ///
142
    /// Return true if anything was removed.
143
6
    pub(crate) fn remove(&mut self, idx: GenIdx) -> bool {
144
6
        self.arena.remove(idx).is_some()
145
6
    }
146

            
147
    /// Testing only: Assert that every invariant for this structure is met.
148
    #[cfg(test)]
149
16
    fn assert_okay(&self) {}
150
}
151

            
152
/// A failure from ObjMap::lookup.
153
///
154
/// (This type is immediately returned into rpc::LookupError before we return it.)
155
#[derive(Clone, Debug, thiserror::Error)]
156
pub(crate) enum LookupError {
157
    /// There was no object with the given ID.
158
    #[error("Object not found")]
159
    NoObject,
160

            
161
    /// The object was present, but it was a weak reference that expired.
162
    #[error("Object expired")]
163
    Expired,
164
}
165

            
166
impl LookupError {
167
    /// Convert this `LookupError` into an [`rpc::LookupError`]
168
    pub(crate) fn to_rpc_lookup_error(&self, id: rpc::ObjectId) -> rpc::LookupError {
169
        match self {
170
            LookupError::NoObject => rpc::LookupError::NoObject(id),
171
            LookupError::Expired => rpc::LookupError::Expired(id),
172
        }
173
    }
174
}
175

            
176
#[cfg(test)]
177
mod test {
178
    // @@ begin test lint list maintained by maint/add_warning @@
179
    #![allow(clippy::bool_assert_comparison)]
180
    #![allow(clippy::clone_on_copy)]
181
    #![allow(clippy::dbg_macro)]
182
    #![allow(clippy::mixed_attributes_style)]
183
    #![allow(clippy::print_stderr)]
184
    #![allow(clippy::print_stdout)]
185
    #![allow(clippy::single_char_pattern)]
186
    #![allow(clippy::unwrap_used)]
187
    #![allow(clippy::unchecked_time_subtraction)]
188
    #![allow(clippy::useless_vec)]
189
    #![allow(clippy::needless_pass_by_value)]
190
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
191

            
192
    use super::*;
193
    use derive_deftly::Deftly;
194
    use tor_rpcbase::templates::*;
195

            
196
    #[derive(Clone, Debug, Deftly)]
197
    #[derive_deftly(Object)]
198
    struct ExampleObject(#[allow(unused)] String);
199

            
200
    #[test]
201
    fn map_basics() {
202
        // Insert an object, make sure it gets inserted twice, and look it up.
203
        let obj1 = Arc::new(ExampleObject("abcdef".to_string()));
204
        let mut map = ObjMap::new();
205
        map.assert_okay();
206
        let id1 = map.insert_strong(obj1.clone());
207
        let id2 = map.insert_strong(obj1.clone());
208
        assert_ne!(id1, id2);
209
        let obj1: Arc<dyn rpc::Object> = obj1;
210
        let obj_out1 = map.lookup(id1).unwrap();
211
        let obj_out2 = map.lookup(id2).unwrap();
212
        assert!(Arc::ptr_eq(&obj1, &obj_out1));
213
        assert!(Arc::ptr_eq(&obj1, &obj_out2));
214
        map.assert_okay();
215

            
216
        map.remove(id1);
217
        assert!(map.lookup(id1).is_err());
218
        let obj_out2b = map.lookup(id2).unwrap();
219
        assert!(Arc::ptr_eq(&obj_out2, &obj_out2b));
220

            
221
        map.assert_okay();
222
    }
223

            
224
    #[test]
225
    fn strong_and_weak() {
226
        // Make sure that a strong object behaves like one, and so does a weak
227
        // object.
228
        let obj1: Arc<dyn rpc::Object> = Arc::new(ExampleObject("hello".to_string()));
229
        let obj2: Arc<dyn rpc::Object> = Arc::new(ExampleObject("world".to_string()));
230
        let mut map = ObjMap::new();
231
        let id1 = map.insert_strong(obj1.clone());
232
        let id2 = map.insert_weak(&obj2);
233

            
234
        {
235
            let out1 = map.lookup(id1).unwrap();
236
            let out2 = map.lookup(id2).unwrap();
237
            assert!(Arc::ptr_eq(&obj1, &out1));
238
            assert!(Arc::ptr_eq(&obj2, &out2));
239
        }
240
        map.assert_okay();
241

            
242
        // Now drop every object we've got, and see what we can still find.
243
        drop(obj1);
244
        drop(obj2);
245
        {
246
            let out1 = map.lookup(id1);
247
            let out2 = map.lookup(id2);
248

            
249
            // This one was strong, so it is still there.
250
            assert!(out1.is_ok());
251

            
252
            // This one is weak so it went away.
253
            assert!(out2.is_err());
254
        }
255
        map.assert_okay();
256
    }
257

            
258
    #[test]
259
    fn remove() {
260
        // Make sure that removing an object makes it go away.
261
        let obj1: Arc<dyn rpc::Object> = Arc::new(ExampleObject("hello".to_string()));
262
        let obj2: Arc<dyn rpc::Object> = Arc::new(ExampleObject("world".to_string()));
263
        let mut map = ObjMap::new();
264
        let id1 = map.insert_strong(obj1.clone());
265
        let id2 = map.insert_weak(&obj2);
266
        map.assert_okay();
267

            
268
        map.remove(id1);
269
        map.assert_okay();
270
        assert!(map.lookup(id1).is_err());
271
        assert!(map.lookup(id2).is_ok());
272

            
273
        map.remove(id2);
274
        map.assert_okay();
275
        assert!(map.lookup(id1).is_err());
276
        assert!(map.lookup(id2).is_err());
277
    }
278

            
279
    #[test]
280
    fn duplicates() {
281
        let obj1: Arc<dyn rpc::Object> = Arc::new(ExampleObject("hello".to_string()));
282
        let obj2: Arc<dyn rpc::Object> = Arc::new(ExampleObject("world".to_string()));
283
        let mut map = ObjMap::new();
284
        let id1 = map.insert_strong(obj1.clone());
285
        let id2 = map.insert_weak(&obj2);
286

            
287
        {
288
            assert_ne!(id2, map.insert_weak(&obj1));
289
            assert_ne!(id2, map.insert_weak(&obj2));
290
        }
291

            
292
        {
293
            assert_ne!(id1, map.insert_strong(obj1.clone()));
294
            assert_ne!(id2, map.insert_strong(obj2.clone()));
295
        }
296
    }
297

            
298
    #[test]
299
    fn objid_encoding() {
300
        use rand::Rng;
301
        fn test_roundtrip(a: u32, b: u32, rng: &mut tor_basic_utils::test_rng::TestingRng) {
302
            let a: u64 = a.into();
303
            let b: u64 = b.into();
304
            let data = KeyData::from_ffi((a << 33) | (1_u64 << 32) | b);
305
            let idx = GenIdx::from(data);
306
            let s1 = idx.encode_with_rng(rng);
307
            let s2 = idx.encode_with_rng(rng);
308
            assert_ne!(s1, s2);
309
            assert_eq!(idx, GenIdx::try_decode(&s1).unwrap());
310
            assert_eq!(idx, GenIdx::try_decode(&s2).unwrap());
311
        }
312
        let mut rng = tor_basic_utils::test_rng::testing_rng();
313

            
314
        test_roundtrip(0, 1, &mut rng);
315
        test_roundtrip(0, 2, &mut rng);
316
        test_roundtrip(1, 1, &mut rng);
317
        test_roundtrip(0xffffffff, 0xffffffff, &mut rng);
318

            
319
        for _ in 0..256 {
320
            test_roundtrip(rng.random(), rng.random(), &mut rng);
321
        }
322
    }
323
}