1use std::path::{Path, PathBuf};
2use std::time::Duration;
3
4use serde::{Deserialize, Serialize};
5use sha2::{Digest, Sha256};
6use time::OffsetDateTime;
7use time::format_description::well_known::Rfc3339;
8use tokio::fs;
9use tokio::io::AsyncWriteExt;
10
11use crate::error::Error;
12use crate::namespace::Namespace;
13use crate::storage::s3::Keyspace;
14
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct Lock {
17 pub id: String,
18 pub path: String,
19 pub locked_at: String,
20 pub owner: Owner,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24pub struct Owner {
25 pub name: String,
26}
27
28const MAX_PATH_BYTES: usize = 4096;
39
40const DEFAULT_CAPACITY: usize = 10_000;
45
46pub struct LockStore {
47 backend: Backend,
48 max_age: Option<Duration>,
49 capacity: usize,
50 conditional_writes: bool,
56}
57
58enum Backend {
59 Local { root: PathBuf },
60 Bucket(Box<Keyspace>),
61}
62
63impl LockStore {
64 pub fn local(root: impl Into<PathBuf>) -> Self {
65 Self::over(Backend::Local { root: root.into() })
66 }
67
68 pub fn bucket(keys: Keyspace) -> Self {
69 Self::over(Backend::Bucket(Box::new(keys)))
70 }
71
72 fn over(backend: Backend) -> Self {
73 Self {
74 backend,
75 max_age: None,
76 capacity: DEFAULT_CAPACITY,
77 conditional_writes: true,
78 }
79 }
80
81 #[cfg(test)]
82 fn with_capacity(mut self, capacity: usize) -> Self {
83 self.capacity = capacity;
84 self
85 }
86
87 pub fn with_max_age(mut self, max_age: Option<Duration>) -> Self {
88 self.max_age = max_age;
89 self
90 }
91
92 pub fn with_conditional_writes(mut self, supported: bool) -> Self {
96 self.conditional_writes = supported;
97 self
98 }
99
100 pub fn max_age(&self) -> Option<Duration> {
101 self.max_age
102 }
103
104 pub fn stale_for(&self, lock: &Lock) -> Option<Duration> {
107 stale_for(lock, self.max_age)
108 }
109
110 pub fn id_of(path: &str) -> String {
111 hex::encode(Sha256::digest(path.as_bytes()))[..32].to_owned()
112 }
113
114 fn prefix(ns: &Namespace) -> String {
117 format!(".locks/{}/{}/", ns.org(), ns.repo())
118 }
119
120 fn key_of(ns: &Namespace, id: &str) -> String {
121 format!("{}{id}.json", Self::prefix(ns))
122 }
123
124 pub async fn create(&self, ns: &Namespace, path: &str, owner: &str) -> Result<Lock, Error> {
125 if path.is_empty() {
126 return Err(Error::MalformedLockPath);
127 }
128 if path.len() > MAX_PATH_BYTES {
129 return Err(Error::LockPathTooLong {
130 actual: path.len(),
131 limit: MAX_PATH_BYTES,
132 });
133 }
134
135 let lock = Lock {
136 id: Self::id_of(path),
137 path: path.to_owned(),
138 locked_at: OffsetDateTime::now_utc()
139 .format(&Rfc3339)
140 .unwrap_or_else(|_| "1970-01-01T00:00:00Z".to_owned()),
141 owner: Owner {
142 name: owner.to_owned(),
143 },
144 };
145 let encoded = serde_json::to_vec(&lock)?;
146
147 if self.get(ns, &lock.id).await?.is_none() && self.count(ns).await? >= self.capacity {
154 return Err(Error::LockLimitReached {
155 limit: self.capacity,
156 });
157 }
158
159 if self.take(ns, &lock, &encoded).await? {
160 return Ok(lock);
161 }
162
163 let Some(held) = self.get(ns, &lock.id).await? else {
168 return Err(Error::LockHeld(Box::new(lock)));
169 };
170
171 let Some(age) = self.stale_for(&held) else {
172 return Err(Error::LockHeld(Box::new(held)));
173 };
174
175 self.discard(ns, &held.id).await?;
180
181 if !self.take(ns, &lock, &encoded).await? {
182 return match self.get(ns, &lock.id).await? {
183 Some(other) => Err(Error::LockHeld(Box::new(other))),
184 None => Err(Error::LockHeld(Box::new(lock))),
185 };
186 }
187
188 crate::audit::audit_log!(
189 actor = lock.owner.name,
190 namespace = %ns,
191 path = lock.path,
192 previous_owner = held.owner.name,
193 untouched_for_seconds = age.as_secs(),
194 "a lock nobody had touched was taken over"
195 );
196
197 Ok(lock)
198 }
199
200 async fn take(&self, ns: &Namespace, lock: &Lock, encoded: &[u8]) -> Result<bool, Error> {
201 match &self.backend {
202 Backend::Local { root } => {
203 Self::write_new(&Self::path_in(root, ns, &lock.id), encoded).await
204 }
205 Backend::Bucket(_) if !self.conditional_writes => Err(Error::Unsupported(
210 "this object store does not refuse a conditional write, so a lock here could be \
211 held by two people at once",
212 )),
213 Backend::Bucket(bucket) => {
214 bucket
215 .put_if_absent(&Self::key_of(ns, &lock.id), encoded.to_vec())
216 .await
217 }
218 }
219 }
220
221 async fn discard(&self, ns: &Namespace, id: &str) -> Result<(), Error> {
224 match self.remove(ns, id).await {
225 Ok(()) | Err(Error::LockNotFound) => Ok(()),
226 Err(error) => Err(error),
227 }
228 }
229
230 async fn write_new(path: &Path, encoded: &[u8]) -> Result<bool, Error> {
231 let parent = path.parent().expect("lock paths have a parent");
232 fs::create_dir_all(parent).await?;
233
234 match fs::OpenOptions::new()
235 .write(true)
236 .create_new(true)
237 .open(path)
238 .await
239 {
240 Ok(mut file) => {
241 file.write_all(encoded).await?;
242 file.sync_all().await?;
243 Ok(true)
244 }
245 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(false),
246 Err(error) => Err(error.into()),
247 }
248 }
249
250 pub async fn get(&self, ns: &Namespace, id: &str) -> Result<Option<Lock>, Error> {
251 if !is_well_formed_id(id) {
252 return Ok(None);
253 }
254
255 let encoded = match &self.backend {
256 Backend::Local { root } => match fs::read(Self::path_in(root, ns, id)).await {
257 Ok(bytes) => Some(bytes),
258 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
259 Err(error) => return Err(error.into()),
260 },
261 Backend::Bucket(bucket) => bucket.get_bytes(&Self::key_of(ns, id)).await?,
262 };
263
264 Ok(encoded.and_then(|bytes| serde_json::from_slice(&bytes).ok()))
265 }
266
267 async fn count(&self, ns: &Namespace) -> Result<usize, Error> {
273 match &self.backend {
274 Backend::Local { root } => {
275 let Ok(mut entries) = fs::read_dir(Self::directory_in(root, ns)).await else {
276 return Ok(0);
277 };
278
279 let mut found = 0;
280 while entries.next_entry().await?.is_some() {
281 found += 1;
282 }
283 Ok(found)
284 }
285 Backend::Bucket(bucket) => Ok(bucket.keys(&Self::prefix(ns)).await?.len()),
286 }
287 }
288
289 pub async fn list(&self, ns: &Namespace) -> Result<Vec<Lock>, Error> {
290 let mut locks = match &self.backend {
291 Backend::Local { root } => Self::list_local(&Self::directory_in(root, ns)).await?,
292 Backend::Bucket(bucket) => Self::list_bucket(bucket, ns).await?,
293 };
294
295 locks.sort_by(|a: &Lock, b: &Lock| a.path.cmp(&b.path));
296 Ok(locks)
297 }
298
299 async fn list_local(directory: &Path) -> Result<Vec<Lock>, Error> {
300 let Ok(mut entries) = fs::read_dir(directory).await else {
301 return Ok(Vec::new());
302 };
303
304 let mut locks = Vec::new();
305 while let Some(entry) = entries.next_entry().await? {
306 if let Ok(bytes) = fs::read(entry.path()).await
307 && let Ok(lock) = serde_json::from_slice(&bytes)
308 {
309 locks.push(lock);
310 }
311 }
312
313 Ok(locks)
314 }
315
316 async fn list_bucket(bucket: &Keyspace, ns: &Namespace) -> Result<Vec<Lock>, Error> {
321 let mut locks = Vec::new();
322
323 for key in bucket.keys(&Self::prefix(ns)).await? {
324 if let Some(bytes) = bucket.get_bytes(&key).await?
325 && let Ok(lock) = serde_json::from_slice(&bytes)
326 {
327 locks.push(lock);
328 }
329 }
330
331 Ok(locks)
332 }
333
334 pub async fn remove(&self, ns: &Namespace, id: &str) -> Result<(), Error> {
335 if !is_well_formed_id(id) {
336 return Err(Error::LockNotFound);
337 }
338
339 let removed = match &self.backend {
340 Backend::Local { root } => match fs::remove_file(Self::path_in(root, ns, id)).await {
341 Ok(()) => true,
342 Err(error) if error.kind() == std::io::ErrorKind::NotFound => false,
343 Err(error) => return Err(error.into()),
344 },
345 Backend::Bucket(bucket) => bucket.delete(&Self::key_of(ns, id)).await?,
346 };
347
348 removed.then_some(()).ok_or(Error::LockNotFound)
349 }
350
351 fn directory_in(root: &Path, ns: &Namespace) -> PathBuf {
352 root.join(".locks").join(ns.org()).join(ns.repo())
353 }
354
355 fn path_in(root: &Path, ns: &Namespace, id: &str) -> PathBuf {
356 Self::directory_in(root, ns).join(format!("{id}.json"))
357 }
358}
359
360pub fn stale_for(lock: &Lock, max_age: Option<Duration>) -> Option<Duration> {
364 let max_age = max_age?;
365 let taken = OffsetDateTime::parse(&lock.locked_at, &Rfc3339).ok()?;
366
367 let age = Duration::try_from(OffsetDateTime::now_utc() - taken).ok()?;
369
370 (age > max_age).then_some(age)
371}
372
373fn is_well_formed_id(id: &str) -> bool {
374 id.len() == 32 && id.bytes().all(|b| b.is_ascii_hexdigit())
375}
376
377#[cfg(test)]
378mod tests;