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
#![allow(dead_code)] // TODO RPC: Remove this once the crate is stable.
9

            
10
use std::any;
11
use std::collections::HashMap;
12
use std::sync::{Arc, Weak};
13

            
14
use slotmap_careful::{Key as _, KeyData, SlotMap};
15
use tor_rpcbase as rpc;
16

            
17
slotmap_careful::new_key_type! {
18
    pub(crate) struct WeakIdx;
19
    pub(crate) struct StrongIdx;
20
}
21

            
22
/// A mechanism to look up RPC `Objects` by their `ObjectId`.
23
#[derive(Default)]
24
pub(crate) struct ObjMap {
25
    /// Generationally indexed arena of strong object references.
26
    strong_arena: SlotMap<StrongIdx, Arc<dyn rpc::Object>>,
27
    /// Generationally indexed arena of weak object references.
28
    ///
29
    /// Invariants:
30
    /// * No object has more than one reference in this arena.
31
    /// * Every `entry` in this arena at position `idx` has a corresponding
32
    ///   entry in `reverse_map` entry such that
33
    ///   `reverse_map[entry.tagged_addr()] == idx`.
34
    weak_arena: SlotMap<WeakIdx, WeakArenaEntry>,
35
    /// Backwards reference to look up weak arena references by the underlying
36
    /// object identity.
37
    ///
38
    /// Invariants:
39
    /// * For every weak `(addr,idx)` entry in this map, there is a
40
    ///   corresponding ArenaEntry in `arena` such that
41
    ///   `arena[idx].tagged_addr() == addr`
42
    reverse_map: HashMap<TaggedAddr, WeakIdx>,
43
    /// Testing only: How many times have we tidied this map?
44
    #[cfg(test)]
45
    n_tidies: usize,
46
}
47

            
48
/// A single entry to a weak Object stored in the generational arena.
49
///
50
struct WeakArenaEntry {
51
    /// The actual Arc or Weak reference for the object that we're storing here.
52
    obj: Weak<dyn rpc::Object>,
53
    ///
54
    /// This contains a strong or weak reference, along with the object's true TypeId.
55
    /// See the [`TaggedAddr`] for more info on
56
    /// why this is needed.
57
    id: any::TypeId,
58
}
59

            
60
/// The raw address of an object held in an Arc or Weak.
61
///
62
/// This will be the same for every clone of an Arc, and the same for every Weak
63
/// derived from an Arc.
64
///
65
/// Note that this is not on its own sufficient to uniquely identify an object;
66
/// we must also know the object's TypeId.  See [`TaggedAddr`] for more information.
67
#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
68
struct RawAddr(usize);
69

            
70
/// An address, type identity, and ownership status, used to identify a `Arc<dyn rpc::Object>`.
71
///
72
/// This type is necessary because of the way that Rust implements `Arc<dyn
73
/// Foo>`. It's represented as a "fat pointer", containing:
74
///    * A pointer to the  object itself (and the reference counts that make the
75
///      `Arc<>` work.)
76
///    * A vtable pointer explaining how to invoke the methods of `Foo` on this
77
///      particular object.
78
///
79
/// The trouble here is that `Arc::ptr_eq()` can give an incorrect result, since
80
/// a single type can be instantiated with multiple instances of its vtable
81
/// pointer, which [breaks pointer comparison on `dyn`
82
/// pointers](https://doc.rust-lang.org/std/ptr/fn.eq.html).
83
///
84
/// Thus, instead of comparing objects by (object pointer, vtable pointer)
85
/// tuples, we have to compare them by (object pointer, type id).
86
///
87
/// (We _do_ have to look at type ids, and not just the pointers, since
88
/// `repr(transparent)` enables people to have two `Arc<dyn Object>`s that have
89
/// the same object pointer but different types.)[^1]
90
///
91
/// # Limitations
92
///
93
/// This type only uniquely identifies an Arc/Weak object for that object's
94
/// lifespan. After the last (strong or weak) reference is dropped, this
95
/// `TaggedAddr` may refer to a different object.
96
///
97
/// [^1]: TODO: Verify whether the necessary transmutation here is actually
98
///     guaranteed to work.
99
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
100
struct TaggedAddr {
101
    /// The address of the object.
102
    addr: RawAddr,
103
    /// The type of the object.
104
    type_id: any::TypeId,
105
}
106

            
107
/// A generational index for [`ObjMap`].
108
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
109
pub(crate) enum GenIdx {
110
    /// An index into the arena of weak references.
111
    Weak(WeakIdx),
112
    /// An index into the arena of strong references
113
    Strong(StrongIdx),
114
}
115

            
116
/// Return the [`RawAddr`] of an arbitrary `Arc<T>`.
117
2310
fn raw_addr_of<T: ?Sized>(arc: &Arc<T>) -> RawAddr {
118
2310
    // I assure you, each one of these 'as'es was needed in the version of
119
2310
    // Rust I wrote them in.
120
2310
    RawAddr(Arc::as_ptr(arc) as *const () as usize)
121
2310
}
122

            
123
/// Return the [`RawAddr`] of an arbitrary `Weak<T>`.
124
64928
fn raw_addr_of_weak<T: ?Sized>(arc: &Weak<T>) -> RawAddr {
125
64928
    RawAddr(Weak::as_ptr(arc) as *const () as usize)
126
64928
}
127

            
128
impl WeakArenaEntry {
129
    /// Create a new `WeakArenaEntry` for a weak reference.
130
2212
    fn new(object: &Arc<dyn rpc::Object>) -> Self {
131
2212
        let id = (**object).type_id();
132
2212
        Self {
133
2212
            obj: Arc::downgrade(object),
134
2212
            id,
135
2212
        }
136
2212
    }
137

            
138
    /// Return true if this `ArenaEntry` is really present.
139
    ///
140
    /// Note that this function can produce false positives (if the entry's
141
    /// last strong reference is dropped in another thread), but it can
142
    /// never produce false negatives.
143
3456
    fn is_present(&self) -> bool {
144
3456
        // This is safe from false negatives because: if we can ever
145
3456
        // observe strong_count == 0, then there is no way for anybody
146
3456
        // else to "resurrect" the object.
147
3456
        self.obj.strong_count() > 0
148
3456
    }
149

            
150
    /// Return a strong reference to the object in this entry, if possible.
151
348
    fn strong(&self) -> Option<Arc<dyn rpc::Object>> {
152
348
        Weak::upgrade(&self.obj)
153
348
    }
154

            
155
    /// Return the [`TaggedAddr`] that can be used to identify this entry's object.
156
64920
    fn tagged_addr(&self) -> TaggedAddr {
157
64920
        TaggedAddr {
158
64920
            addr: raw_addr_of_weak(&self.obj),
159
64920
            type_id: self.id,
160
64920
        }
161
64920
    }
162
}
163

            
164
impl TaggedAddr {
165
    /// Return the `TaggedAddr` to uniquely identify `obj` over the course of
166
    /// its existence.
167
2246
    fn for_object(obj: &Arc<dyn rpc::Object>) -> Self {
168
2246
        let type_id = (*obj).type_id();
169
2246
        let addr = raw_addr_of(obj);
170
2246
        TaggedAddr { addr, type_id }
171
2246
    }
172
}
173

            
174
/// Encoding functions for GenIdx.
175
///
176
/// The encoding is deliberately nondeterministic: we want to avoid situations
177
/// where applications depend on the details of our ObjectIds, or hardcode the
178
/// ObjectIds they expect, or rely on the same  weak generational index getting
179
/// encoded the same way every time they see it.
180
///
181
/// The encoding is deliberately non-cryptographic: we do not want to imply
182
/// that this gives any security. It is just a mild deterrent to misuse.
183
///
184
/// If you find yourself wanting to reverse-engineer this code so that you can
185
/// analyze these object IDs, please contact the Arti developers instead and let
186
/// us give you a better way to do whatever you want.
187
impl GenIdx {
188
    /// The length of a byte-encoded (but not base-64 encoded) GenIdx.
189
    pub(crate) const BYTE_LEN: usize = 16;
190

            
191
    /// Return true if this is a strong (owning) reference.
192
    pub(crate) fn is_strong(&self) -> bool {
193
        matches!(self, GenIdx::Strong(_))
194
    }
195

            
196
    /// Encode `self` into an rpc::ObjectId that we can give to a client.
197
    pub(crate) fn encode(self) -> rpc::ObjectId {
198
        self.encode_with_rng(&mut rand::thread_rng())
199
    }
200

            
201
    /// As `encode`, but take a Rng as an argument. For testing.
202
1040
    fn encode_with_rng<R: rand::RngCore>(self, rng: &mut R) -> rpc::ObjectId {
203
1040
        use base64ct::Encoding;
204
1040
        let bytes = self.to_bytes(rng);
205
1040
        rpc::ObjectId::from(base64ct::Base64UrlUnpadded::encode_string(&bytes[..]))
206
1040
    }
207

            
208
    /// As `encode_with_rng`, but return an array of bytes.
209
1048
    pub(crate) fn to_bytes<R: rand::RngCore>(self, rng: &mut R) -> [u8; Self::BYTE_LEN] {
210
        use rand::Rng;
211
        use tor_bytes::Writer;
212
1048
        let (weak_bit, ffi_idx) = match self {
213
560
            GenIdx::Weak(idx) => (1, idx.data().as_ffi()),
214
488
            GenIdx::Strong(idx) => (0, idx.data().as_ffi()),
215
        };
216
1048
        let x = rng.gen::<u64>() << 1;
217
1048
        let mut bytes = Vec::with_capacity(Self::BYTE_LEN);
218
1048
        bytes.write_u64(x | weak_bit);
219
1048
        bytes.write_u64(ffi_idx.wrapping_add(x));
220
1048

            
221
1048
        bytes.try_into().expect("Length was wrong!")
222
1048
    }
223

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

            
228
1040
        let bytes = base64ct::Base64UrlUnpadded::decode_vec(id.as_ref())
229
1040
            .map_err(|_| rpc::LookupError::NoObject(id.clone()))?;
230
1040
        Self::from_bytes(&bytes).ok_or_else(|| rpc::LookupError::NoObject(id.clone()))
231
1040
    }
