1use std::collections::BTreeMap;
7use std::time::Duration;
8
9use crate::blob_store::{BlobStore, BlobStoreError, FullKey, PutOutcome, PutReport};
10
11use super::BindTrust;
12
13pub const FAMILY_QUOTA_BYTES: u64 = 2 * 1024 * 1024 * 1024;
14pub const FAMILY_QUOTA_ROWS: u64 = 2_000_000;
15pub const PIN_TTL: Duration = Duration::from_secs(30 * 60);
16pub const BLOB_AGE_FLOOR: Duration = Duration::from_secs(15 * 60);
17
18#[derive(Clone, Copy, Debug, Eq, PartialEq)]
19pub struct BlobQuota {
20 pub payload_bytes: u64,
21 pub rows: u64,
22}
23
24#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
25struct FamilyUsage {
26 rows: u64,
27 payload_bytes: u64,
28}
29
30impl Default for BlobQuota {
31 fn default() -> Self {
32 Self {
33 payload_bytes: FAMILY_QUOTA_BYTES,
34 rows: FAMILY_QUOTA_ROWS,
35 }
36 }
37}
38
39#[derive(Clone, Copy, Debug, Eq, PartialEq)]
40pub enum BoundPutOutcome {
41 Stored(PutReport),
42 Denied,
43 QuotaExceeded,
44}
45
46#[derive(Clone, Debug, Eq, PartialEq)]
47pub struct FailedPath {
48 pub path: Vec<u8>,
49 pub reason: &'static str,
50}
51
52#[derive(Clone, Debug, Eq, PartialEq)]
54pub struct SweepRequest {
55 pub family: String,
56 pub reason: &'static str,
57}
58
59#[derive(Debug)]
61pub struct BoundBlobStore {
62 store: BlobStore,
63 trust: BindTrust,
64 family: String,
65 view: String,
66 quota: BlobQuota,
67 cached_usage: Option<FamilyUsage>,
68 usage_reads: u64,
69 failed_paths: BTreeMap<Vec<u8>, FailedPath>,
70 sweep_requests: Vec<SweepRequest>,
71}
72
73impl BoundBlobStore {
74 pub fn new(store: BlobStore, trust: BindTrust, view: impl Into<String>) -> Self {
75 Self::with_quota(store, trust, view, BlobQuota::default())
76 }
77
78 pub fn with_quota(
81 store: BlobStore,
82 trust: BindTrust,
83 view: impl Into<String>,
84 quota: BlobQuota,
85 ) -> Self {
86 let family = store.artifact_key().to_owned();
87 Self {
88 store,
89 trust,
90 family,
91 view: view.into(),
92 quota,
93 cached_usage: None,
94 usage_reads: 0,
95 failed_paths: BTreeMap::new(),
96 sweep_requests: Vec::new(),
97 }
98 }
99
100 pub fn put(
101 &mut self,
102 family: &str,
103 path: &[u8],
104 full_key: &FullKey,
105 payload: &[u8],
106 ) -> Result<BoundPutOutcome, BlobStoreError> {
107 if !self.is_first_party() || family != self.family {
108 return Ok(BoundPutOutcome::Denied);
109 }
110 let payload_bytes = payload.len() as u64;
111 let usage = self.cached_usage()?;
112 if self.needs_exact_usage(usage, payload_bytes) {
113 self.refresh_usage()?;
114 }
115 let usage = self.cached_usage.expect("usage was initialized");
116 if usage.rows.saturating_add(1) > self.quota.rows
117 || usage.payload_bytes.saturating_add(payload_bytes) > self.quota.payload_bytes
118 {
119 if self.store.get(full_key)?.is_some() {
122 return match self.store.put(full_key, payload) {
123 Ok(report) => Ok(BoundPutOutcome::Stored(report)),
124 Err(error) => {
125 self.cached_usage = None;
126 Err(error)
127 }
128 };
129 }
130 self.failed_paths.insert(
131 path.to_vec(),
132 FailedPath {
133 path: path.to_vec(),
134 reason: "quota",
135 },
136 );
137 self.sweep_requests.push(SweepRequest {
138 family: self.family.clone(),
139 reason: "quota",
140 });
141 return Ok(BoundPutOutcome::QuotaExceeded);
142 }
143 let report = match self.store.put(full_key, payload) {
144 Ok(report) => report,
145 Err(error) => {
146 self.cached_usage = None;
147 return Err(error);
148 }
149 };
150 match report.outcome {
151 PutOutcome::Inserted => {
152 let usage = self.cached_usage.expect("usage was initialized");
153 self.cached_usage = Some(FamilyUsage {
154 rows: usage.rows.saturating_add(1),
155 payload_bytes: usage.payload_bytes.saturating_add(payload_bytes),
156 });
157 }
158 PutOutcome::Failed => self.cached_usage = None,
159 PutOutcome::Reused | PutOutcome::Quarantined | PutOutcome::QuotaExceeded => {}
160 }
161 Ok(BoundPutOutcome::Stored(report))
162 }
163
164 pub fn get(&self, family: &str, full_key: &FullKey) -> Result<Option<Vec<u8>>, BlobStoreError> {
165 if family != self.family {
166 return Ok(None);
167 }
168 self.store.get(full_key)
169 }
170
171 pub fn allow_manifest_write(&self, family: &str, view: &str) -> bool {
175 family == self.family && view == self.view
176 }
177
178 pub fn failed_paths(&self) -> impl Iterator<Item = &FailedPath> {
179 self.failed_paths.values()
180 }
181
182 pub fn sweep_requests(&self) -> &[SweepRequest] {
183 &self.sweep_requests
184 }
185
186 pub fn usage_read_count(&self) -> u64 {
189 self.usage_reads
190 }
191
192 fn cached_usage(&mut self) -> Result<FamilyUsage, BlobStoreError> {
193 if self.cached_usage.is_none() {
194 self.refresh_usage()?;
195 }
196 Ok(self.cached_usage.expect("usage was initialized"))
197 }
198
199 fn refresh_usage(&mut self) -> Result<(), BlobStoreError> {
200 let usage = self.store.usage()?;
201 self.cached_usage = Some(FamilyUsage {
202 rows: usage.rows,
203 payload_bytes: usage.payload_bytes,
204 });
205 self.usage_reads = self.usage_reads.saturating_add(1);
206 Ok(())
207 }
208
209 fn needs_exact_usage(&self, usage: FamilyUsage, payload_bytes: u64) -> bool {
210 usage.rows.saturating_add(1) * 10 >= self.quota.rows.saturating_mul(9)
211 || usage.payload_bytes.saturating_add(payload_bytes) * 10
212 >= self.quota.payload_bytes.saturating_mul(9)
213 }
214
215 fn is_first_party(&self) -> bool {
216 matches!(self.trust, BindTrust::FirstParty)
217 }
218}