Skip to main content

snap_control/pg_wap/
grants.rs

1// Copyright 2026 Anapaya Systems
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//   http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! IP based grants.
16//!
17//! A grant is a promise that a client IP may use a target, and a set of private segments
18//! towards it, until a point in time:
19//!
20//! - **Additive**: [`GrantManager::authorize`] adds targets and segments to whatever the IP already
21//!   has, and never shortens an existing grant.
22//! - **Bounded**: an IP holds at most [`GrantManagerConfig::max_targets_per_ip`] target grants, and
23//!   a target holds at most [`GrantManagerConfig::max_segments_per_target`] segment grants. Grants
24//!   over a limit evict unguarded grants, or fail.
25//! - **Not revocable**: apart from that eviction, state leaves the manager only via
26//!   [`GrantManager::clean`], and only once it has expired.
27//! - **Observable**: every grant carries a [`CancellationToken`] that is cancelled when the grant
28//!   is removed, so consumers can tear down or re-path. Cancellation is latched, so a waiter that
29//!   has not been polled yet still observes it.
30//!
31//! A segment grant never outlives the segment it was granted on: `granted_until` is capped at
32//! the segment's own expiry. Without that cap an expired segment would have to be pulled out
33//! from under a live grant, which is indistinguishable from a revocation.
34//!
35//! The invariants we maintain are:
36//! - A segment grant never outlives the segment it was granted on.
37//! - A segment is removed from the authoritative store once its last grant is gone.
38//! - If a segment was granted for an IP and an SNI, it cannot be used for another IP or SNI.
39//! - As soon as a grant is removed, it's connected cancellation token is cancelled.
40//! - An IP never holds more than `max_targets_per_ip` target grants.
41//! - A target never holds more than `max_segments_per_target` segment grants.
42//! - A grant with a live [`SegmentGrantGuard`] or [`TargetGrantGuard`] is never evicted.
43
44use std::{
45    collections::{BTreeMap, HashMap, hash_map::Entry},
46    net::IpAddr,
47    sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
48    time::{Duration, SystemTime},
49};
50
51use sciparse::segment::{SegmentFp, SignedPathSegment};
52use serde::{Deserialize, Serialize};
53use serde_with::{DurationSeconds, serde_as};
54use tokio_util::sync::CancellationToken;
55
56use crate::pg_wap::{
57    crpc::model::AuthSegments,
58    sni::{CustomerDomain, CustomerDomainRef},
59};
60
61/// Configuration for a [`GrantManager`].
62#[serde_as]
63#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
64#[serde(deny_unknown_fields)]
65pub struct GrantManagerConfig {
66    /// How long a grant is handed out for, after which the client has to re-authorize.
67    #[serde_as(as = "DurationSeconds<u64>")]
68    #[serde(default = "GrantManagerConfig::default_grant_duration")]
69    pub grant_duration: Duration,
70    /// Shortest sleep between two [`GrantManager::clean`] runs.
71    #[serde_as(as = "DurationSeconds<u64>")]
72    #[serde(default = "GrantManagerConfig::default_min_clean_interval")]
73    pub min_clean_interval: Duration,
74    /// Longest sleep between two [`GrantManager::clean`] runs.
75    #[serde_as(as = "DurationSeconds<u64>")]
76    #[serde(default = "GrantManagerConfig::default_max_clean_interval")]
77    pub max_clean_interval: Duration,
78    /// Maximum total number of segments a client may grant in one authorization request.
79    #[serde(default = "GrantManagerConfig::default_max_segments_per_request")]
80    pub max_segments_per_request: usize,
81    /// Maximum number of segment grants one target holds per authorized IP.
82    ///
83    /// Grants over the limit are evicted by [`GrantManager::authorize`].
84    #[serde(default = "GrantManagerConfig::default_max_segments_per_target")]
85    pub max_segments_per_target: usize,
86    /// Maximum number of targets a client may grant in one authorization request.
87    #[serde(default = "GrantManagerConfig::default_max_targets_per_request")]
88    pub max_targets_per_request: usize,
89    /// Maximum number of target grants one IP holds.
90    ///
91    /// Grants over the limit are evicted by [`GrantManager::authorize`].
92    #[serde(default = "GrantManagerConfig::default_max_targets_per_ip")]
93    pub max_targets_per_ip: usize,
94}
95
96impl GrantManagerConfig {
97    fn default_grant_duration() -> Duration {
98        Duration::from_secs(2 * 60)
99    }
100
101    fn default_min_clean_interval() -> Duration {
102        Duration::from_secs(10)
103    }
104
105    fn default_max_clean_interval() -> Duration {
106        Duration::from_secs(30)
107    }
108
109    fn default_max_segments_per_request() -> usize {
110        100
111    }
112
113    fn default_max_segments_per_target() -> usize {
114        10
115    }
116
117    fn default_max_targets_per_request() -> usize {
118        10
119    }
120
121    fn default_max_targets_per_ip() -> usize {
122        50
123    }
124
125    /// Validates the configuration, returning an error if it is invalid.
126    pub fn validate(&self) -> anyhow::Result<()> {
127        if self.min_clean_interval > self.max_clean_interval {
128            anyhow::bail!(
129                "min_clean_interval ({:?}) is greater than max_clean_interval ({:?})",
130                self.min_clean_interval,
131                self.max_clean_interval
132            );
133        }
134
135        if self.max_segments_per_request == 0 {
136            anyhow::bail!("max_segments_per_request must be greater than 0");
137        }
138
139        if self.max_segments_per_target == 0 {
140            anyhow::bail!("max_segments_per_target must be greater than 0");
141        }
142
143        if self.max_targets_per_request == 0 {
144            anyhow::bail!("max_targets_per_request must be greater than 0");
145        }
146
147        if self.max_targets_per_ip == 0 {
148            anyhow::bail!("max_targets_per_ip must be greater than 0");
149        }
150
151        if self.grant_duration == Duration::ZERO {
152            anyhow::bail!("grant_duration must be greater than 0");
153        }
154
155        if self.min_clean_interval == Duration::ZERO {
156            anyhow::bail!("min_clean_interval must be greater than 0");
157        }
158
159        if self.max_clean_interval == Duration::ZERO {
160            anyhow::bail!("max_clean_interval must be greater than 0");
161        }
162
163        Ok(())
164    }
165}
166
167impl Default for GrantManagerConfig {
168    fn default() -> Self {
169        Self {
170            grant_duration: Self::default_grant_duration(),
171            min_clean_interval: Self::default_min_clean_interval(),
172            max_clean_interval: Self::default_max_clean_interval(),
173            max_segments_per_request: Self::default_max_segments_per_request(),
174            max_segments_per_target: Self::default_max_segments_per_target(),
175            max_targets_per_request: Self::default_max_targets_per_request(),
176            max_targets_per_ip: Self::default_max_targets_per_ip(),
177        }
178    }
179}
180
181/// Tracks which client IPs are authorized for which targets and segments.
182#[derive(Clone)]
183pub struct GrantManager(Arc<GrantManagerShared>);
184
185struct GrantManagerShared {
186    /// The mutable authorization state.
187    inner: RwLock<GrantManagerInner>,
188    config: GrantManagerConfig,
189}
190
191/// The mutable state of a [`GrantManager`].
192pub struct GrantManagerInner {
193    /// Map of all authenticated IPs.
194    auth_ips: HashMap<IpAddr, IpAuthInfo>,
195
196    /// Authoritative, deduplicated list of all granted non-core segments for all IPs and dst
197    /// domains.
198    ///
199    /// Segments are removed once their last grant is gone.
200    private_non_core_segments: HashMap<SegmentFp, AuthSegmentEntry>,
201    /// Same as `private_non_core_segments`, but for core segments.
202    private_core_segments: HashMap<SegmentFp, AuthSegmentEntry>,
203}
204
205impl GrantManager {
206    /// Creates an empty service handing out grants as `config` describes.
207    ///
208    /// Returns an error if `config` does not [`GrantManagerConfig::validate`].
209    pub fn new(config: GrantManagerConfig) -> anyhow::Result<Self> {
210        config.validate()?;
211
212        Ok(Self(Arc::new(GrantManagerShared {
213            inner: RwLock::new(GrantManagerInner {
214                auth_ips: HashMap::new(),
215                private_non_core_segments: HashMap::new(),
216                private_core_segments: HashMap::new(),
217            }),
218            config,
219        })))
220    }
221
222    fn read(&self) -> RwLockReadGuard<'_, GrantManagerInner> {
223        self.0.inner.read().unwrap()
224    }
225
226    fn write(&self) -> RwLockWriteGuard<'_, GrantManagerInner> {
227        self.0.inner.write().unwrap()
228    }
229
230    /// Checks if the ip has any kind of grant.
231    pub fn ip_is_authorized(&self, ip: IpAddr, now: SystemTime) -> bool {
232        let this = self.read();
233
234        this.auth_ips
235            .get(&ip)
236            .is_some_and(|ip_auth| ip_auth.auth_dsts.values().any(|dst| dst.is_valid(now)))
237    }
238
239    /// Authorizes the given IP address for the given destinations and segments.
240    ///
241    /// Grants are additive: after applying the operation, the set of authorized destinations and
242    /// segments is a superset of the previous set. The expiration of an authorized segment is
243    /// the maximum of an already existing authorization and the provided. Already expired
244    /// segments are not inserted.
245    ///
246    /// A segment grant is capped at the expiry of the segment it is granted on, so refreshing
247    /// only extends it as far as the freshest copy of that segment allows.
248    ///
249    /// An IP holds at most [`GrantManagerConfig::max_targets_per_ip`] destination grants, and a
250    /// destination holds at most [`GrantManagerConfig::max_segments_per_target`] segment grants.
251    /// A request that does not fit evicts destination grants without a [`TargetGrantGuard`], and
252    /// segment grants without a [`SegmentGrantGuard`]. Grants with the least time left are evicted
253    /// first. Evicting a destination grant removes its segment grants as well.
254    ///
255    /// The request is rejected if eviction can not free enough room, or if it supplies too many
256    /// destinations or segments at once. A rejection does not modify the state of the manager.
257    ///
258    /// Returns the time the destination grants hold until, which is when the client has to
259    /// authorize again. A segment grant capped at its segment's expiry ends before that.
260    ///
261    /// Note: Segments are currently not being validated in this function.
262    pub fn authorize(
263        &self,
264        ip: IpAddr,
265        destinations: BTreeMap<CustomerDomain, AuthSegments>,
266        now: SystemTime,
267    ) -> Result<SystemTime, AuthorizeError> {
268        if destinations.len() > self.0.config.max_targets_per_request {
269            return Err(AuthorizeError::TooManyTargetsInRequest {
270                requested: destinations.len(),
271                max: self.0.config.max_targets_per_request,
272            });
273        }
274
275        if destinations.len() > self.0.config.max_targets_per_ip {
276            return Err(AuthorizeError::TooManyTargetsForIp {
277                requested: destinations.len(),
278                max: self.0.config.max_targets_per_ip,
279            });
280        }
281
282        let requested = destinations
283            .values()
284            .map(AuthSegments::count)
285            .sum::<usize>();
286
287        if requested > self.0.config.max_segments_per_request {
288            return Err(AuthorizeError::TooManySegmentsInRequest {
289                requested,
290                max: self.0.config.max_segments_per_request,
291            });
292        }
293
294        // TODO: Auth segments are client input and should be validated.
295
296        let dst_granted_until = now + self.0.config.grant_duration;
297
298        let mut this = self.write();
299
300        // Plan against the current state, before changing any of it.
301        let evicted_targets =
302            plan_targets(&this, self.0.config.max_targets_per_ip, ip, &destinations)?;
303
304        let plans = destinations
305            .into_iter()
306            .map(|(dst, auth_segments)| {
307                plan_destination(
308                    &this,
309                    self.0.config.max_segments_per_target,
310                    ip,
311                    dst,
312                    auth_segments,
313                    dst_granted_until,
314                    now,
315                )
316            })
317            .collect::<Result<Vec<_>, _>>()?;
318
319        let GrantManagerInner {
320            auth_ips,
321            private_non_core_segments,
322            private_core_segments,
323        } = &mut *this;
324
325        let auth_entry = auth_ips.entry(ip).or_default();
326
327        for dst in evicted_targets {
328            evict_target_grant(
329                &mut auth_entry.auth_dsts,
330                private_core_segments,
331                private_non_core_segments,
332                &dst,
333            );
334            tracing::debug!(%ip, %dst, "Evicted an unused target grant");
335        }
336
337        for DstPlan {
338            dst,
339            grants,
340            evictions,
341        } in plans
342        {
343            // Get the grant entry for this (ip, dst) pair, or create a new one if it doesn't
344            // exist. An existing grant is only ever extended.
345            let dst_entry = auth_entry
346                .auth_dsts
347                .entry(dst.clone())
348                .and_modify(|dst_auth| {
349                    dst_auth.granted_until = dst_auth.granted_until.max(dst_granted_until);
350                })
351                .or_insert_with(|| DstAuthInfo::new(dst_granted_until));
352
353            // Evict grants
354            for id in evictions {
355                evict_segment_grant(
356                    &mut dst_entry.segment_grants,
357                    segment_store(id, private_core_segments, private_non_core_segments),
358                    id,
359                );
360                tracing::debug!(%ip, %dst, fp = %id.fp(), "Evicted an unused segment grant");
361            }
362
363            for PlannedGrant {
364                id,
365                segment,
366                segment_expiry,
367                granted_until,
368            } in grants
369            {
370                grant_segment(
371                    &mut dst_entry.segment_grants,
372                    segment_store(id, private_core_segments, private_non_core_segments),
373                    id,
374                    segment,
375                    segment_expiry,
376                    granted_until,
377                );
378            }
379        }
380
381        Ok(dst_granted_until)
382    }
383
384    /// Checks if the IP has a grant for the given SNI.
385    ///
386    /// Returns `None` if no valid grant exists.
387    /// Returns `DstGrant` with the segments currently granted for the given SNI if a valid
388    /// grant exists.
389    pub fn dst_grant(
390        &self,
391        ip: IpAddr,
392        dst: CustomerDomainRef<'_>,
393        now: SystemTime,
394    ) -> Option<DstGrant> {
395        let this = self.read();
396
397        let dst_auth = this.auth_ips.get(&ip)?.auth_dsts.get(dst.as_str())?;
398
399        // Grant has already expired, return None to indicate that the client needs to
400        // reauthenticate.
401        if !dst_auth.is_valid(now) {
402            return None;
403        }
404
405        let mut granted_segments = DstGrant::default();
406
407        for (id, grant) in dst_auth.segment_grants.iter() {
408            // Skip expired grants, they will be cleaned up later.
409            if !grant.is_valid(now) {
410                continue;
411            }
412
413            let (store, out) = match id {
414                GrantedSegmentId::Core(_) => {
415                    (
416                        &this.private_core_segments,
417                        &mut granted_segments.core_segments,
418                    )
419                }
420                GrantedSegmentId::NonCore(_) => {
421                    (
422                        &this.private_non_core_segments,
423                        &mut granted_segments.non_core_segments,
424                    )
425                }
426            };
427
428            let Some(segment_entry) = store.get(&id.fp()) else {
429                continue;
430            };
431
432            out.insert(id.fp(), segment_entry.segment.clone());
433        }
434
435        Some(granted_segments)
436    }
437
438    /// Returns the time when the grant for the given IP and destination expires
439    /// Returns `None` if there is no valid grant for the given IP and destination.
440    pub fn grant_expiry(
441        &self,
442        ip: IpAddr,
443        dst: CustomerDomainRef<'_>,
444        now: SystemTime,
445    ) -> Option<SystemTime> {
446        let this = self.read();
447
448        this.auth_ips
449            .get(&ip)?
450            .auth_dsts
451            .get(dst.as_str())
452            .filter(|dst_auth| dst_auth.is_valid(now))
453            .map(|dst_auth| dst_auth.granted_until)
454    }
455
456    /// Returns when the given IP's grant on the given segment for the given destination expires,
457    /// if that grant is valid.
458    pub fn segment_grant_expiry(
459        &self,
460        ip: IpAddr,
461        dst: CustomerDomainRef<'_>,
462        id: &GrantedSegmentId,
463        now: SystemTime,
464    ) -> Option<SystemTime> {
465        let this = self.read();
466
467        this.auth_ips
468            .get(&ip)?
469            .auth_dsts
470            .get(dst.as_str())?
471            .segment_grants
472            .get(id)
473            .filter(|grant| grant.is_valid(now))
474            .map(|grant| grant.granted_until)
475    }
476
477    /// Guards the grant for the given IP and destination against eviction.
478    ///
479    /// Returns `None` if there is no valid grant to watch.
480    pub fn watch_grant(
481        &self,
482        ip: IpAddr,
483        customer_domain: CustomerDomainRef<'_>,
484        now: SystemTime,
485    ) -> Option<TargetGrantGuard> {
486        let this = self.read();
487
488        let grant = this
489            .auth_ips
490            .get(&ip)
491            .and_then(|auth_entry| auth_entry.auth_dsts.get(customer_domain.as_str()));
492
493        let Some(grant) = grant else {
494            tracing::debug!(%ip, %customer_domain, "No grant found");
495            return None;
496        };
497
498        if !grant.is_valid(now) {
499            tracing::debug!(%ip, %customer_domain, "Grant has already expired");
500            return None;
501        }
502
503        Some(grant.guard())
504    }
505
506    /// Guards the given IP's grant on the given segment for the given destination.
507    ///
508    /// Returns `None` if there is no valid grant.
509    pub fn watch_segment_grant(
510        &self,
511        ip: IpAddr,
512        customer_domain: CustomerDomainRef<'_>,
513        id: &GrantedSegmentId,
514        now: SystemTime,
515    ) -> Option<SegmentGrantGuard> {
516        let this = self.read();
517
518        let segment_grant = this
519            .auth_ips
520            .get(&ip)?
521            .auth_dsts
522            .get(customer_domain.as_str())?
523            .segment_grants
524            .get(id)?;
525
526        if !segment_grant.is_valid(now) {
527            tracing::debug!(%ip, %customer_domain, ?id, "Segment grant has already expired");
528            return None;
529        }
530
531        Some(segment_grant.guard())
532    }
533
534    /// Returns the segment `id` refers to, as long as any grant on it is still valid.
535    ///
536    /// Returns `None` if the segment is not in the authoritative store, or if it has no valid
537    /// grants.
538    pub fn segment(&self, id: &GrantedSegmentId, now: SystemTime) -> Option<SignedPathSegment> {
539        let this: RwLockReadGuard<'_, GrantManagerInner> = self.read();
540
541        let store = match id {
542            GrantedSegmentId::Core(_) => &this.private_core_segments,
543            GrantedSegmentId::NonCore(_) => &this.private_non_core_segments,
544        };
545
546        store
547            .get(&id.fp())
548            .filter(|segment_entry| segment_entry.expiration > now)
549            .filter(|segment_entry| {
550                debug_assert!(
551                    segment_entry.grant_count > 0,
552                    "segment {id:?} is in the store but has no grants referencing it, it should have been removed by clean()"
553                );
554                segment_entry.grant_count > 0
555            })
556            .map(|segment_entry| segment_entry.segment.clone())
557    }
558
559    /// Drops all expired authentication state and notifies everything that waited on it.
560    ///
561    /// Removal cascades from the outside in:
562    /// - A destination grant is removed once its `granted_until` has passed. Removing it removes
563    ///   all of its segment grants, regardless of their own expiry.
564    /// - A segment grant is removed once its own `granted_until` has passed.
565    /// - An IP is removed once it has no destination grants left, i.e. it is fully de-authed.
566    /// - A segment is removed from the authoritative store once its last grant is gone. It cannot
567    ///   outlive its grants, because a grant is capped at the segment's expiry.
568    ///
569    /// Every removed grant has its cancellation token cancelled before it is dropped, so the
570    /// streams watching it via [`Self::watch_grant`] or [`Self::watch_segment_grant`] can tear
571    /// down.
572    ///
573    /// Returns the time the next grant expires, i.e. when calling this again has an effect.
574    /// The functions should be called regularely, to ensure newly added grants are cleaned up
575    /// in time.
576    pub fn clean(&self, now: SystemTime) -> SystemTime {
577        let mut next_expiry = now + self.0.config.grant_duration;
578
579        let mut this = self.write();
580        let GrantManagerInner {
581            auth_ips,
582            private_non_core_segments,
583            private_core_segments,
584        } = &mut *this;
585
586        auth_ips.retain(|ip, ip_auth| {
587            ip_auth.auth_dsts.retain(|dst, dst_auth| {
588                let dst_expired = !dst_auth.is_valid(now);
589
590                dst_auth.segment_grants.retain(|id, segment_grant| {
591                    // Retain if the dst is not expired and the segment grant is not expired.
592                    // Otherwise, remove it.
593                    if !dst_expired && segment_grant.is_valid(now) {
594                        next_expiry = next_expiry.min(segment_grant.granted_until);
595                        return true;
596                    }
597
598                    // Otherwise, clean up the grant.
599                    let store = match id {
600                        GrantedSegmentId::Core(_) => &mut *private_core_segments,
601                        GrantedSegmentId::NonCore(_) => &mut *private_non_core_segments,
602                    };
603
604                    release_segment_grant(store, id.fp());
605                    segment_grant.expired.cancel();
606
607                    tracing::trace!(%ip, %dst, fp = %id.fp(), "Removed expired segment grant");
608                    false
609                });
610
611                if dst_expired {
612                    dst_auth.expired.cancel();
613                    tracing::debug!(%ip, %dst, "Removed expired destination grant");
614                } else {
615                    next_expiry = next_expiry.min(dst_auth.granted_until);
616                }
617
618                !dst_expired
619            });
620
621            let ip_authorized = !ip_auth.auth_dsts.is_empty();
622            if !ip_authorized {
623                tracing::debug!(%ip, "IP has no grants left, de-authorizing");
624            }
625
626            ip_authorized
627        });
628
629        // Drop segments that lost their last grant above.
630        for store in [private_core_segments, private_non_core_segments] {
631            store.retain(|fp, segment_entry| {
632                if segment_entry.grant_count > 0 {
633                    debug_assert!(
634                        segment_entry.expiration > now,
635                        "segment {fp} outlived by a grant, grants must be capped at the \
636                         segment expiry"
637                    );
638                    return true;
639                }
640
641                tracing::debug!(%fp, "Removed unreferenced segment from the authoritative store");
642                false
643            });
644        }
645
646        next_expiry
647    }
648
649    /// Periodically runs [`Self::clean`] so grant expiry is observed close to when it happens.
650    ///
651    /// Reads the wall clock itself, unlike every other entry point. Never returns.
652    pub async fn run(&self) {
653        loop {
654            let now = SystemTime::now();
655            let next = self.clean(now);
656            let sleep = next
657                .duration_since(now)
658                .unwrap_or(Duration::ZERO)
659                .max(self.0.config.min_clean_interval)
660                .min(self.0.config.max_clean_interval);
661
662            tracing::trace!("Next grant manager cleanup in {:?}", sleep);
663            tokio::time::sleep(sleep).await;
664        }
665    }
666}
667
668/// Picks the target grants a [`GrantManager::authorize`] call has to evict to fit `max_targets`.
669///
670/// Changes nothing, so a caller can plan every destination before applying any of them.
671///
672/// The IP holds the requested targets plus the ones the request does not name. Targets over
673/// `max_targets` are evicted, by these rules:
674/// - A requested target is never evicted. It is refreshed instead.
675/// - A target with a live [`TargetGrantGuard`], or with a segment grant that a consumer holds, is
676///   never evicted.
677/// - Of the rest, the ones that expire first are evicted, until the IP fits.
678///
679/// Fails, and plans nothing, if evicting every evictable target still does not fit. Returns
680/// [`AuthorizeError::IpTargetLimitReached`].
681fn plan_targets(
682    inner: &GrantManagerInner,
683    max_targets: usize,
684    ip: IpAddr,
685    requested: &BTreeMap<CustomerDomain, AuthSegments>,
686) -> Result<Vec<CustomerDomain>, AuthorizeError> {
687    let Some(ip_auth) = inner.auth_ips.get(&ip) else {
688        return Ok(Vec::new());
689    };
690
691    // The targets the request refreshes are already counted in `requested`.
692    let untouched = ip_auth
693        .auth_dsts
694        .keys()
695        .filter(|dst| !requested.contains_key(*dst))
696        .count();
697    let over = (untouched + requested.len()).saturating_sub(max_targets);
698
699    if over == 0 {
700        return Ok(Vec::new());
701    }
702
703    // The evictable targets, ordered by the rules above.
704    let mut candidates = ip_auth
705        .auth_dsts
706        .iter()
707        .filter(|(dst, dst_auth)| !requested.contains_key(*dst) && !dst_auth.is_in_use())
708        .map(|(dst, dst_auth)| (dst_auth.granted_until, dst))
709        .collect::<Vec<_>>();
710
711    if candidates.len() < over {
712        return Err(AuthorizeError::IpTargetLimitReached {
713            in_use: untouched - candidates.len(),
714            needed: over,
715            max: max_targets,
716        });
717    }
718
719    candidates.sort_unstable_by(|(left_expiry, left_dst), (right_expiry, right_dst)| {
720        left_expiry
721            .cmp(right_expiry)
722            .then_with(|| left_dst.as_str().cmp(right_dst.as_str()))
723    });
724
725    Ok(candidates
726        .into_iter()
727        .take(over)
728        .map(|(_, dst)| dst.clone())
729        .collect())
730}
731
732/// Plans the changes one destination of a [`GrantManager::authorize`] call turns into.
733///
734/// Changes nothing, so a caller can plan every destination before applying any of them.
735///
736/// The requested segments become the wanted grants:
737/// - Up and down segments become non-core grants, core segments become core grants.
738/// - The same [`GrantedSegmentId`] wanted more than once is one grant, on the longest lived copy.
739/// - A segment that is already expired is dropped.
740///
741/// The destination then holds the wanted grants plus the ones the request does not name.
742/// Grants over `max_grants` are evicted, by these rules:
743/// - A wanted grant is never evicted. It is refreshed instead.
744/// - A grant with a live [`SegmentGrantGuard`] is never evicted.
745/// - Of the rest, the ones that expire first are evicted, until the destination fits.
746///
747/// Fails, and plans nothing, if:
748/// - The request wants more than `max_grants` grants. Returns
749///   [`AuthorizeError::TooManySegmentsForTarget`].
750/// - Evicting every evictable grant still does not fit. Returns
751///   [`AuthorizeError::TargetGrantLimitReached`].
752fn plan_destination(
753    inner: &GrantManagerInner,
754    max_grants: usize,
755    ip: IpAddr,
756    dst: CustomerDomain,
757    auth_segments: AuthSegments,
758    dst_granted_until: SystemTime,
759    now: SystemTime,
760) -> Result<DstPlan, AuthorizeError> {
761    let AuthSegments {
762        up_segments,
763        down_segments,
764        core_segments,
765    } = auth_segments;
766
767    // Up and down segments both become non-core grants, see `GrantedSegmentId`.
768    let tagged = core_segments
769        .into_iter()
770        .map(|segment| (segment, true))
771        .chain(
772            up_segments
773                .into_iter()
774                .chain(down_segments)
775                .map(|segment| (segment, false)),
776        );
777
778    // Deduplicate segments.
779    let mut wanted: HashMap<GrantedSegmentId, PlannedGrant> = HashMap::new();
780
781    for (segment, is_core) in tagged {
782        let fp = segment.fingerprint();
783        let segment_expiry = segment.expires_earliest();
784
785        // A grant can never outlive the segment it is granted on, see the module docs.
786        let granted_until = dst_granted_until.min(segment_expiry);
787        if granted_until <= now {
788            tracing::debug!(
789                %ip, %dst, %fp,
790                "Ignoring grant for a segment that has already expired"
791            );
792            continue;
793        }
794
795        let id = if is_core {
796            GrantedSegmentId::Core(fp)
797        } else {
798            GrantedSegmentId::NonCore(fp)
799        };
800
801        let planned = PlannedGrant {
802            id,
803            segment,
804            segment_expiry,
805            granted_until,
806        };
807
808        match wanted.entry(id) {
809            // Keep the copy that lives longest.
810            Entry::Occupied(mut existing) => {
811                if existing.get().segment_expiry < segment_expiry {
812                    existing.insert(planned);
813                }
814            }
815            Entry::Vacant(vacant) => {
816                vacant.insert(planned);
817            }
818        }
819    }
820
821    if wanted.len() > max_grants {
822        return Err(AuthorizeError::TooManySegmentsForTarget {
823            dst,
824            requested: wanted.len(),
825            max: max_grants,
826        });
827    }
828
829    let held = inner
830        .auth_ips
831        .get(&ip)
832        .and_then(|ip_auth| ip_auth.auth_dsts.get(&dst));
833
834    let mut evictions = Vec::new();
835
836    if let Some(dst_auth) = held {
837        // The grants the request refreshes are already counted in `wanted`.
838        let untouched = dst_auth
839            .segment_grants
840            .keys()
841            .filter(|id| !wanted.contains_key(id))
842            .count();
843        let over = (untouched + wanted.len()).saturating_sub(max_grants);
844
845        if over > 0 {
846            // The evictable grants, ordered by the rules above.
847            let mut candidates = dst_auth
848                .segment_grants
849                .iter()
850                .filter(|(id, grant)| !wanted.contains_key(id) && !grant.is_in_use())
851                .map(|(id, grant)| (grant.granted_until, *id))
852                .collect::<Vec<_>>();
853
854            if candidates.len() < over {
855                return Err(AuthorizeError::TargetGrantLimitReached {
856                    dst,
857                    in_use: untouched - candidates.len(),
858                    needed: over,
859                    max: max_grants,
860                });
861            }
862
863            candidates.sort_unstable();
864            evictions.extend(candidates.into_iter().take(over).map(|(_, id)| id));
865        }
866    }
867
868    Ok(DstPlan {
869        dst,
870        grants: wanted.into_values().collect(),
871        evictions,
872    })
873}
874
875/// Picks the authoritative store `id` belongs to.
876fn segment_store<'store>(
877    id: GrantedSegmentId,
878    core: &'store mut HashMap<SegmentFp, AuthSegmentEntry>,
879    non_core: &'store mut HashMap<SegmentFp, AuthSegmentEntry>,
880) -> &'store mut HashMap<SegmentFp, AuthSegmentEntry> {
881    if id.is_core() { core } else { non_core }
882}
883
884/// Records a grant for `id` and inserts or refreshes the granted segment.
885///
886/// - The grant for `id` is only ever extended, never shortened.
887/// - The stored copy of the segment is replaced only by a longer lived one.
888/// - The grant count grows by one only if this call adds a grant that did not exist before.
889fn grant_segment(
890    grants: &mut HashMap<GrantedSegmentId, SegmentGrant>,
891    store: &mut HashMap<SegmentFp, AuthSegmentEntry>,
892    id: GrantedSegmentId,
893    segment: SignedPathSegment,
894    segment_expiry: SystemTime,
895    granted_until: SystemTime,
896) {
897    let is_new_grant = match grants.entry(id) {
898        // Extend the existing grant.
899        Entry::Occupied(mut existing) => {
900            let grant = existing.get_mut();
901            grant.granted_until = grant.granted_until.max(granted_until);
902            false
903        }
904        // Insert a new grant.
905        Entry::Vacant(vacant) => {
906            vacant.insert(SegmentGrant::new(granted_until));
907            true
908        }
909    };
910
911    match store.entry(id.fp()) {
912        // Segment already exists, extend its expiry, and increase the grant count if a new
913        // grant was added above.
914        Entry::Occupied(mut existing) => {
915            let existing = existing.get_mut();
916
917            // Keep the copy that lives longest, a refresh may carry a fresher one.
918            if existing.expiration < segment_expiry {
919                existing.segment = segment;
920                existing.expiration = segment_expiry;
921            }
922
923            if is_new_grant {
924                existing.grant_count += 1;
925            }
926        }
927        // Segment does not exist yet, insert it with a grant count of 1. For the new grant
928        // added above
929        Entry::Vacant(vacant) => {
930            debug_assert!(
931                is_new_grant,
932                "segment store was missing a segment that was already referenced by a grant"
933            );
934            vacant.insert(AuthSegmentEntry::new(segment, segment_expiry, 1));
935        }
936    }
937}
938
939/// Drops the grant on `id` before it expired.
940///
941/// - Cancels the grant's token.
942/// - Releases the grant's hold on the segment.
943/// - Drops the segment from the store if it lost its last grant.
944fn evict_segment_grant(
945    grants: &mut HashMap<GrantedSegmentId, SegmentGrant>,
946    store: &mut HashMap<SegmentFp, AuthSegmentEntry>,
947    id: GrantedSegmentId,
948) {
949    let Some(grant) = grants.remove(&id) else {
950        debug_assert!(
951            false,
952            "grant {id:?} was planned for eviction but is not there"
953        );
954        return;
955    };
956
957    drop_segment_grant(store, id.fp(), grant);
958}
959
960/// Drops the target grant on `dst`, and all of its segment grants, before they expired.
961///
962/// Cancels the token of the target grant and of every segment grant it held.
963fn evict_target_grant(
964    auth_dsts: &mut HashMap<CustomerDomain, DstAuthInfo>,
965    core: &mut HashMap<SegmentFp, AuthSegmentEntry>,
966    non_core: &mut HashMap<SegmentFp, AuthSegmentEntry>,
967    dst: &CustomerDomain,
968) {
969    let Some(dst_auth) = auth_dsts.remove(dst) else {
970        debug_assert!(
971            false,
972            "target {dst} was planned for eviction but is not there"
973        );
974        return;
975    };
976
977    for (id, grant) in dst_auth.segment_grants {
978        drop_segment_grant(segment_store(id, core, non_core), id.fp(), grant);
979    }
980
981    dst_auth.expired.cancel();
982}
983
984/// Releases `grant`'s hold on the segment `fp` refers to and cancels the grant.
985///
986/// Drops the segment from the store if it lost its last grant. The caller has already taken the
987/// grant out of the map that held it.
988fn drop_segment_grant(
989    store: &mut HashMap<SegmentFp, AuthSegmentEntry>,
990    fp: SegmentFp,
991    grant: SegmentGrant,
992) {
993    release_segment_grant(store, fp);
994
995    // Drop the segment right away; there is no sweep after this, unlike in `clean`.
996    if store
997        .get(&fp)
998        .is_some_and(|segment_entry| segment_entry.grant_count == 0)
999    {
1000        store.remove(&fp);
1001        tracing::debug!(%fp, "Removed unreferenced segment from the authoritative store");
1002    }
1003
1004    grant.expired.cancel();
1005}
1006
1007/// Decrements the grant count of the segment `fp` refers to.
1008///
1009/// The entry itself is left in place; it is dropped by the store sweep in
1010/// [`GrantManager::clean`] once its count has reached zero.
1011fn release_segment_grant(store: &mut HashMap<SegmentFp, AuthSegmentEntry>, fp: SegmentFp) {
1012    let Some(segment_entry) = store.get_mut(&fp) else {
1013        debug_assert!(false, "segment {fp} has a grant but is not in the store");
1014        return;
1015    };
1016
1017    debug_assert!(
1018        segment_entry.grant_count > 0,
1019        "segment {fp} is in the store but has no grants referencing it"
1020    );
1021    segment_entry.grant_count = segment_entry.grant_count.saturating_sub(1);
1022}
1023
1024/// Failures of [`GrantManager::authorize`].
1025#[derive(Debug, thiserror::Error)]
1026pub enum AuthorizeError {
1027    /// The request is over [`GrantManagerConfig::max_segments_per_request`].
1028    #[error("request grants {requested} segments, at most {max} are allowed per request")]
1029    TooManySegmentsInRequest {
1030        /// The number of segments the request wants to grant, over all its targets.
1031        requested: usize,
1032        /// The maximum number of segments allowed per request.
1033        max: usize,
1034    },
1035    /// The request is over [`GrantManagerConfig::max_targets_per_request`].
1036    #[error("request has {requested} targets, at most {max} are allowed per request")]
1037    TooManyTargetsInRequest {
1038        /// The number of targets in the request.
1039        requested: usize,
1040        /// The maximum number of targets allowed per request.
1041        max: usize,
1042    },
1043    /// One destination of the request is over [`GrantManagerConfig::max_segments_per_target`].
1044    #[error("request grants {requested} segments for {dst}, at most {max} are allowed per target")]
1045    TooManySegmentsForTarget {
1046        /// The destination that is over the limit.
1047        dst: CustomerDomain,
1048        /// The number of segments the request wants to grant for that destination.
1049        requested: usize,
1050        /// The maximum number of segments allowed per destination.
1051        max: usize,
1052    },
1053    /// One destination is at its limit, and too many of its grants are guarded to evict.
1054    #[error(
1055        "target {dst} needs {needed} of its {max} segment grants freed, but {in_use} of the \
1056         candidates are in use"
1057    )]
1058    TargetGrantLimitReached {
1059        /// The destination that is over the limit.
1060        dst: CustomerDomain,
1061        /// The number of grants that are currently in use and cannot be evicted.
1062        in_use: usize,
1063        /// The number of grants that need to be freed to fit the request.
1064        needed: usize,
1065        /// The maximum number of grants allowed per destination.
1066        max: usize,
1067    },
1068    /// The request alone is over [`GrantManagerConfig::max_targets_per_ip`].
1069    #[error("request has {requested} targets, an IP may hold at most {max}")]
1070    TooManyTargetsForIp {
1071        /// The number of targets the request wants to grant for the IP.
1072        requested: usize,
1073        /// The maximum number of targets allowed per IP.
1074        max: usize,
1075    },
1076    /// The IP is at its limit, and too many of its target grants are in use to evict.
1077    #[error(
1078        "the IP needs {needed} of its {max} target grants freed, but {in_use} of the candidates \
1079         are in use"
1080    )]
1081    IpTargetLimitReached {
1082        /// The number of target grants that are in use and cannot be evicted.
1083        in_use: usize,
1084        /// The number of target grants that have to be freed to fit the request.
1085        needed: usize,
1086        /// The maximum number of target grants allowed per IP.
1087        max: usize,
1088    },
1089}
1090
1091/// The changes [`GrantManager::authorize`] applies to one destination.
1092struct DstPlan {
1093    dst: CustomerDomain,
1094    /// The grants to add or extend, deduplicated by id.
1095    grants: Vec<PlannedGrant>,
1096    /// The grants to evict first.
1097    evictions: Vec<GrantedSegmentId>,
1098}
1099
1100/// One grant of a [`DstPlan`].
1101struct PlannedGrant {
1102    id: GrantedSegmentId,
1103    segment: SignedPathSegment,
1104    segment_expiry: SystemTime,
1105    granted_until: SystemTime,
1106}
1107
1108/// Identification of a segment of the authoritative store.
1109#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
1110pub enum GrantedSegmentId {
1111    /// A core segment with the given fingerprint.
1112    Core(SegmentFp),
1113    /// An up or down segment with the given fingerprint.
1114    NonCore(SegmentFp),
1115}
1116
1117impl GrantedSegmentId {
1118    /// The fingerprint of the segment.
1119    pub fn fp(&self) -> SegmentFp {
1120        match self {
1121            Self::Core(fp) | Self::NonCore(fp) => *fp,
1122        }
1123    }
1124
1125    /// Returns true if this identifies a core segment.
1126    pub fn is_core(&self) -> bool {
1127        matches!(self, Self::Core(_))
1128    }
1129}
1130
1131/// A segment in the authoritative store, with the number of grants referencing it.
1132struct AuthSegmentEntry {
1133    segment: SignedPathSegment,
1134    expiration: SystemTime,
1135    grant_count: usize,
1136}
1137
1138impl AuthSegmentEntry {
1139    fn new(segment: SignedPathSegment, expiration: SystemTime, initial_grant_count: usize) -> Self {
1140        Self {
1141            segment,
1142            expiration,
1143            grant_count: initial_grant_count,
1144        }
1145    }
1146}
1147
1148/// Everything one IP is authorized for.
1149#[derive(Default)]
1150struct IpAuthInfo {
1151    auth_dsts: HashMap<CustomerDomain, DstAuthInfo>,
1152}
1153
1154/// A grant for one (IP, destination) pair.
1155struct DstAuthInfo {
1156    segment_grants: HashMap<GrantedSegmentId, SegmentGrant>,
1157    granted_until: SystemTime,
1158    /// Cancelled when this grant is removed, either by [`GrantManager::clean`] or by eviction.
1159    expired: CancellationToken,
1160    /// Holds one strong reference per live [`TargetGrantGuard`], plus the one kept here.
1161    in_use: Arc<()>,
1162}
1163
1164impl DstAuthInfo {
1165    fn new(granted_until: SystemTime) -> Self {
1166        Self {
1167            segment_grants: HashMap::new(),
1168            granted_until,
1169            expired: CancellationToken::new(),
1170            in_use: Arc::new(()),
1171        }
1172    }
1173
1174    fn is_valid(&self, now: SystemTime) -> bool {
1175        now < self.granted_until
1176    }
1177
1178    /// Whether a consumer holds this grant, or one of the segment grants it would take with it.
1179    ///
1180    /// May stay true for a moment after the last guard is dropped.
1181    fn is_in_use(&self) -> bool {
1182        Arc::strong_count(&self.in_use) > 1
1183            || self.segment_grants.values().any(SegmentGrant::is_in_use)
1184    }
1185
1186    /// Guards this grant from eviction.
1187    fn guard(&self) -> TargetGrantGuard {
1188        TargetGrantGuard {
1189            expired: self.expired.clone(),
1190            _in_use: Arc::clone(&self.in_use),
1191        }
1192    }
1193}
1194
1195/// A grant for one segment within a [`DstAuthInfo`].
1196struct SegmentGrant {
1197    granted_until: SystemTime,
1198    /// Cancelled when this grant is removed, either by [`GrantManager::clean`] or by eviction.
1199    expired: CancellationToken,
1200    /// Holds one strong reference per live [`SegmentGrantGuard`], plus the one kept here.
1201    in_use: Arc<()>,
1202}
1203
1204impl SegmentGrant {
1205    fn new(granted_until: SystemTime) -> Self {
1206        Self {
1207            granted_until,
1208            expired: CancellationToken::new(),
1209            in_use: Arc::new(()),
1210        }
1211    }
1212
1213    fn is_valid(&self, now: SystemTime) -> bool {
1214        now < self.granted_until
1215    }
1216
1217    /// Whether a [`SegmentGrantGuard`] on this grant is alive.
1218    ///
1219    /// May stay true for a moment after the last guard is dropped.
1220    fn is_in_use(&self) -> bool {
1221        Arc::strong_count(&self.in_use) > 1
1222    }
1223
1224    /// Guards this grant from eviction.
1225    fn guard(&self) -> SegmentGrantGuard {
1226        SegmentGrantGuard {
1227            expired: self.expired.clone(),
1228            _in_use: Arc::clone(&self.in_use),
1229        }
1230    }
1231}
1232
1233/// Guards a target grant from eviction, and allows a client to observe expiry.
1234///
1235/// Handed out by [`GrantManager::watch_grant`].
1236/// Use [`Self::expired`] to await expiry.
1237pub struct TargetGrantGuard {
1238    expired: CancellationToken,
1239    /// Counted by [`DstAuthInfo::is_in_use`] for as long as this guard lives.
1240    _in_use: Arc<()>,
1241}
1242
1243impl TargetGrantGuard {
1244    /// The token that is cancelled once the grant is gone.
1245    ///
1246    /// Cancellation is latched.
1247    pub fn expired(&self) -> &CancellationToken {
1248        &self.expired
1249    }
1250}
1251
1252/// Guards a segment grant from eviction, and allows a client to observe expiry.
1253///
1254/// Handed out by [`GrantManager::watch_segment_grant`].
1255/// Use [`Self::expired`] to await expiry.
1256pub struct SegmentGrantGuard {
1257    expired: CancellationToken,
1258    /// Counted by [`SegmentGrant::is_in_use`] for as long as this guard lives.
1259    _in_use: Arc<()>,
1260}
1261
1262impl SegmentGrantGuard {
1263    /// The token that is cancelled once the grant is gone.
1264    ///
1265    /// Cancellation is latched.
1266    pub fn expired(&self) -> &CancellationToken {
1267        &self.expired
1268    }
1269}
1270
1271// XXX: this should not copy the segments, just an iterator over the granted segments, but that
1272// would require locking magic.
1273/// A grany for one destination in [`GrantManager::dst_grant`].
1274#[derive(Default)]
1275pub struct DstGrant {
1276    core_segments: HashMap<SegmentFp, SignedPathSegment>,
1277    non_core_segments: HashMap<SegmentFp, SignedPathSegment>,
1278}
1279
1280impl DstGrant {
1281    /// Iterates over the granted core segments.
1282    pub fn iter_core_segments(&self) -> impl Iterator<Item = &SignedPathSegment> {
1283        self.core_segments.values()
1284    }
1285
1286    /// Iterates over the granted up and down segments.
1287    pub fn iter_non_core_segments(&self) -> impl Iterator<Item = &SignedPathSegment> {
1288        self.non_core_segments.values()
1289    }
1290
1291    /// Returns whether a core segment with the given fingerprint is granted.
1292    pub fn has_core_segment(&self, fp: SegmentFp) -> bool {
1293        self.core_segments.contains_key(&fp)
1294    }
1295
1296    /// Returns whether an up or down segment with the given fingerprint is granted.
1297    pub fn has_non_core_segment(&self, fp: SegmentFp) -> bool {
1298        self.non_core_segments.contains_key(&fp)
1299    }
1300}
1301
1302#[cfg(test)]
1303mod tests {
1304    use std::collections::BTreeMap;
1305
1306    use super::*;
1307    use crate::pg_wap::test_util::{
1308        Fixture, MockFetcher, at, client_ip, granted_id, nth_sni, nth_up_segment, other_client_ip,
1309        other_sni, other_up_segment, secs_since_epoch, sni, up_segment,
1310    };
1311
1312    #[test]
1313    fn refreshing_a_grant_adds_to_it_without_shortening_it() {
1314        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1315
1316        let first = up_segment(0);
1317        let second = other_up_segment(0);
1318
1319        fixture.grant_non_core(vec![first.clone()], at(0));
1320        let first_expiry = fixture
1321            .grant
1322            .segment_grant_expiry(
1323                client_ip(),
1324                sni().customer_domain(),
1325                &granted_id(&first),
1326                at(0),
1327            )
1328            .expect("the first segment is granted");
1329
1330        // A refresh 10s later that only mentions the second segment must not take the first away.
1331        fixture.grant_non_core(vec![second.clone()], at(10));
1332
1333        assert_eq!(
1334            fixture.grant.segment_grant_expiry(
1335                client_ip(),
1336                sni().customer_domain(),
1337                &granted_id(&first),
1338                at(10)
1339            ),
1340            Some(first_expiry),
1341            "the first grant must be kept, with its original expiry"
1342        );
1343        assert_eq!(
1344            fixture.grant.segment_grant_expiry(
1345                client_ip(),
1346                sni().customer_domain(),
1347                &granted_id(&second),
1348                at(10)
1349            ),
1350            Some(at(110)),
1351            "the second grant runs for the full grant duration from the refresh"
1352        );
1353        assert_eq!(
1354            fixture
1355                .grant
1356                .grant_expiry(client_ip(), sni().customer_domain(), at(10)),
1357            Some(at(110)),
1358            "the target grant is extended by the refresh"
1359        );
1360    }
1361
1362    #[test]
1363    fn a_segment_grant_never_outlives_its_segment() {
1364        // The grant duration is far longer than the segment lives.
1365        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(10 * 24 * 3600));
1366
1367        let segment = up_segment(0);
1368        fixture.grant_non_core(vec![segment.clone()], at(0));
1369
1370        assert_eq!(
1371            fixture.grant.segment_grant_expiry(
1372                client_ip(),
1373                sni().customer_domain(),
1374                &granted_id(&segment),
1375                at(0)
1376            ),
1377            Some(segment.expires_earliest()),
1378            "the grant must be capped at the expiry of the segment it is granted on"
1379        );
1380    }
1381
1382    #[test]
1383    fn a_grant_on_an_expired_segment_is_ignored() {
1384        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1385
1386        let segment = up_segment(0);
1387        let expired = at(secs_since_epoch(segment.expires_earliest()) + 1);
1388        fixture.grant_non_core(vec![segment.clone()], expired);
1389
1390        assert_eq!(
1391            fixture.grant.segment_grant_expiry(
1392                client_ip(),
1393                sni().customer_domain(),
1394                &granted_id(&segment),
1395                expired
1396            ),
1397            None
1398        );
1399    }
1400
1401    #[tokio::test]
1402    async fn expired_grants_are_cleaned_up_and_observable() {
1403        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1404
1405        let segment = up_segment(0);
1406        fixture.grant_non_core(vec![segment.clone()], at(0));
1407
1408        let target_expired = fixture
1409            .grant
1410            .watch_grant(client_ip(), sni().customer_domain(), at(0))
1411            .expect("the target grant is live");
1412        let segment_expired = fixture
1413            .grant
1414            .watch_segment_grant(
1415                client_ip(),
1416                sni().customer_domain(),
1417                &granted_id(&segment),
1418                at(0),
1419            )
1420            .expect("the segment grant is live");
1421
1422        assert!(fixture.grant.ip_is_authorized(client_ip(), at(0)));
1423
1424        // Nothing has expired yet, so cleaning changes nothing.
1425        assert_eq!(
1426            fixture.grant.clean(at(0)),
1427            at(100),
1428            "clean reports when the next grant expires"
1429        );
1430        assert!(
1431            fixture
1432                .grant
1433                .dst_grant(client_ip(), sni().customer_domain(), at(0))
1434                .is_some()
1435        );
1436
1437        fixture.grant.clean(at(101));
1438
1439        // Both notifiers were handed out before the expiry and still resolve.
1440        target_expired.expired().cancelled().await;
1441        segment_expired.expired().cancelled().await;
1442
1443        assert!(
1444            !fixture.grant.ip_is_authorized(client_ip(), at(101)),
1445            "an IP without grants is de-authorized"
1446        );
1447        assert!(
1448            fixture
1449                .grant
1450                .dst_grant(client_ip(), sni().customer_domain(), at(101))
1451                .is_none()
1452        );
1453        assert!(
1454            fixture
1455                .grant
1456                .segment(&granted_id(&segment), at(101))
1457                .is_none()
1458        );
1459        assert!(
1460            fixture
1461                .grant
1462                .watch_grant(client_ip(), sni().customer_domain(), at(101))
1463                .is_none(),
1464            "there is no live grant left to wait on"
1465        );
1466    }
1467
1468    /// One IP with two destination grants: [`sni`] runs until t=100 with a shared and a private
1469    /// segment, [`other_sni`] is refreshed at t=50 and runs until t=150 with only the shared one.
1470    ///
1471    /// Returns the fixture, the shared segment and the segment only [`sni`] is granted.
1472    fn two_destination_fixture() -> (Fixture, SignedPathSegment, SignedPathSegment) {
1473        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1474
1475        let shared = up_segment(0);
1476        let only_sni = other_up_segment(0);
1477
1478        fixture
1479            .grant
1480            .authorize(
1481                client_ip(),
1482                BTreeMap::from([
1483                    (
1484                        sni().customer_domain().into(),
1485                        AuthSegments {
1486                            up_segments: vec![shared.clone(), only_sni.clone()],
1487                            ..AuthSegments::default()
1488                        },
1489                    ),
1490                    (
1491                        other_sni().customer_domain().into(),
1492                        AuthSegments {
1493                            up_segments: vec![shared.clone()],
1494                            ..AuthSegments::default()
1495                        },
1496                    ),
1497                ]),
1498                at(0),
1499            )
1500            .expect("both destinations fit");
1501
1502        // Only the second destination re-authenticates, so the first one lapses at t=100.
1503        fixture
1504            .grant
1505            .authorize(
1506                client_ip(),
1507                BTreeMap::from([(
1508                    other_sni().customer_domain().into(),
1509                    AuthSegments {
1510                        up_segments: vec![shared.clone()],
1511                        ..AuthSegments::default()
1512                    },
1513                )]),
1514                at(50),
1515            )
1516            .expect("the refresh fits");
1517
1518        (fixture, shared, only_sni)
1519    }
1520
1521    #[test]
1522    fn an_expired_destination_takes_only_its_own_grants_with_it() {
1523        let (fixture, shared, only_sni) = two_destination_fixture();
1524
1525        let sni_expired = fixture
1526            .grant
1527            .watch_grant(client_ip(), sni().customer_domain(), at(50))
1528            .expect("the lapsing destination is granted");
1529        let other_expired = fixture
1530            .grant
1531            .watch_grant(client_ip(), other_sni().customer_domain(), at(50))
1532            .expect("the refreshed destination is granted");
1533        let shared_for_sni_expired = fixture
1534            .grant
1535            .watch_segment_grant(
1536                client_ip(),
1537                sni().customer_domain(),
1538                &granted_id(&shared),
1539                at(50),
1540            )
1541            .expect("the shared segment is granted for the lapsing destination");
1542        let shared_for_other_expired = fixture
1543            .grant
1544            .watch_segment_grant(
1545                client_ip(),
1546                other_sni().customer_domain(),
1547                &granted_id(&shared),
1548                at(50),
1549            )
1550            .expect("the shared segment is granted for the refreshed destination");
1551
1552        assert_eq!(
1553            fixture.grant.clean(at(101)),
1554            at(150),
1555            "the remaining destination is the next thing to expire"
1556        );
1557
1558        assert!(
1559            fixture
1560                .grant
1561                .dst_grant(client_ip(), sni().customer_domain(), at(101))
1562                .is_none(),
1563            "the lapsed destination is gone"
1564        );
1565        assert!(sni_expired.expired().is_cancelled());
1566        assert!(
1567            shared_for_sni_expired.expired().is_cancelled(),
1568            "an expiring destination takes its segment grants with it, however long they had left"
1569        );
1570
1571        let other_grant = fixture
1572            .grant
1573            .dst_grant(client_ip(), other_sni().customer_domain(), at(101))
1574            .expect("the refreshed destination is untouched");
1575        assert_eq!(
1576            other_grant
1577                .iter_non_core_segments()
1578                .map(SignedPathSegment::fingerprint)
1579                .collect::<Vec<_>>(),
1580            vec![shared.fingerprint()]
1581        );
1582        assert!(!other_expired.expired().is_cancelled());
1583        assert!(!shared_for_other_expired.expired().is_cancelled());
1584
1585        assert!(
1586            fixture.grant.ip_is_authorized(client_ip(), at(101)),
1587            "an IP with one destination left stays authorized"
1588        );
1589        assert!(
1590            fixture
1591                .grant
1592                .segment(&granted_id(&shared), at(101))
1593                .is_some(),
1594            "the shared segment is still held by the destination that refreshed"
1595        );
1596        assert!(
1597            fixture
1598                .grant
1599                .segment(&granted_id(&only_sni), at(101))
1600                .is_none(),
1601            "the segment only the lapsed destination held lost its last grant"
1602        );
1603    }
1604
1605    #[test]
1606    fn an_ip_is_dropped_with_its_last_destination() {
1607        let (fixture, shared, _only_sni) = two_destination_fixture();
1608
1609        // The first destination goes at t=100, the second one at t=150.
1610        fixture.grant.clean(at(101));
1611        assert!(fixture.grant.ip_is_authorized(client_ip(), at(101)));
1612
1613        fixture.grant.clean(at(151));
1614
1615        assert!(
1616            !fixture.grant.ip_is_authorized(client_ip(), at(151)),
1617            "the IP is de-authorized once its last destination is gone"
1618        );
1619        assert!(
1620            fixture
1621                .grant
1622                .dst_grant(client_ip(), other_sni().customer_domain(), at(151))
1623                .is_none(),
1624            "and so is that last destination"
1625        );
1626        assert!(
1627            fixture
1628                .grant
1629                .segment(&granted_id(&shared), at(151))
1630                .is_none(),
1631            "a segment cannot outlive the last grant referencing it"
1632        );
1633    }
1634
1635    #[test]
1636    fn a_refresh_that_omits_a_segment_keeps_it_until_its_own_grant_expires() {
1637        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1638
1639        let kept = up_segment(0);
1640        let omitted = other_up_segment(0);
1641        fixture.grant_non_core(vec![kept.clone(), omitted.clone()], at(0));
1642
1643        let omitted_expired = fixture
1644            .grant
1645            .watch_segment_grant(
1646                client_ip(),
1647                sni().customer_domain(),
1648                &granted_id(&omitted),
1649                at(0),
1650            )
1651            .expect("both segments are granted");
1652
1653        // The re-grant carries only one of the two segments; the other keeps running out.
1654        fixture.grant_non_core(vec![kept.clone()], at(50));
1655
1656        assert!(
1657            fixture
1658                .grant
1659                .segment(&granted_id(&omitted), at(50))
1660                .is_some(),
1661            "an omitted segment stays usable while its own grant is valid"
1662        );
1663        assert!(!omitted_expired.expired().is_cancelled());
1664
1665        fixture.grant.clean(at(101));
1666
1667        assert!(
1668            omitted_expired.expired().is_cancelled(),
1669            "the omitted segment's grant expires on its own schedule"
1670        );
1671        assert!(
1672            fixture
1673                .grant
1674                .segment(&granted_id(&omitted), at(101))
1675                .is_none(),
1676            "and the segment goes with its last grant"
1677        );
1678
1679        let grant = fixture
1680            .grant
1681            .dst_grant(client_ip(), sni().customer_domain(), at(101))
1682            .expect("the destination was extended by the refresh");
1683        assert_eq!(
1684            grant
1685                .iter_non_core_segments()
1686                .map(SignedPathSegment::fingerprint)
1687                .collect::<Vec<_>>(),
1688            vec![kept.fingerprint()],
1689            "only the refreshed segment is left"
1690        );
1691        assert!(fixture.grant.segment(&granted_id(&kept), at(101)).is_some());
1692    }
1693
1694    #[test]
1695    fn granted_segments_are_returned_per_target() {
1696        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1697
1698        let granted = up_segment(0);
1699        let other = other_up_segment(0);
1700        fixture
1701            .grant
1702            .authorize(
1703                client_ip(),
1704                BTreeMap::from([
1705                    (
1706                        sni().customer_domain().into(),
1707                        AuthSegments {
1708                            up_segments: vec![granted.clone()],
1709                            ..AuthSegments::default()
1710                        },
1711                    ),
1712                    (
1713                        other_sni().customer_domain().into(),
1714                        AuthSegments {
1715                            core_segments: vec![other.clone()],
1716                            ..AuthSegments::default()
1717                        },
1718                    ),
1719                ]),
1720                at(0),
1721            )
1722            .expect("both destinations fit");
1723
1724        let segments = fixture
1725            .grant
1726            .dst_grant(client_ip(), sni().customer_domain(), at(0))
1727            .expect("the target is granted");
1728        assert_eq!(
1729            segments
1730                .iter_non_core_segments()
1731                .map(SignedPathSegment::fingerprint)
1732                .collect::<Vec<_>>(),
1733            vec![granted.fingerprint()]
1734        );
1735        assert_eq!(segments.iter_core_segments().count(), 0);
1736
1737        let other_segments = fixture
1738            .grant
1739            .dst_grant(client_ip(), other_sni().customer_domain(), at(0))
1740            .expect("the other target is granted");
1741        assert_eq!(
1742            other_segments
1743                .iter_core_segments()
1744                .map(SignedPathSegment::fingerprint)
1745                .collect::<Vec<_>>(),
1746            vec![other.fingerprint()],
1747            "a segment granted for another target is only visible there"
1748        );
1749
1750        let ungranted =
1751            CustomerDomain::new("ungranted.example.com".to_owned()).expect("a valid domain");
1752        assert!(
1753            fixture
1754                .grant
1755                .dst_grant(client_ip(), ungranted.as_domain(), at(0))
1756                .is_none()
1757        );
1758    }
1759
1760    /// A fixture whose targets hold at most `max_grants` segment grants, with a 100s duration.
1761    fn capped_fixture(max_grants: usize) -> Fixture {
1762        Fixture::with_grant_config(
1763            MockFetcher::empty(),
1764            GrantManagerConfig {
1765                grant_duration: Duration::from_secs(100),
1766                max_segments_per_target: max_grants,
1767                ..Fixture::grant_config_defaults()
1768            },
1769        )
1770    }
1771
1772    /// When the test client's grant on `segment` for [`sni`] expires, if it holds one.
1773    fn expiry_of(
1774        fixture: &Fixture,
1775        segment: &SignedPathSegment,
1776        now: SystemTime,
1777    ) -> Option<SystemTime> {
1778        fixture.grant.segment_grant_expiry(
1779            client_ip(),
1780            sni().customer_domain(),
1781            &granted_id(segment),
1782            now,
1783        )
1784    }
1785
1786    #[test]
1787    fn a_target_at_its_limit_evicts_its_oldest_unused_grant() {
1788        let fixture = capped_fixture(2);
1789
1790        let oldest = nth_up_segment(0);
1791        let newer = nth_up_segment(1);
1792        let extra = nth_up_segment(2);
1793
1794        fixture.grant_non_core(vec![oldest.clone()], at(0));
1795        fixture.grant_non_core(vec![newer.clone()], at(10));
1796
1797        // Watching a grant holds it, so keep only the token: this one is meant to be evictable.
1798        let watch = fixture
1799            .grant
1800            .watch_segment_grant(
1801                client_ip(),
1802                sni().customer_domain(),
1803                &granted_id(&oldest),
1804                at(10),
1805            )
1806            .expect("the first segment is granted");
1807        let oldest_expired = watch.expired().clone();
1808        drop(watch);
1809
1810        fixture.grant_non_core(vec![extra.clone()], at(20));
1811
1812        assert_eq!(
1813            expiry_of(&fixture, &oldest, at(20)),
1814            None,
1815            "the grant with the least left to lose made room for the new one"
1816        );
1817        assert!(
1818            oldest_expired.is_cancelled(),
1819            "an evicted grant is cancelled, just like an expired one"
1820        );
1821        assert!(
1822            fixture
1823                .grant
1824                .segment(&granted_id(&oldest), at(20))
1825                .is_none(),
1826            "and the segment goes with its last grant"
1827        );
1828
1829        assert_eq!(expiry_of(&fixture, &newer, at(20)), Some(at(110)));
1830        assert_eq!(expiry_of(&fixture, &extra, at(20)), Some(at(120)));
1831    }
1832
1833    #[test]
1834    fn a_grant_a_consumer_holds_is_not_evicted() {
1835        let fixture = capped_fixture(2);
1836
1837        let held = nth_up_segment(0);
1838        let unheld = nth_up_segment(1);
1839        let extra = nth_up_segment(2);
1840
1841        // The held grant is the older of the two, so it would go first if nobody had it.
1842        fixture.grant_non_core(vec![held.clone()], at(0));
1843        fixture.grant_non_core(vec![unheld.clone()], at(10));
1844
1845        let hold = fixture
1846            .grant
1847            .watch_segment_grant(
1848                client_ip(),
1849                sni().customer_domain(),
1850                &granted_id(&held),
1851                at(10),
1852            )
1853            .expect("the first segment is granted");
1854
1855        fixture.grant_non_core(vec![extra.clone()], at(20));
1856
1857        assert_eq!(
1858            expiry_of(&fixture, &held, at(20)),
1859            Some(at(100)),
1860            "a grant a consumer still holds is spared, however little it has left"
1861        );
1862        assert!(!hold.expired().is_cancelled());
1863        assert_eq!(
1864            expiry_of(&fixture, &unheld, at(20)),
1865            None,
1866            "the next candidate goes instead"
1867        );
1868        assert_eq!(expiry_of(&fixture, &extra, at(20)), Some(at(120)));
1869    }
1870
1871    #[test]
1872    fn dropping_a_hold_frees_the_grant_for_eviction_again() {
1873        let fixture = capped_fixture(1);
1874
1875        let held = nth_up_segment(0);
1876        let extra = nth_up_segment(1);
1877
1878        fixture.grant_non_core(vec![held.clone()], at(0));
1879
1880        let hold = fixture
1881            .grant
1882            .watch_segment_grant(
1883                client_ip(),
1884                sni().customer_domain(),
1885                &granted_id(&held),
1886                at(0),
1887            )
1888            .expect("the segment is granted");
1889
1890        fixture
1891            .grant
1892            .authorize(
1893                client_ip(),
1894                BTreeMap::from([(
1895                    sni().customer_domain().into(),
1896                    AuthSegments {
1897                        up_segments: vec![extra.clone()],
1898                        ..AuthSegments::default()
1899                    },
1900                )]),
1901                at(10),
1902            )
1903            .expect_err("the only grant of the target is in use");
1904
1905        drop(hold);
1906
1907        fixture.grant_non_core(vec![extra.clone()], at(20));
1908
1909        assert_eq!(expiry_of(&fixture, &held, at(20)), None);
1910        assert_eq!(expiry_of(&fixture, &extra, at(20)), Some(at(120)));
1911    }
1912
1913    #[test]
1914    fn a_request_that_cannot_free_room_is_rejected_whole() {
1915        let fixture = capped_fixture(2);
1916
1917        let first = nth_up_segment(0);
1918        let second = nth_up_segment(1);
1919        let extra = nth_up_segment(2);
1920
1921        fixture.grant_non_core(vec![first.clone(), second.clone()], at(0));
1922
1923        let holds = [&first, &second].map(|segment| {
1924            fixture
1925                .grant
1926                .watch_segment_grant(
1927                    client_ip(),
1928                    sni().customer_domain(),
1929                    &granted_id(segment),
1930                    at(0),
1931                )
1932                .expect("both segments are granted")
1933        });
1934
1935        // The same request carries a second destination, which would have fit on its own.
1936        let error = fixture
1937            .grant
1938            .authorize(
1939                client_ip(),
1940                BTreeMap::from([
1941                    (
1942                        sni().customer_domain().into(),
1943                        AuthSegments {
1944                            up_segments: vec![extra.clone()],
1945                            ..AuthSegments::default()
1946                        },
1947                    ),
1948                    (
1949                        other_sni().customer_domain().into(),
1950                        AuthSegments::default(),
1951                    ),
1952                ]),
1953                at(10),
1954            )
1955            .expect_err("a full target whose grants are all in use takes nothing more");
1956
1957        assert!(
1958            matches!(
1959                error,
1960                AuthorizeError::TargetGrantLimitReached {
1961                    needed: 1,
1962                    in_use: 2,
1963                    max: 2,
1964                    ..
1965                }
1966            ),
1967            "{error:?}"
1968        );
1969
1970        assert!(holds.iter().all(|hold| !hold.expired().is_cancelled()));
1971        assert_eq!(expiry_of(&fixture, &first, at(10)), Some(at(100)));
1972        assert_eq!(expiry_of(&fixture, &second, at(10)), Some(at(100)));
1973        assert_eq!(
1974            expiry_of(&fixture, &extra, at(10)),
1975            None,
1976            "the segment that did not fit was not granted"
1977        );
1978        assert_eq!(
1979            fixture
1980                .grant
1981                .grant_expiry(client_ip(), sni().customer_domain(), at(10)),
1982            Some(at(100)),
1983            "and the target grant it came with was not extended either"
1984        );
1985        assert!(
1986            fixture
1987                .grant
1988                .grant_expiry(client_ip(), other_sni().customer_domain(), at(10))
1989                .is_none(),
1990            "the destination that would have fit is rejected with the request it came in"
1991        );
1992    }
1993
1994    #[test]
1995    fn a_request_granting_more_than_a_target_can_hold_is_rejected() {
1996        let fixture = capped_fixture(2);
1997
1998        let error = fixture
1999            .grant
2000            .authorize(
2001                client_ip(),
2002                BTreeMap::from([(
2003                    sni().customer_domain().into(),
2004                    AuthSegments {
2005                        up_segments: vec![nth_up_segment(0), nth_up_segment(1)],
2006                        core_segments: vec![nth_up_segment(2)],
2007                        ..AuthSegments::default()
2008                    },
2009                )]),
2010                at(0),
2011            )
2012            .expect_err("three segments never fit into two grants");
2013
2014        assert!(
2015            matches!(
2016                error,
2017                AuthorizeError::TooManySegmentsForTarget {
2018                    requested: 3,
2019                    max: 2,
2020                    ..
2021                }
2022            ),
2023            "{error:?}"
2024        );
2025        assert!(
2026            !fixture.grant.ip_is_authorized(client_ip(), at(0)),
2027            "a rejected request leaves the IP as unauthorized as it was"
2028        );
2029    }
2030
2031    #[test]
2032    fn a_request_over_the_request_limit_is_rejected() {
2033        let fixture = Fixture::with_grant_config(
2034            MockFetcher::empty(),
2035            GrantManagerConfig {
2036                grant_duration: Duration::from_secs(100),
2037                max_segments_per_request: 1,
2038                ..Fixture::grant_config_defaults()
2039            },
2040        );
2041
2042        let error = fixture
2043            .grant
2044            .authorize(
2045                client_ip(),
2046                BTreeMap::from([
2047                    (
2048                        sni().customer_domain().into(),
2049                        AuthSegments {
2050                            up_segments: vec![up_segment(0)],
2051                            ..AuthSegments::default()
2052                        },
2053                    ),
2054                    (
2055                        other_sni().customer_domain().into(),
2056                        AuthSegments {
2057                            up_segments: vec![other_up_segment(0)],
2058                            ..AuthSegments::default()
2059                        },
2060                    ),
2061                ]),
2062                at(0),
2063            )
2064            .expect_err("one segment per destination is two over the whole request");
2065
2066        assert!(
2067            matches!(
2068                error,
2069                AuthorizeError::TooManySegmentsInRequest {
2070                    requested: 2,
2071                    max: 1
2072                }
2073            ),
2074            "{error:?}"
2075        );
2076        assert!(!fixture.grant.ip_is_authorized(client_ip(), at(0)));
2077    }
2078
2079    #[test]
2080    fn the_same_segment_granted_twice_costs_one_grant() {
2081        let fixture = capped_fixture(1);
2082
2083        // Up and down segments both become the same non-core grant, see `GrantedSegmentId`.
2084        let segment = nth_up_segment(0);
2085        fixture.grant_for(
2086            client_ip(),
2087            sni().customer_domain(),
2088            AuthSegments {
2089                up_segments: vec![segment.clone()],
2090                down_segments: vec![segment.clone()],
2091                ..AuthSegments::default()
2092            },
2093            at(0),
2094        );
2095
2096        assert_eq!(expiry_of(&fixture, &segment, at(0)), Some(at(100)));
2097    }
2098
2099    #[test]
2100    fn a_core_and_a_non_core_grant_on_one_segment_are_two_grants() {
2101        let fixture = capped_fixture(1);
2102
2103        let segment = nth_up_segment(0);
2104        let error = fixture
2105            .grant
2106            .authorize(
2107                client_ip(),
2108                BTreeMap::from([(
2109                    sni().customer_domain().into(),
2110                    AuthSegments {
2111                        up_segments: vec![segment.clone()],
2112                        core_segments: vec![segment.clone()],
2113                        ..AuthSegments::default()
2114                    },
2115                )]),
2116                at(0),
2117            )
2118            .expect_err("the two grants are told apart by more than the fingerprint");
2119
2120        assert!(
2121            matches!(
2122                error,
2123                AuthorizeError::TooManySegmentsForTarget {
2124                    requested: 2,
2125                    max: 1,
2126                    ..
2127                }
2128            ),
2129            "{error:?}"
2130        );
2131    }
2132
2133    #[test]
2134    fn refreshing_a_target_at_its_limit_evicts_nothing() {
2135        let fixture = capped_fixture(2);
2136
2137        let first = nth_up_segment(0);
2138        let second = nth_up_segment(1);
2139
2140        fixture.grant_non_core(vec![first.clone(), second.clone()], at(0));
2141        fixture.grant_non_core(vec![first.clone(), second.clone()], at(10));
2142
2143        assert_eq!(expiry_of(&fixture, &first, at(10)), Some(at(110)));
2144        assert_eq!(expiry_of(&fixture, &second, at(10)), Some(at(110)));
2145    }
2146
2147    #[test]
2148    fn the_limit_is_per_target_not_per_ip() {
2149        let fixture = capped_fixture(1);
2150
2151        let for_sni = nth_up_segment(0);
2152        let for_other = nth_up_segment(1);
2153
2154        fixture.grant_for(
2155            client_ip(),
2156            sni().customer_domain(),
2157            AuthSegments {
2158                up_segments: vec![for_sni.clone()],
2159                ..AuthSegments::default()
2160            },
2161            at(0),
2162        );
2163        fixture.grant_for(
2164            client_ip(),
2165            other_sni().customer_domain(),
2166            AuthSegments {
2167                up_segments: vec![for_other.clone()],
2168                ..AuthSegments::default()
2169            },
2170            at(0),
2171        );
2172
2173        assert_eq!(expiry_of(&fixture, &for_sni, at(0)), Some(at(100)));
2174        assert_eq!(
2175            fixture.grant.segment_grant_expiry(
2176                client_ip(),
2177                other_sni().customer_domain(),
2178                &granted_id(&for_other),
2179                at(0)
2180            ),
2181            Some(at(100))
2182        );
2183    }
2184    /// A fixture whose IPs hold at most `max_targets` target grants, with a 100s duration.
2185    fn target_capped_fixture(max_targets: usize) -> Fixture {
2186        Fixture::with_grant_config(
2187            MockFetcher::empty(),
2188            GrantManagerConfig {
2189                grant_duration: Duration::from_secs(100),
2190                max_targets_per_ip: max_targets,
2191                ..Fixture::grant_config_defaults()
2192            },
2193        )
2194    }
2195
2196    /// Grants `ip` the `n`th target over the `n`th up segment.
2197    fn grant_nth_target(fixture: &Fixture, ip: IpAddr, n: usize, now: SystemTime) {
2198        fixture.grant_for(
2199            ip,
2200            nth_sni(n).customer_domain(),
2201            AuthSegments {
2202                up_segments: vec![nth_up_segment(n as u16)],
2203                ..AuthSegments::default()
2204            },
2205            now,
2206        );
2207    }
2208
2209    /// When the test client's grant on the `n`th target expires, if it holds one.
2210    fn target_expiry_of(fixture: &Fixture, n: usize, now: SystemTime) -> Option<SystemTime> {
2211        fixture
2212            .grant
2213            .grant_expiry(client_ip(), nth_sni(n).customer_domain(), now)
2214    }
2215
2216    #[test]
2217    fn an_ip_at_its_target_limit_evicts_its_oldest_unused_target() {
2218        let fixture = target_capped_fixture(2);
2219
2220        grant_nth_target(&fixture, client_ip(), 0, at(0));
2221        grant_nth_target(&fixture, client_ip(), 1, at(10));
2222
2223        // Watching a target holds it, so keep only the token: this one is meant to be evictable.
2224        let watch = fixture
2225            .grant
2226            .watch_grant(client_ip(), nth_sni(0).customer_domain(), at(10))
2227            .expect("the first target is granted");
2228        let oldest_expired = watch.expired().clone();
2229        drop(watch);
2230
2231        grant_nth_target(&fixture, client_ip(), 2, at(20));
2232
2233        assert_eq!(
2234            target_expiry_of(&fixture, 0, at(20)),
2235            None,
2236            "the target with the least left to lose made room for the new one"
2237        );
2238        assert!(
2239            oldest_expired.is_cancelled(),
2240            "an evicted target is cancelled, just like an expired one"
2241        );
2242        assert!(
2243            fixture
2244                .grant
2245                .segment(&granted_id(&nth_up_segment(0)), at(20))
2246                .is_none(),
2247            "an evicted target takes its segment grants with it"
2248        );
2249
2250        assert_eq!(target_expiry_of(&fixture, 1, at(20)), Some(at(110)));
2251        assert_eq!(target_expiry_of(&fixture, 2, at(20)), Some(at(120)));
2252    }
2253
2254    #[test]
2255    fn a_target_a_consumer_holds_is_not_evicted() {
2256        let fixture = target_capped_fixture(2);
2257
2258        // The held target is the older of the two, so it would go first if nobody had it.
2259        grant_nth_target(&fixture, client_ip(), 0, at(0));
2260        grant_nth_target(&fixture, client_ip(), 1, at(10));
2261
2262        let hold = fixture
2263            .grant
2264            .watch_grant(client_ip(), nth_sni(0).customer_domain(), at(10))
2265            .expect("the first target is granted");
2266
2267        grant_nth_target(&fixture, client_ip(), 2, at(20));
2268
2269        assert_eq!(
2270            target_expiry_of(&fixture, 0, at(20)),
2271            Some(at(100)),
2272            "a target a consumer still holds is spared, however little it has left"
2273        );
2274        assert!(!hold.expired().is_cancelled());
2275        assert_eq!(
2276            target_expiry_of(&fixture, 1, at(20)),
2277            None,
2278            "the next candidate goes instead"
2279        );
2280        assert_eq!(target_expiry_of(&fixture, 2, at(20)), Some(at(120)));
2281    }
2282
2283    #[test]
2284    fn a_target_whose_segment_grant_is_in_use_is_not_evicted() {
2285        let fixture = target_capped_fixture(2);
2286
2287        grant_nth_target(&fixture, client_ip(), 0, at(0));
2288        grant_nth_target(&fixture, client_ip(), 1, at(10));
2289
2290        // Only the segment grant is held, the target grant itself is not.
2291        let hold = fixture
2292            .grant
2293            .watch_segment_grant(
2294                client_ip(),
2295                nth_sni(0).customer_domain(),
2296                &granted_id(&nth_up_segment(0)),
2297                at(10),
2298            )
2299            .expect("the segment of the first target is granted");
2300
2301        grant_nth_target(&fixture, client_ip(), 2, at(20));
2302
2303        assert_eq!(
2304            target_expiry_of(&fixture, 0, at(20)),
2305            Some(at(100)),
2306            "evicting the target would evict a segment grant a consumer holds"
2307        );
2308        assert!(!hold.expired().is_cancelled());
2309        assert_eq!(target_expiry_of(&fixture, 1, at(20)), None);
2310        assert_eq!(target_expiry_of(&fixture, 2, at(20)), Some(at(120)));
2311    }
2312
2313    #[test]
2314    fn dropping_a_target_hold_frees_it_for_eviction_again() {
2315        let fixture = target_capped_fixture(1);
2316
2317        grant_nth_target(&fixture, client_ip(), 0, at(0));
2318
2319        let hold = fixture
2320            .grant
2321            .watch_grant(client_ip(), nth_sni(0).customer_domain(), at(0))
2322            .expect("the target is granted");
2323
2324        fixture
2325            .grant
2326            .authorize(
2327                client_ip(),
2328                BTreeMap::from([(nth_sni(1).customer_domain().into(), AuthSegments::default())]),
2329                at(10),
2330            )
2331            .expect_err("the only target of the IP is in use");
2332
2333        drop(hold);
2334
2335        grant_nth_target(&fixture, client_ip(), 1, at(20));
2336
2337        assert_eq!(target_expiry_of(&fixture, 0, at(20)), None);
2338        assert_eq!(target_expiry_of(&fixture, 1, at(20)), Some(at(120)));
2339    }
2340
2341    #[test]
2342    fn a_request_that_cannot_free_a_target_is_rejected_whole() {
2343        let fixture = target_capped_fixture(2);
2344
2345        grant_nth_target(&fixture, client_ip(), 0, at(0));
2346        grant_nth_target(&fixture, client_ip(), 1, at(0));
2347
2348        let holds = [0, 1].map(|n| {
2349            fixture
2350                .grant
2351                .watch_grant(client_ip(), nth_sni(n).customer_domain(), at(0))
2352                .expect("both targets are granted")
2353        });
2354
2355        let error = fixture
2356            .grant
2357            .authorize(
2358                client_ip(),
2359                BTreeMap::from([(
2360                    nth_sni(2).customer_domain().into(),
2361                    AuthSegments {
2362                        up_segments: vec![nth_up_segment(2)],
2363                        ..AuthSegments::default()
2364                    },
2365                )]),
2366                at(10),
2367            )
2368            .expect_err("a full IP whose targets are all in use takes nothing more");
2369
2370        assert!(
2371            matches!(
2372                error,
2373                AuthorizeError::IpTargetLimitReached {
2374                    needed: 1,
2375                    in_use: 2,
2376                    max: 2,
2377                }
2378            ),
2379            "{error:?}"
2380        );
2381
2382        assert!(holds.iter().all(|hold| !hold.expired().is_cancelled()));
2383        assert_eq!(target_expiry_of(&fixture, 0, at(10)), Some(at(100)));
2384        assert_eq!(target_expiry_of(&fixture, 1, at(10)), Some(at(100)));
2385        assert_eq!(
2386            target_expiry_of(&fixture, 2, at(10)),
2387            None,
2388            "the target that did not fit was not granted"
2389        );
2390        assert!(
2391            fixture
2392                .grant
2393                .segment(&granted_id(&nth_up_segment(2)), at(10))
2394                .is_none(),
2395            "and neither were the segments it came with"
2396        );
2397    }
2398
2399    #[test]
2400    fn a_request_naming_more_targets_than_an_ip_can_hold_is_rejected() {
2401        let fixture = target_capped_fixture(2);
2402
2403        let error = fixture
2404            .grant
2405            .authorize(
2406                client_ip(),
2407                BTreeMap::from([
2408                    (nth_sni(0).customer_domain().into(), AuthSegments::default()),
2409                    (nth_sni(1).customer_domain().into(), AuthSegments::default()),
2410                    (nth_sni(2).customer_domain().into(), AuthSegments::default()),
2411                ]),
2412                at(0),
2413            )
2414            .expect_err("three targets never fit into two grants");
2415
2416        assert!(
2417            matches!(
2418                error,
2419                AuthorizeError::TooManyTargetsForIp {
2420                    requested: 3,
2421                    max: 2
2422                }
2423            ),
2424            "{error:?}"
2425        );
2426        assert!(
2427            !fixture.grant.ip_is_authorized(client_ip(), at(0)),
2428            "a rejected request leaves the IP as unauthorized as it was"
2429        );
2430    }
2431
2432    #[test]
2433    fn refreshing_the_targets_of_a_full_ip_evicts_nothing() {
2434        let fixture = target_capped_fixture(2);
2435
2436        grant_nth_target(&fixture, client_ip(), 0, at(0));
2437        grant_nth_target(&fixture, client_ip(), 1, at(0));
2438
2439        fixture
2440            .grant
2441            .authorize(
2442                client_ip(),
2443                BTreeMap::from([
2444                    (nth_sni(0).customer_domain().into(), AuthSegments::default()),
2445                    (nth_sni(1).customer_domain().into(), AuthSegments::default()),
2446                ]),
2447                at(10),
2448            )
2449            .expect("a refresh of every target fits");
2450
2451        assert_eq!(target_expiry_of(&fixture, 0, at(10)), Some(at(110)));
2452        assert_eq!(target_expiry_of(&fixture, 1, at(10)), Some(at(110)));
2453        assert!(
2454            fixture
2455                .grant
2456                .segment(&granted_id(&nth_up_segment(0)), at(10))
2457                .is_some(),
2458            "a refreshed target keeps the segment grants the refresh omits"
2459        );
2460    }
2461
2462    #[test]
2463    fn the_target_limit_is_per_ip_not_global() {
2464        let fixture = target_capped_fixture(1);
2465
2466        grant_nth_target(&fixture, client_ip(), 0, at(0));
2467        grant_nth_target(&fixture, other_client_ip(), 1, at(0));
2468
2469        assert_eq!(target_expiry_of(&fixture, 0, at(0)), Some(at(100)));
2470        assert_eq!(
2471            fixture
2472                .grant
2473                .grant_expiry(other_client_ip(), nth_sni(1).customer_domain(), at(0)),
2474            Some(at(100)),
2475            "another IP at its own limit is untouched"
2476        );
2477    }
2478}