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
pub(crate) mod methods;
18

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
223
1048
        bytes.try_into().expect("Length was wrong!")
224
1048
    }
225

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

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

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

            
245
1044
        let ffi_idx = ffi_idx.wrapping_sub(x);
246
1044

            
247
1044
        if is_weak {
248
430
            Some(GenIdx::Weak(WeakIdx::from(KeyData::from_ffi(ffi_idx))))
249
        } else {
250
614
            Some(GenIdx::Strong(StrongIdx::from(KeyData::from_ffi(ffi_idx))))
251
        }
252
1044
    }
253
}
254

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

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

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

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

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

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

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

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

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

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

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

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

            
390
    use super::*;
391
    use derive_deftly::Deftly;
392
    use tor_rpcbase::templates::*;
393

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

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

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

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

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

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

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

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

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

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

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

            
465
        assert_eq!(
466
            TaggedAddr::for_object(&wrapped_dyn),
467
            TaggedAddr::for_object(&wrapped_dyn2)
468
        );
469

            
470
        assert_ne!(
471
            TaggedAddr::for_object(&object_dyn),
472
            TaggedAddr::for_object(&wrapped_dyn)
473
        );
474

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
692
        test_roundtrip(0, 1, &mut rng);
693
        test_roundtrip(0, 2, &mut rng);
694
        test_roundtrip(1, 1, &mut rng);
695
        test_roundtrip(0xffffffff, 0xffffffff, &mut rng);
696

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