Skip to content

fix(pty): allow ticket-authorized ownership takeover #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jlongster wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix(pty): allow ticket-authorized ownership takeover #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,15 @@ Losing the owner connection stops the daemon and its terminals unless the owner
first sends `{"token":"...","request":{"op":"prepare_handoff"}}` on that same
socket. The response is `{"type":"handoff","ticket":"...","expires_at":123}`,
where `expires_at` is Unix milliseconds, 120 seconds from preparation. Repeated
preparation during that window returns the same ticket and deadline. After the
old owner disconnects, a successor claims the same instance with the ticket in
its `own` request. A live owner cannot be displaced, and successful adoption
consumes the ticket. Expiry stops an unowned daemon; if the old owner is still
connected, expiry simply cancels the handoff. `shutdown` always stops the daemon,
including during handoff. No ownership or handoff state is persisted.
preparation during that window returns the same ticket and deadline. A successor
claims the same instance with the ticket in its `own` request, even while the old
owner is still connected. Successful adoption consumes the ticket atomically;
only one successor can claim it. The superseded connection can no longer prepare
handoffs or shut down the daemon, and its disconnect does not affect the new owner.
Expiry stops an unowned daemon; if the old owner is still connected, expiry simply
cancels the handoff. An ordinary authenticated `shutdown` request, or one from the
current owner, stops the daemon even during handoff. No ownership or handoff state
is persisted.

## Architecture

