1use 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#[serde_as]
63#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
64#[serde(deny_unknown_fields)]
65pub struct GrantManagerConfig {
66 #[serde_as(as = "DurationSeconds<u64>")]
68 #[serde(default = "GrantManagerConfig::default_grant_duration")]
69 pub grant_duration: Duration,
70 #[serde_as(as = "DurationSeconds<u64>")]
72 #[serde(default = "GrantManagerConfig::default_min_clean_interval")]
73 pub min_clean_interval: Duration,
74 #[serde_as(as = "DurationSeconds<u64>")]
76 #[serde(default = "GrantManagerConfig::default_max_clean_interval")]
77 pub max_clean_interval: Duration,
78 #[serde(default = "GrantManagerConfig::default_max_segments_per_request")]
80 pub max_segments_per_request: usize,
81 #[serde(default = "GrantManagerConfig::default_max_segments_per_target")]
85 pub max_segments_per_target: usize,
86 #[serde(default = "GrantManagerConfig::default_max_targets_per_request")]
88 pub max_targets_per_request: usize,
89 #[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 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#[derive(Clone)]
183pub struct GrantManager(Arc<GrantManagerShared>);
184
185struct GrantManagerShared {
186 inner: RwLock<GrantManagerInner>,
188 config: GrantManagerConfig,
189}
190
191pub struct GrantManagerInner {
193 auth_ips: HashMap<IpAddr, IpAuthInfo>,
195
196 private_non_core_segments: HashMap<SegmentFp, AuthSegmentEntry>,
201 private_core_segments: HashMap<SegmentFp, AuthSegmentEntry>,
203}
204
205impl GrantManager {
206 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 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 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 let dst_granted_until = now + self.0.config.grant_duration;
297
298 let mut this = self.write();
299
300 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 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 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 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 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 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 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 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 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 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 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 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 if !dst_expired && segment_grant.is_valid(now) {
594 next_expiry = next_expiry.min(segment_grant.granted_until);
595 return true;
596 }
597
598 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 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 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
668fn 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 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 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
732fn 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 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 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 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 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 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 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
875fn 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
884fn 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 Entry::Occupied(mut existing) => {
900 let grant = existing.get_mut();
901 grant.granted_until = grant.granted_until.max(granted_until);
902 false
903 }
904 Entry::Vacant(vacant) => {
906 vacant.insert(SegmentGrant::new(granted_until));
907 true
908 }
909 };
910
911 match store.entry(id.fp()) {
912 Entry::Occupied(mut existing) => {
915 let existing = existing.get_mut();
916
917 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 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
939fn 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
960fn 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
984fn drop_segment_grant(
989 store: &mut HashMap<SegmentFp, AuthSegmentEntry>,
990 fp: SegmentFp,
991 grant: SegmentGrant,
992) {
993 release_segment_grant(store, fp);
994
995 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
1007fn 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#[derive(Debug, thiserror::Error)]
1026pub enum AuthorizeError {
1027 #[error("request grants {requested} segments, at most {max} are allowed per request")]
1029 TooManySegmentsInRequest {
1030 requested: usize,
1032 max: usize,
1034 },
1035 #[error("request has {requested} targets, at most {max} are allowed per request")]
1037 TooManyTargetsInRequest {
1038 requested: usize,
1040 max: usize,
1042 },
1043 #[error("request grants {requested} segments for {dst}, at most {max} are allowed per target")]
1045 TooManySegmentsForTarget {
1046 dst: CustomerDomain,
1048 requested: usize,
1050 max: usize,
1052 },
1053 #[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 dst: CustomerDomain,
1061 in_use: usize,
1063 needed: usize,
1065 max: usize,
1067 },
1068 #[error("request has {requested} targets, an IP may hold at most {max}")]
1070 TooManyTargetsForIp {
1071 requested: usize,
1073 max: usize,
1075 },
1076 #[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 in_use: usize,
1084 needed: usize,
1086 max: usize,
1088 },
1089}
1090
1091struct DstPlan {
1093 dst: CustomerDomain,
1094 grants: Vec<PlannedGrant>,
1096 evictions: Vec<GrantedSegmentId>,
1098}
1099
1100struct PlannedGrant {
1102 id: GrantedSegmentId,
1103 segment: SignedPathSegment,
1104 segment_expiry: SystemTime,
1105 granted_until: SystemTime,
1106}
1107
1108#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
1110pub enum GrantedSegmentId {
1111 Core(SegmentFp),
1113 NonCore(SegmentFp),
1115}
1116
1117impl GrantedSegmentId {
1118 pub fn fp(&self) -> SegmentFp {
1120 match self {
1121 Self::Core(fp) | Self::NonCore(fp) => *fp,
1122 }
1123 }
1124
1125 pub fn is_core(&self) -> bool {
1127 matches!(self, Self::Core(_))
1128 }
1129}
1130
1131struct 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#[derive(Default)]
1150struct IpAuthInfo {
1151 auth_dsts: HashMap<CustomerDomain, DstAuthInfo>,
1152}
1153
1154struct DstAuthInfo {
1156 segment_grants: HashMap<GrantedSegmentId, SegmentGrant>,
1157 granted_until: SystemTime,
1158 expired: CancellationToken,
1160 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 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 fn guard(&self) -> TargetGrantGuard {
1188 TargetGrantGuard {
1189 expired: self.expired.clone(),
1190 _in_use: Arc::clone(&self.in_use),
1191 }
1192 }
1193}
1194
1195struct SegmentGrant {
1197 granted_until: SystemTime,
1198 expired: CancellationToken,
1200 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 fn is_in_use(&self) -> bool {
1221 Arc::strong_count(&self.in_use) > 1
1222 }
1223
1224 fn guard(&self) -> SegmentGrantGuard {
1226 SegmentGrantGuard {
1227 expired: self.expired.clone(),
1228 _in_use: Arc::clone(&self.in_use),
1229 }
1230 }
1231}
1232
1233pub struct TargetGrantGuard {
1238 expired: CancellationToken,
1239 _in_use: Arc<()>,
1241}
1242
1243impl TargetGrantGuard {
1244 pub fn expired(&self) -> &CancellationToken {
1248 &self.expired
1249 }
1250}
1251
1252pub struct SegmentGrantGuard {
1257 expired: CancellationToken,
1258 _in_use: Arc<()>,
1260}
1261
1262impl SegmentGrantGuard {
1263 pub fn expired(&self) -> &CancellationToken {
1267 &self.expired
1268 }
1269}
1270
1271#[derive(Default)]
1275pub struct DstGrant {
1276 core_segments: HashMap<SegmentFp, SignedPathSegment>,
1277 non_core_segments: HashMap<SegmentFp, SignedPathSegment>,
1278}
1279
1280impl DstGrant {
1281 pub fn iter_core_segments(&self) -> impl Iterator<Item = &SignedPathSegment> {
1283 self.core_segments.values()
1284 }
1285
1286 pub fn iter_non_core_segments(&self) -> impl Iterator<Item = &SignedPathSegment> {
1288 self.non_core_segments.values()
1289 }
1290
1291 pub fn has_core_segment(&self, fp: SegmentFp) -> bool {
1293 self.core_segments.contains_key(&fp)
1294 }
1295
1296 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}