Skip to content

Commit 33a76e8

Browse files
committed
Fixes from LLM PR review
1 parent 8ccb8e1 commit 33a76e8

15 files changed

Lines changed: 533 additions & 230 deletions

File tree

Cargo.lock

Lines changed: 21 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ hyper = "1"
2222
indicatif = "0.18"
2323
kms-agent = { git = "https://github.com/oxidecomputer/kms-agent", branch = "lib", default-features = false, features = [] }
2424
libc = "0.2"
25+
lru = "0.18"
2526
memmap2 = "0.9"
2627
p256 = { version = "0.13", features = ["ecdsa"] }
2728
pem-rfc7468 = { version = "0.7", features = ["std"] }

api/src/lib.rs

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -72,13 +72,21 @@ pub trait SushApi {
7272
// Session management.
7373

7474
/// Start a new support session.
75+
///
76+
/// There may only be one session active on the rack at a time.
77+
/// If a session is already running when this request is made,
78+
/// the old session is stopped (regardless of ownership).
7579
#[endpoint { method = POST, path = "/sessions" }]
7680
async fn session_start(
7781
ctx: RequestContext<Self::Context>,
7882
headers: Header<Authorization>,
7983
) -> Result<HttpResponseOk<SessionId>, HttpError>;
8084

8185
/// End a support session.
86+
///
87+
/// Since there may be only one session active on the rack at a time
88+
/// and we do not want to prevent starting new sessions, anyone is
89+
/// allowed to stop anyone else's sesssion and start their own.
8290
#[endpoint { method = POST, path = "/sessions/{session_id}/stop" }]
8391
async fn session_stop(
8492
ctx: RequestContext<Self::Context>,
@@ -98,8 +106,8 @@ pub trait SushApi {
98106
body: TypedBody<SignedJob>,
99107
) -> Result<HttpResponseOk<JobStatus>, HttpError>;
100108

101-
/// Abort a started job.
102-
#[endpoint { method = GET, path = "/jobs/{job_id}/abort" }]
109+
/// Stop a (running) job.
110+
#[endpoint { method = POST, path = "/jobs/{job_id}/stop" }]
103111
async fn job_stop(
104112
ctx: RequestContext<Self::Context>,
105113
headers: Header<Authorization>,
@@ -118,7 +126,7 @@ pub trait SushApi {
118126
#[endpoint { method = GET, path = "/jobs/{job_id}/output/{stream}" }]
119127
async fn job_output(
120128
ctx: RequestContext<Self::Context>,
121-
headers: Header<RangeRequest>,
129+
headers: Header<AuthorizedRangeRequest>,
122130
params: PathParams<JobOutputParams>,
123131
) -> Result<Response<Body>, HttpError>;
124132

@@ -127,7 +135,7 @@ pub trait SushApi {
127135
#[endpoint { method = DELETE, path = "/jobs/{job_id}/output/{stream}" }]
128136
async fn job_output_delete(
129137
ctx: RequestContext<Self::Context>,
130-
headers: Header<RangeRequest>,
138+
headers: Header<AuthorizedRangeRequest>,
131139
params: PathParams<JobOutputParams>,
132140
) -> Result<HttpResponseOk<u64>, HttpError>;
133141

@@ -234,7 +242,7 @@ pub struct JobOutputParams {
234242

235243
/// Authorized `Range` request headers.
236244
#[derive(Debug, Deserialize, JsonSchema)]
237-
pub struct RangeRequest {
245+
pub struct AuthorizedRangeRequest {
238246
/// Authorization to access the range.
239247
///
240248
/// See: <https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Authorization>
@@ -246,7 +254,7 @@ pub struct RangeRequest {
246254
range: Option<String>,
247255
}
248256

249-
impl RangeRequest {
257+
impl AuthorizedRangeRequest {
250258
/// Extract a single range from the `Range` request header.
251259
/// This is just to avoid the complexity of encoding multiple
252260
/// ranges as output, not an inherent limitation.

client/src/cli.rs

Lines changed: 41 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use x509_cert::Certificate;
1313
use x509_cert::der::Encode as _;
1414

1515
use sush_common::authn::{Credentials, Identity};
16-
use sush_common::jobs::{JobId, JobOutputStream, JobStatus, SessionId, SignedJob};
16+
use sush_common::jobs::{JobId, JobOutputStream, JobStatus, Session, SessionId, SignedJob};
1717
use sush_common::keys::{KeyId, Signature, SshPublicKey};
1818

1919
use crate::commands::{CommandError, GlobalArgs};
@@ -23,8 +23,7 @@ use crate::context::{CommandContext, OutputFormat};
2323
pub struct Cli {
2424
output: OutputFormat,
2525
progress: Option<ProgressBar>,
26-
session: Option<SessionId>,
27-
last_job: Option<SignedJob>,
26+
session: Option<Session>,
2827
credentials: Option<(Credentials, SshPublicKey)>,
2928
}
3029

@@ -34,7 +33,6 @@ impl Cli {
3433
output,
3534
progress: None,
3635
session: None,
37-
last_job: None,
3836
credentials: None,
3937
}
4038
}
@@ -84,22 +82,23 @@ impl CommandContext for Cli {
8482
}
8583

8684
fn session_id(&self) -> Option<&SessionId> {
87-
self.session.as_ref()
85+
self.session.as_ref().map(|s| s.session_id())
8886
}
8987

90-
fn job_id(&mut self) -> Result<JobId, CommandError> {
91-
let Some(session_id) = self.session.as_ref() else {
92-
return Err(CommandError::MissingSessionId);
93-
};
94-
if let Some(job) = self.last_job.as_ref() {
95-
Ok(session_id.next_job_id(job)?)
88+
fn next_job_id(&self) -> Result<JobId, CommandError> {
89+
if let Some(session) = self.session.as_ref() {
90+
Ok(session.next_job_id()?)
9691
} else {
97-
Ok(session_id.first_job_id())
92+
Err(CommandError::MissingSession)
9893
}
9994
}
10095

101-
fn session_started(&mut self, session_id: &SessionId) -> Result<(), CommandError> {
102-
self.session = Some(session_id.to_owned());
96+
fn session_started(
97+
&mut self,
98+
session_id: &SessionId,
99+
key_id: &KeyId,
100+
) -> Result<(), CommandError> {
101+
self.session = Some(Session::new(session_id.to_owned(), key_id.to_owned()));
103102
match self.get_output_format() {
104103
OutputFormat::Json => println!("{}", json!({"session_started": session_id})),
105104
OutputFormat::Text => println!("✅ Session is now `{session_id}`"),
@@ -108,7 +107,11 @@ impl CommandContext for Cli {
108107
}
109108

110109
fn session_stopped(&mut self, session_id: &SessionId) -> Result<(), CommandError> {
111-
let _ = self.session.take();
110+
if let Some(session) = self.session.as_ref()
111+
&& session.session_id() == session_id
112+
{
113+
let _ = self.session.take();
114+
}
112115
match self.get_output_format() {
113116
OutputFormat::Json => println!("{}", json!({"session_ended": session_id})),
114117
OutputFormat::Text => println!("✅ Ended session `{session_id}`"),
@@ -173,10 +176,19 @@ impl CommandContext for Cli {
173176

174177
// Job management
175178

179+
fn job_started(&mut self, job: &SignedJob) -> Result<(), CommandError> {
180+
if let Some(session) = self.session.as_mut() {
181+
session.job_started(job.to_owned());
182+
Ok(())
183+
} else {
184+
Err(CommandError::MissingSession)
185+
}
186+
}
187+
176188
fn job_stopped(&mut self, job_id: &JobId) -> Result<(), CommandError> {
177189
match self.get_output_format() {
178-
OutputFormat::Json => println!("{job_id}"),
179-
OutputFormat::Text => println!("✅ Aborted job `{job_id}`"),
190+
OutputFormat::Json => println!("{}", json!(job_id)),
191+
OutputFormat::Text => println!("✅ Stopped job `{job_id}`"),
180192
}
181193
Ok(())
182194
}
@@ -285,7 +297,7 @@ impl CommandContext for Cli {
285297
if matches!(self.get_output_format(), OutputFormat::Text) && self.progress.is_none() {
286298
let bar = ProgressBar::new_spinner();
287299
bar.set_elapsed(elapsed);
288-
bar.set_prefix(format!("Waiting for `{job_id}`"));
300+
bar.set_prefix(format!("Waiting for job `{job_id}`"));
289301
bar.set_style(
290302
ProgressStyle::with_template(
291303
"{spinner} \
@@ -386,7 +398,6 @@ impl CommandContext for Cli {
386398
}
387399

388400
fn job_signed(&mut self, job: &SignedJob, show: bool) -> Result<(), CommandError> {
389-
self.last_job = Some(job.to_owned());
390401
if show {
391402
match self.get_output_format() {
392403
OutputFormat::Json => println!("{}", to_json_string(&job)?),
@@ -402,9 +413,10 @@ impl CommandContext for Cli {
402413
OutputFormat::Text => match status {
403414
JobStatus::Unknown { job_id } => println!(
404415
"✅ Job ID:\t{job_id}\n \
405-
Job status:\tUnknown"
416+
Status:\tUnknown"
406417
),
407418
JobStatus::Started {
419+
session_id,
408420
time_started,
409421
stdout_len,
410422
stderr_len,
@@ -414,14 +426,16 @@ impl CommandContext for Cli {
414426
let stderr_len = byte_size(*stderr_len);
415427
println!(
416428
"✅ Job ID:\t{job_id}\n \
417-
Job status:\tStarted\n \
429+
Session ID:\t{session_id}\n \
430+
Status:\tStarted\n \
418431
Started at:\t{time_started}\n \
419432
Stdout len:\t{stdout_len}\n \
420433
Stderr len:\t{stderr_len}"
421434
)
422435
}
423436
JobStatus::Ended {
424437
job: _,
438+
session_id,
425439
time_started,
426440
time_ended,
427441
status: Some(exit_status),
@@ -435,7 +449,8 @@ impl CommandContext for Cli {
435449
let stderr_len = byte_size(*stderr_len);
436450
println!(
437451
"✅ Job ID:\t{job_id}\n \
438-
Job status:\tEnded\n \
452+
Session ID:\t{session_id}\n \
453+
Status:\tEnded\n \
439454
Started at:\t{time_started}\n \
440455
Ended at:\t{time_ended} ({duration})\n \
441456
Status:\t{exit_status}\n \
@@ -447,6 +462,7 @@ impl CommandContext for Cli {
447462
}
448463
JobStatus::Ended {
449464
job: _,
465+
session_id,
450466
time_started,
451467
time_ended,
452468
status: None,
@@ -460,9 +476,10 @@ impl CommandContext for Cli {
460476
let stderr_len = byte_size(*stderr_len);
461477
println!(
462478
"✅ Job ID:\t{job_id}\n \
463-
Job status:\tAborted\n \
479+
Session ID:\t{session_id}\n \
480+
Status:\tStopped\n \
464481
Started at:\t{time_started}\n \
465-
Aborted at:\t{time_ended} ({duration})\n \
482+
Stopped at:\t{time_ended} ({duration})\n \
466483
Stdout len:\t{stdout_len}\n \
467484
Stderr len:\t{stderr_len}\n \
468485
Stdout hash:\t{stdout_hash}\n \

0 commit comments

Comments
 (0)