Expand Down
46 changes: 35 additions & 11 deletions src/daemon.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,19 +233,23 @@ mod unix {
.map_err(|_| anyhow!("ownership lock poisoned"))?
.claim(ticket.as_deref(), Instant::now())
};
if let Err(error) = claim {
return write_frame(
&mut stream,
&Response::Error {
message: error.to_string(),
},
);
}
let result = owner_connection(&mut stream, registration, shutdown, ownership);
let generation = match claim {
Ok(generation) => generation,
Err(error) => {
return write_frame(
&mut stream,
&Response::Error {
message: error.to_string(),
},
);
}
};
let result =
owner_connection(&mut stream, registration, shutdown, ownership, generation);
ownership
.lock()
.map_err(|_| anyhow!("ownership lock poisoned"))?
.disconnect(Instant::now());
.disconnect(generation, Instant::now());
return result;
}
if let Request::Subscribe {
Expand Down Expand Up @@ -287,10 +291,23 @@ mod unix {
registration: &Registration,
shutdown: &AtomicBool,
ownership: &Mutex<Ownership>,
generation: u64,
) -> Result<()> {
write_frame(&mut *stream, &Response::Owned)?;
while !shutdown.load(Ordering::Acquire) {
let envelope: Envelope = read_frame(&mut *stream)?;
if !ownership
.lock()
.map_err(|_| anyhow!("ownership lock poisoned"))?
.is_owner(generation)
{
return write_frame(
&mut *stream,
&Response::Error {
message: "owner connection has been superseded".to_string(),
},
);
}
let stopping = envelope.token == registration.token
&& matches!(envelope.request, Request::Shutdown);
let response = if envelope.token != registration.token {
Expand All @@ -304,6 +321,7 @@ mod unix {
.lock()
.map_err(|_| anyhow!("ownership lock poisoned"))?
.prepare(
generation,
Instant::now(),
SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64,
)?;
Expand All @@ -321,7 +339,13 @@ mod unix {
};
let result = write_frame(&mut *stream, &response);
if stopping {
shutdown.store(true, Ordering::Release);
// A takeover may have occurred while writing the response.
let owner = ownership
.lock()
.map_err(|_| anyhow!("ownership lock poisoned"))?;
if owner.is_owner(generation) {
shutdown.store(true, Ordering::Release);
}
return result;
}
result?;
Expand Down
130 changes: 94 additions & 36 deletions src/ownership.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,12 @@ pub(crate) struct Handoff {
deadline: Instant,
}

pub(crate) enum Ownership {
pub(crate) struct Ownership {
state: State,
generation: u64,
}

enum State {
Starting(Instant),
Owned(Option<Handoff>),
Waiting(Handoff),
Expand All @@ -21,35 +26,49 @@ pub(crate) enum Ownership {

impl Ownership {
pub fn new(now: Instant) -> Self {
Self::Starting(now + ACQUIRE_TIMEOUT)
Self {
state: State::Starting(now + ACQUIRE_TIMEOUT),
generation: 0,
}
}

pub fn tick(&mut self, now: Instant) -> bool {
match self {
Self::Starting(deadline) if now >= *deadline => *self = Self::Stopped,
Self::Waiting(handoff) if now >= handoff.deadline => *self = Self::Stopped,
Self::Owned(Some(handoff)) if now >= handoff.deadline => *self = Self::Owned(None),
match &self.state {
State::Starting(deadline) if now >= *deadline => self.state = State::Stopped,
State::Waiting(handoff) if now >= handoff.deadline => self.state = State::Stopped,
State::Owned(Some(handoff)) if now >= handoff.deadline => {
self.state = State::Owned(None)
}
_ => {}
}
matches!(self, Self::Stopped)
matches!(self.state, State::Stopped)
}

pub fn claim(&mut self, ticket: Option<&str>, now: Instant) -> Result<()> {
pub fn claim(&mut self, ticket: Option<&str>, now: Instant) -> Result<u64> {
self.tick(now);
match &*self {
Self::Starting(_) if ticket.is_none() => {}
Self::Waiting(handoff) if ticket == Some(handoff.ticket.as_str()) => {}
Self::Owned(_) => bail!("daemon already has a live owner"),
Self::Stopped => bail!("daemon ownership deadline expired or daemon is stopping"),
match &self.state {
State::Starting(_) if ticket.is_none() => {}
State::Owned(Some(handoff)) | State::Waiting(handoff)
if ticket == Some(handoff.ticket.as_str()) => {}
State::Owned(_) => bail!("daemon already has a live owner"),
State::Stopped => bail!("daemon ownership deadline expired or daemon is stopping"),
_ => bail!("invalid handoff ticket"),
}
*self = Self::Owned(None);
Ok(())
self.generation += 1;
self.state = State::Owned(None);
Ok(self.generation)
}

pub fn prepare(&mut self, now: Instant, unix_ms: u64) -> Result<Handoff> {
pub fn is_owner(&self, generation: u64) -> bool {
self.generation == generation && matches!(self.state, State::Owned(_))
}

pub fn prepare(&mut self, generation: u64, now: Instant, unix_ms: u64) -> Result<Handoff> {
self.tick(now);
let Self::Owned(handoff) = self else {
if self.generation != generation {
bail!("handoff requires the owner connection");
}
let State::Owned(handoff) = &mut self.state else {
bail!("handoff requires the owner connection");
};
Ok(handoff
Expand All @@ -61,11 +80,15 @@ impl Ownership {
.clone())
}

pub fn disconnect(&mut self, now: Instant) {
pub fn disconnect(&mut self, generation: u64, now: Instant) {
// A ticket may transfer ownership before the previous connection closes.
if !self.is_owner(generation) {
return;
}
self.tick(now);
*self = match self {
Self::Owned(Some(handoff)) => Self::Waiting(handoff.clone()),
_ => Self::Stopped,
self.state = match &self.state {
State::Owned(Some(handoff)) => State::Waiting(handoff.clone()),
_ => State::Stopped,
};
}
}
Expand All @@ -79,9 +102,9 @@ mod tests {
let now = Instant::now();
let mut owner = Ownership::new(now);
assert!(owner.claim(Some("unexpected"), now).is_err());
owner.claim(None, now).unwrap();
let generation = owner.claim(None, now).unwrap();
assert!(owner.claim(None, now).is_err());
owner.disconnect(now);
owner.disconnect(generation, now);
assert!(owner.tick(now));
assert!(owner.claim(None, now).is_err());

Expand All @@ -95,40 +118,75 @@ mod tests {
fn handoff_is_nonrenewing_and_consumed_by_claim() {
let now = Instant::now();
let mut owner = Ownership::new(now);
owner.claim(None, now).unwrap();
let handoff = owner.prepare(now, 1000).unwrap();
let generation = owner.claim(None, now).unwrap();
let handoff = owner.prepare(generation, now, 1000).unwrap();
assert_eq!(handoff.expires_at, 121_000);
assert_eq!(
owner
.prepare(now + Duration::from_secs(60), 61_000)
.prepare(generation, now + Duration::from_secs(60), 61_000)
.unwrap(),
handoff
);
owner.disconnect(generation, now);
assert!(owner.claim(None, now).is_err());
assert!(owner.claim(Some("wrong"), now).is_err());
let successor = owner.claim(Some(&handoff.ticket), now).unwrap();
assert!(owner.claim(Some(&handoff.ticket), now).is_err());
owner.disconnect(now);
owner.disconnect(successor, now);
assert!(owner.tick(now));
}

#[test]
fn valid_ticket_replaces_live_owner_and_fences_old_connection() {
let now = Instant::now();
let mut owner = Ownership::new(now);
let generation = owner.claim(None, now).unwrap();
let handoff = owner.prepare(generation, now, 1000).unwrap();
assert!(owner.claim(None, now).is_err());
assert!(owner.claim(Some("wrong"), now).is_err());
owner.claim(Some(&handoff.ticket), now).unwrap();
assert!(owner.is_owner(generation));

let successor = owner.claim(Some(&handoff.ticket), now).unwrap();
assert!(!owner.is_owner(generation));
assert!(owner.is_owner(successor));
assert!(owner.claim(Some(&handoff.ticket), now).is_err());
owner.disconnect(now);
assert!(owner.prepare(generation, now, 1000).is_err());
owner.disconnect(generation, now);
assert!(owner.is_owner(successor));
assert!(!owner.tick(now));

let next = owner.prepare(successor, now, 1000).unwrap();
owner.disconnect(generation, now);
assert!(owner.is_owner(successor));
owner.disconnect(successor, now);
// Stale cleanup must not discard the successor's pending handoff either.
owner.disconnect(generation, now);
let last = owner.claim(Some(&next.ticket), now).unwrap();
owner.disconnect(last, now);
assert!(owner.tick(now));
}

#[test]
fn expiry_only_stops_a_disconnected_owner() {
let now = Instant::now();
let mut owner = Ownership::new(now);
owner.claim(None, now).unwrap();
let handoff = owner.prepare(now, 1000).unwrap();
let generation = owner.claim(None, now).unwrap();
let handoff = owner.prepare(generation, now, 1000).unwrap();
assert!(!owner.tick(handoff.deadline));
assert!(matches!(owner, Ownership::Owned(None)));
owner.disconnect(handoff.deadline);
assert!(matches!(owner.state, State::Owned(None)));
assert!(
owner
.claim(Some(&handoff.ticket), handoff.deadline)
.is_err()
);
assert!(owner.is_owner(generation));
owner.disconnect(generation, handoff.deadline);
assert!(owner.tick(handoff.deadline));

let mut owner = Ownership::new(now);
owner.claim(None, now).unwrap();
let handoff = owner.prepare(now, 1000).unwrap();
owner.disconnect(now);
let generation = owner.claim(None, now).unwrap();
let handoff = owner.prepare(generation, now, 1000).unwrap();
owner.disconnect(generation, now);
assert!(!owner.tick(handoff.deadline - Duration::from_millis(1)));
assert!(
owner
Expand Down
Loading
Loading