232

            
233
    /// As `try_decode`, but take a slice of bytes.
234
1044
    pub(crate) fn from_bytes(bytes: &[u8]) -> Option<Self> {
235
1044
        use tor_bytes::Reader;
236
1044
        let mut r = Reader::from_slice(bytes);
237
1044
        let x = r.take_u64().ok()?;
238
1044
        let is_weak = (x & 1) == 1;
239
1044
        let x = x & !1;
240
1044
        let ffi_idx = r.take_u64().ok()?;
241
1044
        r.should_be_exhausted().ok()?;
242

            
243
1044
        let ffi_idx = ffi_idx.wrapping_sub(x);
244
1044

            
245
1044
        if is_weak {
246
558
            Some(GenIdx::Weak(WeakIdx::from(KeyData::from_ffi(ffi_idx))))
247
        } else {
248
486
            Some(GenIdx::Strong(StrongIdx::from(KeyData::from_ffi(ffi_idx))))
249
        }
250
1044
    }
251
}
252

            
253
impl ObjMap {
254
    /// Create a new empty ObjMap.
255
14
    pub(crate) fn new() -> Self {
256
14
        Self::default()
257
14
    }
258

            
259
    /// Reclaim unused space in this map's weak arena.
260
    ///
261
    /// This runs in `O(n)` time.
262
56
    fn tidy(&mut self) {
263
56
        #[cfg(test)]
264
56
        {
265
56
            self.n_tidies += 1;
266
56
        }
267
3484
        self.weak_arena.retain(|index, entry| {
268
3456
            let present = entry.is_present();
269
3456
            if !present {
270
                // For everything we are removing from the `arena`, we must also
271
                // remove it from `reverse_map`.
272
2000
                let ptr = entry.tagged_addr();
273
2000
                let found = self.reverse_map.remove(&ptr);
274
2000
                debug_assert_eq!(found, Some(index));
275
1456
            }
276
3456
            present
277
3484
        });
278
56
    }
279

            
280
    /// If needed, clean the weak arena and resize it.
281
    ///
282
    /// (We call this whenever we're about to add an entry.  This ensures that
283
    /// our insertion operations run in `O(1)` time.)
284
2212
    fn adjust_size(&mut self) {
285
2212
        // If we're about to fill the arena...
286
2212
        if self.weak_arena.len() >= self.weak_arena.capacity() {
287
            // ... we delete any dead `Weak` entries.
288
54
            self.tidy();
289
54
            // Then, if the arena is still above half-full, we double the
290
54
            // capacity of the arena.
291
54
            //
292
54
            // (We have to grow the arena this even if tidy() removed _some_
293
54
            // entries, or else we might re-run tidy() too soon.  But we don't
294
54
            // want to grow the arena if tidy() removed _most_ entries, or some
295
54
            // normal usage patterns will lead to unbounded growth.)
296
54
            if self.weak_arena.len() > self.weak_arena.capacity() / 2 {
297
12
                self.weak_arena.reserve(self.weak_arena.capacity());
298
42
            }
299
2158
        }
300
2212
    }
301

            
302
    /// Unconditionally insert a strong entry for `value` in self, and return its index.
303
22
    pub(crate) fn insert_strong(&mut self, value: Arc<dyn rpc::Object>) -> GenIdx {
304
22
        GenIdx::Strong(self.strong_arena.insert(value))
305
22
    }
306

            
307
    /// Ensure that there is a weak entry for `value` in self, and return an
308
    /// index for it.
309
    /// If there is no entry, create a weak entry.
310
    #[allow(clippy::needless_pass_by_value)] // TODO: Decide whether to make this take a reference.
311
2214
    pub(crate) fn insert_weak(&mut self, value: Arc<dyn rpc::Object>) -> GenIdx {
312
2214
        let ptr = TaggedAddr::for_object(&value);
313
2214
        if let Some(idx) = self.reverse_map.get(&ptr) {
314
            #[cfg(debug_assertions)]
315
2
            match self.weak_arena.get(*idx) {
316
2
                Some(entry) => debug_assert!(entry.tagged_addr() == ptr),
317
                None => panic!("Found a dangling reference"),
318
            }
319
2
            return GenIdx::Weak(*idx);
320
2212
        }
321
2212

            
322
2212
        self.adjust_size();
323
2212
        let idx = self.weak_arena.insert(WeakArenaEntry::new(&value));
324
2212
        self.reverse_map.insert(ptr, idx);
325
2212
        GenIdx::Weak(idx)
326
2214
    }
327

            
328
    /// Return the entry from this ObjMap for `idx`.
329
2224
    pub(crate) fn lookup(&self, idx: GenIdx) -> Option<Arc<dyn rpc::Object>> {
330
2224
        match idx {
331
2210
            GenIdx::Weak(idx) => self.weak_arena.get(idx).and_then(WeakArenaEntry::strong),
332
14
            GenIdx::Strong(idx) => self.strong_arena.get(idx).cloned(),
333
        }
334
2224
    }
335

            
336
    /// Remove and return the entry at `idx`, if any.
337
4
    pub(crate) fn remove(&mut self, idx: GenIdx) -> Option<Arc<dyn rpc::Object>> {
338
4
        match idx {
339
2
            GenIdx::Weak(idx) => {
340
2
                if let Some(entry) = self.weak_arena.remove(idx) {
341
2
                    let old_idx = self.reverse_map.remove(&entry.tagged_addr());
342
2
                    debug_assert_eq!(old_idx, Some(idx));
343
2
                    entry.obj.upgrade()
344
                } else {
345
                    None
346
                }
347
            }
348
2
            GenIdx::Strong(idx) => self.strong_arena.remove(idx),
349
        }
350
4
    }
351

            
352
    /// Testing only: Assert that every invariant for this structure is met.
353
    #[cfg(test)]
354
220
    fn assert_okay(&self) {
355
20972
        for (index, entry) in self.weak_arena.iter() {
356
20972
            let ptr = entry.tagged_addr();
357
20972
            assert_eq!(self.reverse_map.get(&ptr), Some(&index));
358
20972
            assert_eq!(ptr, entry.tagged_addr());
359
        }
360

            
361
20972
        for (ptr, idx) in self.reverse_map.iter() {
362
20972
            let entry = self
363
20972
                .weak_arena
364
20972
                .get(*idx)
365
20972
                .expect("Dangling pointer in reverse map");
366
20972

            
367
20972
            assert_eq!(&entry.tagged_addr(), ptr);
368
        }
369
220
    }
370
}
371

            
372
#[cfg(test)]
373
mod test {
374
    // @@ begin test lint list maintained by maint/add_warning @@
375
    #![allow(clippy::bool_assert_comparison)]
376
    #![allow(clippy::clone_on_copy)]
377
    #![allow(clippy::dbg_macro)]
378
    #![allow(clippy::mixed_attributes_style)]
379
    #![allow(clippy::print_stderr)]
380
    #![allow(clippy::print_stdout)]
381
    #![allow(clippy::single_char_pattern)]
382
    #![allow(clippy::unwrap_used)]
383
    #![allow(clippy::unchecked_duration_subtraction)]
384
    #![allow(clippy::useless_vec)]
385
    #![allow(clippy::needless_pass_by_value)]
386
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
387

            
388
    use super::*;
389
    use derive_deftly::Deftly;
390
    use tor_rpcbase::templates::*;
391

            
392
    #[derive(Clone, Debug, Deftly)]
393
    #[derive_deftly(Object)]
394
    struct ExampleObject(String);
395

            
396
    impl ExampleObject {
397
        fn wrap_arc(self: Arc<Self>) -> Arc<Wrapper> {
398
            // SAFETY: Using `repr(transparent)` on Wrapper guarantees that
399
            // transmuting an ExampleObject to a Wrapper will work correctly.
400
            //
401
            // Given this, and the fact that they have the same alignment and
402
            // size, the documentation for `Arc::from_raw` says that this should
403
            // be safe.
404
            //
405
            // Also this is only a test.
406
            unsafe { Arc::from_raw(Arc::into_raw(self) as *const Wrapper) }
407
        }
408
    }
409

            
410
    #[derive(Clone, Debug, Deftly)]
411
    #[derive_deftly(Object)]
412
    #[repr(transparent)]
413
    struct Wrapper(ExampleObject);
414

            
415
    #[test]
416
    fn arc_to_addr() {
417
        let a1 = Arc::new("Hello world");
418
        let a2 = Arc::clone(&a1);
419
        let a3 = Arc::new("Hello world");
420
        let w1 = Arc::downgrade(&a2);
421

            
422
        assert_eq!(raw_addr_of(&a1), raw_addr_of(&a2));
423
        assert_eq!(raw_addr_of(&a1), raw_addr_of_weak(&w1));
424
        assert_ne!(raw_addr_of(&a1), raw_addr_of(&a3));
425

            
426
        let obj1: Arc<dyn rpc::Object> = Arc::new(ExampleObject("Hello world".into()));
427
        let obj2 = Arc::clone(&obj1);
428
        let obj3: Arc<dyn rpc::Object> = Arc::new(ExampleObject("Hello world".into()));
429
        let obj4 = Arc::clone(&obj3);
430
        let weak1 = Arc::downgrade(&obj1);
431
        let weak2 = Arc::downgrade(&obj3);
432

            
433
        assert_eq!(raw_addr_of(&obj1), raw_addr_of(&obj2));
434
        assert_eq!(raw_addr_of(&obj1), raw_addr_of_weak(&weak1));
435
        assert_eq!(raw_addr_of(&obj3), raw_addr_of(&obj4));
436
        assert_eq!(raw_addr_of(&obj3), raw_addr_of_weak(&weak2));
437
        assert_ne!(raw_addr_of(&obj1), raw_addr_of(&obj3));
438
        assert_ne!(raw_addr_of(&obj1), raw_addr_of(&a1));
439
    }
440

            
441
    #[test]
442
    fn obj_ptr() {
443
        let object = Arc::new(ExampleObject("Ten tons of flax".into()));
444
        let object2: Arc<dyn rpc::Object> = Arc::new(ExampleObject("Ten tons of flax".into()));
445

            
446
        let wrapped: Arc<Wrapper> = object.clone().wrap_arc();
447
        let object_dyn = object.clone() as Arc<dyn rpc::Object>;
448
        let wrapped_dyn = wrapped.clone() as Arc<dyn rpc::Object>;
449

            
450
        let object_dyn2 = Arc::clone(&object_dyn);
451
        let wrapped_dyn2 = Arc::clone(&wrapped_dyn);
452
        let wrapped_weak = Arc::downgrade(&wrapped_dyn);
453

            
454
        assert_eq!(
455
            TaggedAddr::for_object(&object_dyn),
456
            TaggedAddr::for_object(&object_dyn2)
457
        );
458
        assert_ne!(
459
            TaggedAddr::for_object(&object_dyn),
460
            TaggedAddr::for_object(&object2)
461
        );
462

            
463
        assert_eq!(
464
            TaggedAddr::for_object(&wrapped_dyn),
465
            TaggedAddr::for_object(&wrapped_dyn2)
466
        );
467

            
468
        assert_ne!(
469
            TaggedAddr::for_object(&object_dyn),
470
            TaggedAddr::for_object(&wrapped_dyn)
471
        );
472

            
473
        assert_eq!(
474
            TaggedAddr::for_object(&object_dyn).addr,
475
            TaggedAddr::for_object(&wrapped_dyn).addr
476
        );
477
        assert_eq!(
478
            TaggedAddr::for_object(&wrapped_dyn).addr,
479
            raw_addr_of_weak(&wrapped_weak)
480
        );
481

            
482
        assert_eq!(
483
            TaggedAddr::for_object(&object_dyn).type_id,
484
            any::TypeId::of::<ExampleObject>()
485
        );
486
        assert_eq!(
487
            TaggedAddr::for_object(&wrapped_dyn).type_id,
488
            any::TypeId::of::<Wrapper>()
489
        );
490

            
491
        assert_eq!(
492
            TaggedAddr::for_object(&object_dyn).addr,
493
            raw_addr_of(&object)
494
        );
495
        assert_eq!(
496
            TaggedAddr::for_object(&wrapped_dyn).addr,
497
            raw_addr_of(&wrapped)
498
        );
499
        assert_ne!(
500
            TaggedAddr::for_object(&object_dyn).addr,
501
            raw_addr_of(&object2)
502
        );
503
    }
504

            
505
    #[test]
506
    fn map_basics() {
507
        // Insert an object, make sure it only gets inserted once, and look it up.
508
        let obj1 = Arc::new(ExampleObject("abcdef".to_string()));
509
        let mut map = ObjMap::new();
510
        map.assert_okay();
511
        let id1 = map.insert_strong(obj1.clone());
512
        let id2 = map.insert_strong(obj1.clone());
513
        assert_ne!(id1, id2);
514
        let obj_out1 = map.lookup(id1).unwrap();
515
        let obj_out2 = map.lookup(id2).unwrap();
516
        assert_eq!(raw_addr_of(&obj1), raw_addr_of(&obj_out1));
517
        assert_eq!(raw_addr_of(&obj1), raw_addr_of(&obj_out2));
518
        map.assert_okay();
519
    }
520

            
521
    #[test]
522
    fn strong_and_weak() {
523
        // Make sure that a strong object behaves like one, and so does a weak
524
        // object.
525
        let obj1: Arc<dyn rpc::Object> = Arc::new(ExampleObject("hello".to_string()));
526
        let obj2: Arc<dyn rpc::Object> = Arc::new(ExampleObject("world".to_string()));
527
        let mut map = ObjMap::new();
528
        let id1 = map.insert_strong(obj1.clone());
529
        let id2 = map.insert_weak(obj2.clone());
530

            
531
        {
532
            let out1 = map.lookup(id1);
533
            let out2 = map.lookup(id2);
534
            assert_eq!(raw_addr_of(&obj1), raw_addr_of(&out1.unwrap()));
535
            assert_eq!(raw_addr_of(&obj2), raw_addr_of(&out2.unwrap()));
536
        }
537
        let addr1 = raw_addr_of(&obj1);
538
        map.assert_okay();
539

            
540
        // Now drop every object we've got, and see what we can still find.
541
        drop(obj1);
542
        drop(obj2);
543
        {
544
            let out1 = map.lookup(id1);
545
            let out2 = map.lookup(id2);
546

            
547
            // This one was strong, so it is still there.
548
            assert!(out1.is_some());
549
            assert_eq!(raw_addr_of(&out1.unwrap()), addr1);
550

            
551
            // This one is weak so it went away.
552
            assert!(out2.is_none());
553
        }
554
        map.assert_okay();
555
    }
556

            
557
    #[test]
558
    fn remove() {
559
        // Make sure that removing an object makes it go away.
560
        let obj1: Arc<dyn rpc::Object> = Arc::new(ExampleObject("hello".to_string()));
561
        let obj2: Arc<dyn rpc::Object> = Arc::new(ExampleObject("world".to_string()));
562
        let mut map = ObjMap::new();
563
        let id1 = map.insert_strong(obj1.clone());
564
        let id2 = map.insert_weak(obj2.clone());
565
        map.assert_okay();
566

            
567
        map.remove(id1);
568
        map.assert_okay();
569
        assert!(map.lookup(id1).is_none());
570
        assert!(map.lookup(id2).is_some());
571

            
572
        map.remove(id2);
573
        map.assert_okay();
574
        assert!(map.lookup(id1).is_none());
575
        assert!(map.lookup(id2).is_none());
576
    }
577

            
578
    #[test]
579
    fn duplicates() {
580
        // Make sure that inserting duplicate objects behaves right.
581
        let obj1: Arc<dyn rpc::Object> = Arc::new(ExampleObject("hello".to_string()));
582
        let obj2: Arc<dyn rpc::Object> = Arc::new(ExampleObject("world".to_string()));
583
        let mut map = ObjMap::new();
584
        let id1 = map.insert_strong(obj1.clone());
585
        let id2 = map.insert_weak(obj2.clone());
586

            
587
        {
588
            assert_ne!(id2, map.insert_weak(obj1.clone()));
589
            assert_eq!(id2, map.insert_weak(obj2.clone()));
590
        }
591

            
592
        {
593
            assert_ne!(id1, map.insert_strong(obj1.clone()));
594
            assert_ne!(id2, map.insert_strong(obj2.clone()));
595
        }
596
    }
597

            
598
    #[test]
599
    fn upgrade() {
600
        // Make sure that inserting an object as weak and strong (in either
601
        // order) makes two separate entries.
602
        let obj1: Arc<dyn rpc::Object> = Arc::new(ExampleObject("hello".to_string()));
603
        let obj2: Arc<dyn rpc::Object> = Arc::new(ExampleObject("world".to_string()));
604
        let addr1 = raw_addr_of(&obj1);
605
        let addr2 = raw_addr_of(&obj2);
606

            
607
        let mut map = ObjMap::new();
608
        let id1 = map.insert_strong(obj1.clone());
609
        let id2 = map.insert_weak(obj2.clone());
610

            
611
        assert_ne!(id2, map.insert_weak(obj1.clone()));
612
        assert_ne!(id1, map.insert_strong(obj2.clone()));
613
        map.assert_okay();
614

            
615
        drop(obj1);
616
        drop(obj2);
617
        let out1 = map.lookup(id1).unwrap();
618
        let out2 = map.lookup(id2).unwrap();
619
        assert_eq!(raw_addr_of(&out1), addr1);
620
        assert_eq!(raw_addr_of(&out2), addr2);
621
    }
622

            
623
    #[test]
624
    fn tidy() {
625
        let mut map = ObjMap::new();
626
        let mut keep_these = vec![];
627
        let mut s = vec![];
628
        let mut w = vec![];
629
        for _ in 0..100 {
630
            let mut t = vec![];
631
            for _ in 0..10 {
632
                let o = Arc::new(ExampleObject("dump".into()));
633
                w.push(map.insert_weak(o.clone()));
634
                t.push(o);
635
            }
636
            let obj = Arc::new(ExampleObject("cafe".into()));
637
            keep_these.push(obj.clone());
638
            s.push(map.insert_weak(obj));
639
            drop(t);
640
            map.assert_okay();
641
        }
642

            
643
        assert_eq!(s.len(), 100);
644
        assert_eq!(w.len(), 1000);
645
        assert!(w.iter().all(|id| map.lookup(*id).is_none()));
646
        assert!(s.iter().all(|id| map.lookup(*id).is_some()));
647

            
648
        assert_ne!(map.weak_arena.len() + map.strong_arena.len(), 1100);
649
        map.assert_okay();
650
        map.tidy();
651
        map.assert_okay();
652
        assert_eq!(map.weak_arena.len() + map.strong_arena.len(), 100);
653

            
654
        // This number is a bit arbitrary.
655
        assert!(dbg!(map.n_tidies) < 30);
656
    }
657

            
658
    #[test]
659
    fn wrapper_magic() {
660
        // Make sure that the wrapper transmutation trick works well.
661
        let obj = Arc::new(ExampleObject("dump".into()));
662
        let wrap = obj.clone().wrap_arc();
663

            
664
        let mut map = ObjMap::new();
665
        map.insert_strong(obj);
666
        map.insert_strong(wrap);
667
        assert_eq!(map.strong_arena.len(), 2);
668
    }
669

            
670
    #[test]
671
    fn objid_encoding() {
672
        use rand::Rng;
673
        fn test_roundtrip(a: u32, b: u32, rng: &mut tor_basic_utils::test_rng::TestingRng) {
674
            let a: u64 = a.into();
675
            let b: u64 = b.into();
676
            let data = KeyData::from_ffi((a << 33) | 1_u64 << 32 | b);
677
            let idx = if rng.gen_bool(0.5) {
678
                GenIdx::Strong(StrongIdx::from(data))
679
            } else {
680
                GenIdx::Weak(WeakIdx::from(data))
681
            };
682
            let s1 = idx.encode_with_rng(rng);
683
            let s2 = idx.encode_with_rng(rng);
684
            assert_ne!(s1, s2);
685
            assert_eq!(idx, GenIdx::try_decode(&s1).unwrap());
686
            assert_eq!(idx, GenIdx::try_decode(&s2).unwrap());
687
        }
688
        let mut rng = tor_basic_utils::test_rng::testing_rng();
689

            
690
        test_roundtrip(0, 1, &mut rng);
691
        test_roundtrip(0, 2, &mut rng);
692
        test_roundtrip(1, 1, &mut rng);
693
        test_roundtrip(0xffffffff, 0xffffffff, &mut rng);
694

            
695
        for _ in 0..256 {
696
            test_roundtrip(rng.gen(), rng.gen(), &mut rng);
697
        }
698
    }
699
}