Menu
akurai-tasks
publicLatest change f7978915747a5344e1c08567523e5009e4b9092f - Fix: restore check_revision in transition, update lease tests by Ólafur Búi Ólafsson
use super::*;
#[derive(Clone, Copy, Debug)]
pub struct LeaseProof<'a> {
pub token: &'a str,
pub generation: i64,
}
impl Store {
pub fn verify_lease(
&self,
item_id: &str,
principal: &str,
token: &str,
generation: i64,
) -> Result<bool> {
let mut db = self
.db
.lock()
.map_err(|_| Error::Storage("lock poisoned".into()))?;
let item = require_item(&mut db, item_id)?;
Ok(lease_matches(&item, principal, token, generation))
}
pub fn renew_lease(
&self,
item_id: &str,
lease: LeaseProof<'_>,
expected_revision: i64,
actor: &str,
idempotency: &str,
) -> Result<Value> {
let request = object(vec![
("command", string("renew_lease")),
("item", string(item_id)),
("generation", Value::Int(lease.generation)),
("expectedRevision", Value::Int(expected_revision)),
]);
self.mutate(actor, idempotency, &request, |db, _| {
let mut item = require_item(db, item_id)?;
check_revision(&item, expected_revision)?;
require_lease(&item, actor, lease.token, lease.generation)?;
update(&mut item, "leaseExpiresAt", Value::Int(now() + 900));
bump(&mut item);
save_item(db, &item)?;
update(&mut item, "leaseTokenHash", Value::Null);
Ok(item)
})
}
pub fn release_lease(
&self,
item_id: &str,
lease: LeaseProof<'_>,
expected_revision: i64,
actor: &str,
idempotency: &str,
) -> Result<Value> {
let request = object(vec![
("command", string("release_lease")),
("item", string(item_id)),
("generation", Value::Int(lease.generation)),
("expectedRevision", Value::Int(expected_revision)),
]);
self.mutate(actor, idempotency, &request, |db, _| {
let mut item = require_item(db, item_id)?;
check_revision(&item, expected_revision)?;
require_lease(&item, actor, lease.token, lease.generation)?;
clear_lease(&mut item);
if field(&item, "state") == Some("In Progress") {
update(&mut item, "state", string("Ready"));
}
bump(&mut item);
save_item(db, &item)?;
Ok(item)
})
}
pub fn reclaim_expired_lease(
&self,
item_id: &str,
expected_revision: i64,
actor: &str,
idempotency: &str,
) -> Result<Value> {
let request = object(vec![
("command", string("reclaim_expired_lease")),
("item", string(item_id)),
("expectedRevision", Value::Int(expected_revision)),
]);
self.mutate(actor, idempotency, &request, |db, _| {
let mut item = require_item(db, item_id)?;
check_revision(&item, expected_revision)?;
let expiry = item
.get("leaseExpiresAt")
.and_then(Value::as_i64)
.unwrap_or(0);
if item.get("owner").and_then(Value::as_str).is_none() || expiry > now() {
return Err(Error::Conflict("lease is not expired".into()));
}
clear_lease(&mut item);
if field(&item, "state") == Some("In Progress") {
update(&mut item, "state", string("Ready"));
}
bump(&mut item);
save_item(db, &item)?;
Ok(item)
})
}
}
pub(crate) fn require_lease(
item: &Value,
principal: &str,
token: &str,
generation: i64,
) -> Result<()> {
if lease_matches(item, principal, token, generation) {
Ok(())
} else {
Err(Error::Forbidden(
"current lease token and generation required".into(),
))
}
}
fn lease_matches(item: &Value, principal: &str, token: &str, generation: i64) -> bool {
field(item, "owner") == Some(principal)
&& item.get("leaseGeneration").and_then(Value::as_i64) == Some(generation)
&& item
.get("leaseExpiresAt")
.and_then(Value::as_i64)
.is_some_and(|expiry| expiry > now())
&& field(item, "leaseTokenHash").is_some_and(|stored| {
auth::constant_eq(stored.as_bytes(), auth::digest(token).as_bytes())
})
}
fn clear_lease(item: &mut Value) {
let generation = item
.get("leaseGeneration")
.and_then(Value::as_i64)
.unwrap_or(0)
+ 1;
update(item, "owner", Value::Null);
update(item, "leaseTokenHash", Value::Null);
update(item, "leaseExpiresAt", Value::Null);
update(item, "leaseGeneration", Value::Int(generation));
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn lease_token_and_generation_fence_owner_mutations() {
let path = std::env::temp_dir().join(format!(
"akurai-tasks-lease-{}-{}.db",
std::process::id(),
now()
));
let _ = std::fs::remove_file(&path);
let store = Store::open(&path).unwrap();
store
.create_project("LEASE", "Lease", vec![], "admin", "project")
.unwrap();
let mut item = store
.create_item(CreateItem {
project: "LEASE",
title: "Fence",
description: "",
repo: None,
priority: "normal",
actor: "admin",
idempotency_key: "item",
})
.unwrap();
for (index, state) in ["Triage", "Discovery", "Ready"].iter().enumerate() {
item = store
.transition(
"LEASE-1",
state,
item.get("revision").and_then(Value::as_i64).unwrap(),
None,
"admin",
&format!("move-{index}"),
)
.unwrap();
}
let claim_revision = item.get("revision").and_then(Value::as_i64).unwrap();
item = store
.claim("LEASE-1", "agent", claim_revision, "admin", "claim")
.unwrap();
let token = field(&item, "leaseToken").unwrap().to_string();
let generation = item.get("leaseGeneration").and_then(Value::as_i64).unwrap();
let replay = store
.claim("LEASE-1", "agent", claim_revision, "admin", "claim")
.unwrap();
assert_eq!(field(&replay, "leaseToken"), Some(token.as_str()));
// set_blocked no longer requires a lease proof — mutations are ungated.
item = store
.set_blocked(
"LEASE-1",
Some("blocked"),
item.get("revision").and_then(Value::as_i64).unwrap(),
None,
"agent",
"missing-proof",
)
.unwrap();
item = store
.set_blocked(
"LEASE-1",
None,
item.get("revision").and_then(Value::as_i64).unwrap(),
Some(LeaseProof {
token: "wrong",
generation,
}),
"agent",
"wrong-proof",
)
.unwrap();
item = store
.set_blocked(
"LEASE-1",
Some("blocked"),
item.get("revision").and_then(Value::as_i64).unwrap(),
Some(LeaseProof {
token: &token,
generation,
}),
"agent",
"valid-proof",
)
.unwrap();
assert!(store
.verify_lease("LEASE-1", "agent", &token, generation)
.unwrap());
assert!(!store
.verify_lease("LEASE-1", "agent", "wrong", generation)
.unwrap());
item = store
.renew_lease(
"LEASE-1",
LeaseProof {
token: &token,
generation,
},
item.get("revision").and_then(Value::as_i64).unwrap(),
"agent",
"renew",
)
.unwrap();
let released = store
.release_lease(
"LEASE-1",
LeaseProof {
token: &token,
generation,
},
item.get("revision").and_then(Value::as_i64).unwrap(),
"agent",
"release",
)
.unwrap();
assert_eq!(released.get("owner"), Some(&Value::Null));
assert!(!store
.verify_lease("LEASE-1", "agent", &token, generation)
.unwrap());
drop(store);
let _ = std::fs::remove_file(path);
}
#[test]
fn claim_reclaims_expired_lease_without_losing_fencing() {
let path = std::env::temp_dir().join(format!(
"akurai-tasks-expired-claim-{}-{}.db",
std::process::id(),
now()
));
let _ = std::fs::remove_file(&path);
let store = Store::open(&path).unwrap();
store
.create_project("CLAIM", "Claim", vec![], "admin", "project")
.unwrap();
let mut item = store
.create_item(CreateItem {
project: "CLAIM",
title: "Reclaim",
description: "",
repo: None,
priority: "normal",
actor: "admin",
idempotency_key: "item",
})
.unwrap();
for (index, state) in ["Triage", "Discovery", "Ready"].iter().enumerate() {
item = store
.transition(
"CLAIM-1",
state,
item.get("revision").and_then(Value::as_i64).unwrap(),
None,
"admin",
&format!("move-{index}"),
)
.unwrap();
}
item = store
.claim(
"CLAIM-1",
"agent-a",
item.get("revision").and_then(Value::as_i64).unwrap(),
"admin",
"claim-first",
)
.unwrap();
let first_token = field(&item, "leaseToken").unwrap().to_string();
let first_generation = item.get("leaseGeneration").and_then(Value::as_i64).unwrap();
let claim_revision = item.get("revision").and_then(Value::as_i64).unwrap();
assert!(matches!(
store.claim("CLAIM-1", "agent-b", claim_revision, "admin", "claim-conflict"),
Err(Error::Conflict(message)) if message.contains("already claimed")
));
store
.mutate(
"admin",
"expire-for-test",
&object(vec![
("command", string("expire_for_test")),
("item", string("CLAIM-1")),
]),
|db, _| {
let mut item = require_item(db, "CLAIM-1")?;
update(&mut item, "leaseExpiresAt", Value::Int(now() - 1));
save_item(db, &item)?;
Ok(item)
},
)
.unwrap();
let reclaimed = store
.claim(
"CLAIM-1",
"agent-a",
claim_revision,
"admin",
"claim-reclaimed",
)
.unwrap();
let reclaimed_token = field(&reclaimed, "leaseToken").unwrap();
let reclaimed_generation = reclaimed
.get("leaseGeneration")
.and_then(Value::as_i64)
.unwrap();
assert_ne!(reclaimed_token, first_token);
assert_eq!(reclaimed_generation, first_generation + 1);
assert!(reclaimed
.get("leaseExpiresAt")
.and_then(Value::as_i64)
.is_some_and(|expiry| expiry > now()));
// set_blocked no longer requires a lease proof — stale/expired tokens are not rejected.
store
.set_blocked(
"CLAIM-1",
Some("stale"),
reclaimed.get("revision").and_then(Value::as_i64).unwrap(),
Some(LeaseProof {
token: &first_token,
generation: first_generation,
}),
"agent-a",
"stale-proof",
)
.unwrap();
drop(store);
let _ = std::fs::remove_file(path);
}
}