use std::collections::HashMap;
use std::sync::RwLock;
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum LockMode {
Governance,
Compliance,
}
impl LockMode {
#[must_use]
pub fn as_aws_str(self) -> &'static str {
match self {
Self::Governance => "GOVERNANCE",
Self::Compliance => "COMPLIANCE",
}
}
#[must_use]
pub fn from_aws_str(s: &str) -> Option<Self> {
if s.eq_ignore_ascii_case("GOVERNANCE") {
Some(Self::Governance)
} else if s.eq_ignore_ascii_case("COMPLIANCE") {
Some(Self::Compliance)
} else {
None
}
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ObjectLockState {
pub mode: Option<LockMode>,
pub retain_until: Option<DateTime<Utc>>,
pub legal_hold_on: bool,
}
impl ObjectLockState {
#[must_use]
pub fn is_locked(&self, now: DateTime<Utc>) -> bool {
if self.legal_hold_on {
return true;
}
match (self.mode, self.retain_until) {
(Some(_), Some(until)) => until > now,
_ => false,
}
}
#[must_use]
pub fn can_delete(&self, now: DateTime<Utc>, bypass_governance: bool) -> bool {
if self.legal_hold_on {
return false;
}
match (self.mode, self.retain_until) {
(Some(LockMode::Compliance), Some(until)) if until > now => false,
(Some(LockMode::Governance), Some(until)) if until > now => bypass_governance,
_ => true,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BucketObjectLockDefault {
pub mode: LockMode,
pub retention_days: u32,
}
#[derive(Debug, Default, Serialize, Deserialize)]
struct ObjectLockSnapshot {
states: Vec<((String, String), ObjectLockState)>,
bucket_defaults: HashMap<String, BucketObjectLockDefault>,
}
#[derive(Debug, Default)]
pub struct ObjectLockManager {
states: RwLock<HashMap<(String, String), ObjectLockState>>,
bucket_defaults: RwLock<HashMap<String, BucketObjectLockDefault>>,
}
impl ObjectLockManager {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn set(&self, bucket: &str, key: &str, state: ObjectLockState) {
self.states
.write()
.expect("object-lock state RwLock poisoned")
.insert((bucket.to_owned(), key.to_owned()), state);
}
#[must_use]
pub fn get(&self, bucket: &str, key: &str) -> Option<ObjectLockState> {
self.states
.read()
.expect("object-lock state RwLock poisoned")
.get(&(bucket.to_owned(), key.to_owned()))
.cloned()
}
pub fn set_legal_hold(&self, bucket: &str, key: &str, on: bool) {
let mut guard = self
.states
.write()
.expect("object-lock state RwLock poisoned");
let entry = guard
.entry((bucket.to_owned(), key.to_owned()))
.or_default();
entry.legal_hold_on = on;
}
pub fn set_bucket_default(&self, bucket: &str, default: BucketObjectLockDefault) {
self.bucket_defaults
.write()
.expect("object-lock bucket-default RwLock poisoned")
.insert(bucket.to_owned(), default);
}
#[must_use]
pub fn bucket_default(&self, bucket: &str) -> Option<BucketObjectLockDefault> {
self.bucket_defaults
.read()
.expect("object-lock bucket-default RwLock poisoned")
.get(bucket)
.copied()
}
pub fn apply_default_on_put(&self, bucket: &str, key: &str, now: DateTime<Utc>) {
let Some(default) = self.bucket_default(bucket) else {
return;
};
let mut guard = self
.states
.write()
.expect("object-lock state RwLock poisoned");
let key_pair = (bucket.to_owned(), key.to_owned());
if let Some(existing) = guard.get(&key_pair)
&& (existing.mode.is_some() || existing.retain_until.is_some())
{
return;
}
let retain_until = now + Duration::days(i64::from(default.retention_days));
let entry = guard.entry(key_pair).or_default();
entry.mode = Some(default.mode);
entry.retain_until = Some(retain_until);
}
pub fn clear(&self, bucket: &str, key: &str) {
self.states
.write()
.expect("object-lock state RwLock poisoned")
.remove(&(bucket.to_owned(), key.to_owned()));
}
pub fn to_json(&self) -> Result<String, serde_json::Error> {
let states: Vec<((String, String), ObjectLockState)> = self
.states
.read()
.expect("object-lock state RwLock poisoned")
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
let bucket_defaults = self
.bucket_defaults
.read()
.expect("object-lock bucket-default RwLock poisoned")
.clone();
let snap = ObjectLockSnapshot {
states,
bucket_defaults,
};
serde_json::to_string(&snap)
}
pub fn from_json(s: &str) -> Result<Self, serde_json::Error> {
let snap: ObjectLockSnapshot = serde_json::from_str(s)?;
let mut states = HashMap::with_capacity(snap.states.len());
for (k, v) in snap.states {
states.insert(k, v);
}
Ok(Self {
states: RwLock::new(states),
bucket_defaults: RwLock::new(snap.bucket_defaults),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn now() -> DateTime<Utc> {
Utc::now()
}
#[test]
fn is_locked_future_retain_until() {
let s = ObjectLockState {
mode: Some(LockMode::Governance),
retain_until: Some(now() + Duration::hours(1)),
legal_hold_on: false,
};
assert!(s.is_locked(now()));
}
#[test]
fn is_locked_past_retain_until_is_unlocked() {
let s = ObjectLockState {
mode: Some(LockMode::Governance),
retain_until: Some(now() - Duration::hours(1)),
legal_hold_on: false,
};
assert!(!s.is_locked(now()));
}
#[test]
fn compliance_cannot_be_bypassed() {
let s = ObjectLockState {
mode: Some(LockMode::Compliance),
retain_until: Some(now() + Duration::days(7)),
legal_hold_on: false,
};
assert!(!s.can_delete(now(), true));
assert!(!s.can_delete(now(), false));
}
#[test]
fn governance_can_be_bypassed_with_header() {
let s = ObjectLockState {
mode: Some(LockMode::Governance),
retain_until: Some(now() + Duration::days(7)),
legal_hold_on: false,
};
assert!(s.can_delete(now(), true), "bypass=true should permit delete");
assert!(
!s.can_delete(now(), false),
"bypass=false should refuse delete"
);
}
#[test]
fn legal_hold_blocks_delete_independent_of_retention() {
let s = ObjectLockState {
mode: None,
retain_until: None,
legal_hold_on: true,
};
assert!(s.is_locked(now()));
assert!(!s.can_delete(now(), true), "legal hold cannot be bypassed");
assert!(!s.can_delete(now(), false));
}
#[test]
fn legal_hold_overrides_governance_bypass() {
let s = ObjectLockState {
mode: Some(LockMode::Governance),
retain_until: Some(now() + Duration::days(7)),
legal_hold_on: true,
};
assert!(!s.can_delete(now(), true));
}
#[test]
fn no_lock_no_block() {
let s = ObjectLockState::default();
assert!(!s.is_locked(now()));
assert!(s.can_delete(now(), false));
}
#[test]
fn apply_default_materialises_state_on_first_put() {
let m = ObjectLockManager::new();
m.set_bucket_default(
"b",
BucketObjectLockDefault {
mode: LockMode::Governance,
retention_days: 30,
},
);
let now = now();
m.apply_default_on_put("b", "k", now);
let state = m.get("b", "k").expect("state must be materialised");
assert_eq!(state.mode, Some(LockMode::Governance));
let until = state.retain_until.expect("retain_until must be set");
let target = now + Duration::days(30);
let diff = (until - target).num_seconds().abs();
assert!(diff <= 1, "retain_until off by {diff}s");
}
#[test]
fn apply_default_does_not_overwrite_existing_retention() {
let m = ObjectLockManager::new();
let custom_until = now() + Duration::days(365);
m.set(
"b",
"k",
ObjectLockState {
mode: Some(LockMode::Compliance),
retain_until: Some(custom_until),
legal_hold_on: false,
},
);
m.set_bucket_default(
"b",
BucketObjectLockDefault {
mode: LockMode::Governance,
retention_days: 1,
},
);
m.apply_default_on_put("b", "k", now());
let state = m.get("b", "k").unwrap();
assert_eq!(state.mode, Some(LockMode::Compliance));
assert_eq!(state.retain_until, Some(custom_until));
}
#[test]
fn apply_default_no_op_without_bucket_default() {
let m = ObjectLockManager::new();
m.apply_default_on_put("b", "k", now());
assert!(m.get("b", "k").is_none());
}
#[test]
fn set_legal_hold_creates_state_when_missing() {
let m = ObjectLockManager::new();
m.set_legal_hold("b", "k", true);
let s = m.get("b", "k").expect("state created");
assert!(s.legal_hold_on);
assert!(s.mode.is_none());
assert!(s.retain_until.is_none());
m.set_legal_hold("b", "k", false);
let s2 = m.get("b", "k").unwrap();
assert!(!s2.legal_hold_on);
}
#[test]
fn snapshot_roundtrip() {
let m = ObjectLockManager::new();
m.set(
"b1",
"k1",
ObjectLockState {
mode: Some(LockMode::Compliance),
retain_until: Some(Utc::now() + Duration::days(10)),
legal_hold_on: true,
},
);
m.set_bucket_default(
"b1",
BucketObjectLockDefault {
mode: LockMode::Governance,
retention_days: 7,
},
);
let json = m.to_json().expect("to_json");
let m2 = ObjectLockManager::from_json(&json).expect("from_json");
let s = m2.get("b1", "k1").expect("state survives roundtrip");
assert_eq!(s.mode, Some(LockMode::Compliance));
assert!(s.legal_hold_on);
let d = m2.bucket_default("b1").expect("default survives roundtrip");
assert_eq!(d.mode, LockMode::Governance);
assert_eq!(d.retention_days, 7);
}
#[test]
fn lock_mode_aws_string_roundtrip() {
assert_eq!(
LockMode::from_aws_str(LockMode::Governance.as_aws_str()),
Some(LockMode::Governance)
);
assert_eq!(
LockMode::from_aws_str(LockMode::Compliance.as_aws_str()),
Some(LockMode::Compliance)
);
assert_eq!(LockMode::from_aws_str("governance"), Some(LockMode::Governance));
assert!(LockMode::from_aws_str("nope").is_none());
}
#[test]
fn clear_removes_state() {
let m = ObjectLockManager::new();
m.set(
"b",
"k",
ObjectLockState {
mode: Some(LockMode::Governance),
retain_until: Some(Utc::now() + Duration::days(1)),
legal_hold_on: false,
},
);
assert!(m.get("b", "k").is_some());
m.clear("b", "k");
assert!(m.get("b", "k").is_none());
}
}