Skip to content

Commit 74c413e

Browse files
plotnickplotnick
committed
Still more fixes
1 parent 04aa41a commit 74c413e

3 files changed

Lines changed: 104 additions & 94 deletions

File tree

‎client/src/commands.rs‎

Lines changed: 74 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use futures::{FutureExt as _, StreamExt as _};
1818
use http::header::WWW_AUTHENTICATE;
1919
use http::status::StatusCode;
2020
use memmap2::Mmap;
21+
use progenitor_client::ResponseValue;
2122
use reqwest::Upgraded;
2223
use rustix::termios::tcgetwinsize;
2324
use thiserror::Error;
@@ -28,7 +29,7 @@ use tokio_tungstenite::WebSocketStream;
2829
use tokio_tungstenite::tungstenite::error::Error as WebSocketError;
2930
use tokio_tungstenite::tungstenite::protocol::Role;
3031

31-
use sush_common::authn::{AuthnError, Challenge, ChallengeResponse, Credentials};
32+
use sush_common::authn::{AuthnError, Challenge, ChallengeResponse, Credentials, Identity};
3233
use sush_common::interactive::InteractiveSessionError;
3334
use sush_common::jobs::JobOutputStream::{self, Stderr, Stdout};
3435
use sush_common::jobs::{
@@ -65,6 +66,9 @@ const DEFAULT_CHUNK_SIZE: ByteSize = ByteSize::mib(32);
6566
/// Default number of elements in a page of results.
6667
const DEFAULT_PAGE_LIMIT: NonZeroU32 = NonZeroU32::new(100).unwrap();
6768

69+
/// Maximum number of authentication retries.
70+
const MAX_AUTHN_RETRIES: usize = 3;
71+
6872
/// Default number of simultaneous downloads for large output.
6973
const PARALLEL_CHUNKS: NonZeroU8 = NonZeroU8::new(8).unwrap();
7074

@@ -464,9 +468,52 @@ impl ClientCommand {
464468
}
465469
}
466470

471+
async fn authz<E>(
472+
ctx: &mut impl CommandContext,
473+
client: &Client,
474+
ssh_auth_sock: &Option<String>,
475+
ssh_key_id: &Option<KeyId>,
476+
response: ResponseValue<E>,
477+
) -> Result<Identity, CommandError> {
478+
let mut ssh_agent = if let Some(ssh_auth_sock) = ssh_auth_sock {
479+
SshAgentConnection::connect(ssh_auth_sock).await?
480+
} else {
481+
return Err(CommandError::MissingSshAuthSock);
482+
};
483+
let public_key = ssh_agent.identity(ssh_key_id.as_ref()).await?;
484+
let challenge = response
485+
.headers()
486+
.get(WWW_AUTHENTICATE)
487+
.ok_or(CommandError::InvalidAuthorization)?
488+
.to_str()
489+
.map_err(|_| CommandError::InvalidAuthorization)?
490+
.parse::<Challenge>()?;
491+
let response = ChallengeResponse::new(challenge);
492+
ctx.please_touch(&public_key)?;
493+
let signed = select! {
494+
s = ssh_agent.sign(response) => s?,
495+
_ = ctrl_c() => return Err(CommandError::Canceled),
496+
};
497+
let verified = signed.verify_with_ssh_public_key(&public_key)?;
498+
let credentials = Credentials::new(verified);
499+
Ok(client
500+
.iam()
501+
.authorization(credentials.to_string())
502+
.body(public_key.to_string())
503+
.send()
504+
.await?
505+
.into_inner())
506+
}
507+
467508
macro_rules! with_authz {
468-
($ctx:ident, $client:ident, $ssh_auth_sock:expr, $ssh_key_id:expr, $authz:ident => $req:expr) => {
509+
($ctx:ident, $client:ident, $ssh_auth_sock:expr, $ssh_key_id:expr, $authz:ident => $req:expr) => {{
510+
let mut i = 0;
469511
loop {
512+
i += 1;
513+
if i == MAX_AUTHN_RETRIES {
514+
return Err(CommandError::InvalidAuthorization);
515+
}
516+
470517
let $authz = $ctx
471518
.get_identity()
472519
.map(|i| i.to_owned().into_credentials().to_string())
@@ -476,40 +523,13 @@ macro_rules! with_authz {
476523
Err(ClientError::ErrorResponse(err))
477524
if err.status() == StatusCode::UNAUTHORIZED =>
478525
{
479-
let mut ssh_agent = if let Some(ssh_auth_sock) = $ssh_auth_sock {
480-
SshAgentConnection::connect(ssh_auth_sock).await?
481-
} else {
482-
return Err(CommandError::MissingSshAuthSock);
483-
};
484-
let public_key = ssh_agent.identity($ssh_key_id).await?;
485-
let challenge = err
486-
.headers()
487-
.get(WWW_AUTHENTICATE)
488-
.ok_or(CommandError::InvalidAuthorization)?
489-
.to_str()
490-
.map_err(|_| CommandError::InvalidAuthorization)?
491-
.parse::<Challenge>()?;
492-
let response = ChallengeResponse::new(challenge);
493-
$ctx.please_touch(&public_key)?;
494-
let signed = select! {
495-
s = ssh_agent.sign(response) => s?,
496-
_ = ctrl_c() => return Err(CommandError::Canceled),
497-
};
498-
let verified = signed.verify_with_ssh_public_key(&public_key)?;
499-
let credentials = Credentials::new(verified);
500-
let identity = $client
501-
.iam()
502-
.authorization(credentials.to_string())
503-
.body(public_key.to_string())
504-
.send()
505-
.await?
506-
.into_inner();
526+
let identity = authz($ctx, $client, $ssh_auth_sock, $ssh_key_id, err).await?;
507527
$ctx.set_identity(Some(identity));
508528
}
509529
res => break res,
510530
}
511531
}
512-
};
532+
}};
513533
}
514534

515535
async fn iam(
@@ -543,7 +563,7 @@ async fn iam(
543563
ctx,
544564
client,
545565
ssh_auth_sock,
546-
ssh_key_id.as_ref(),
566+
ssh_key_id,
547567
authz => client
548568
.identities()
549569
.limit(limit)
@@ -562,7 +582,7 @@ async fn iam(
562582
ctx,
563583
client,
564584
ssh_auth_sock,
565-
ssh_key_id.as_ref(),
585+
ssh_key_id,
566586
authz => client
567587
.identities()
568588
.authorization(&authz)
@@ -578,16 +598,16 @@ async fn iam(
578598
}
579599

580600
IdentityCommand::Login => {
581-
ctx.set_identity(None);
582-
let identity = with_authz!(
583-
ctx,
584-
client,
585-
ssh_auth_sock,
586-
ssh_key_id.as_ref(),
587-
authz => client.iam().authorization(&authz).body(None).send()
588-
)?
589-
.into_inner();
590-
ctx.iam(&identity)
601+
if let Err(ClientError::ErrorResponse(err)) = client.identities().send().await
602+
&& err.status() == StatusCode::UNAUTHORIZED
603+
{
604+
let identity = authz(ctx, client, ssh_auth_sock, ssh_key_id, err).await?;
605+
ctx.iam(&identity)?;
606+
ctx.set_identity(Some(identity));
607+
Ok(())
608+
} else {
609+
Err(CommandError::InvalidAuthorization)
610+
}
591611
}
592612

593613
IdentityCommand::Revoke { revoke } => {
@@ -596,7 +616,7 @@ async fn iam(
596616
ctx,
597617
client,
598618
ssh_auth_sock,
599-
ssh_key_id.as_ref(),
619+
ssh_key_id,
600620
authz => client
601621
.revoke_identity()
602622
.authorization(&authz)
@@ -625,7 +645,7 @@ async fn cert(
625645
ctx,
626646
client,
627647
ssh_auth_sock,
628-
ssh_key_id.as_ref(),
648+
ssh_key_id,
629649
authz => client
630650
.import_cert()
631651
.authorization(authz)
@@ -641,7 +661,7 @@ async fn cert(
641661
ctx,
642662
client,
643663
ssh_auth_sock,
644-
ssh_key_id.as_ref(),
664+
ssh_key_id,
645665
authz => client
646666
.cert_chain()
647667
.authorization(authz)
@@ -667,7 +687,7 @@ async fn session(
667687
ctx,
668688
client,
669689
ssh_auth_sock,
670-
ssh_key_id.as_ref(),
690+
ssh_key_id,
671691
authz => client
672692
.session_start()
673693
.authorization(&authz)
@@ -697,7 +717,7 @@ async fn session(
697717
ctx,
698718
client,
699719
ssh_auth_sock,
700-
ssh_key_id.as_ref(),
720+
ssh_key_id,
701721
authz => client
702722
.session_stop()
703723
.session_id(session_id.clone())
@@ -804,7 +824,7 @@ async fn job(
804824
ctx,
805825
client,
806826
ssh_auth_sock,
807-
ssh_key_id.as_ref(),
827+
ssh_key_id,
808828
authz => client
809829
.job_stop()
810830
.job_id(&job_id)
@@ -819,7 +839,7 @@ async fn job(
819839
ctx,
820840
client,
821841
ssh_auth_sock,
822-
ssh_key_id.as_ref(),
842+
ssh_key_id,
823843
authz => client
824844
.job_status()
825845
.job_id(&job_id)
@@ -1028,7 +1048,7 @@ async fn job_output(
10281048
ctx,
10291049
client,
10301050
ssh_auth_sock,
1031-
ssh_key_id.as_ref(),
1051+
ssh_key_id,
10321052
authz => client
10331053
.job_output_delete()
10341054
.authorization(&authz)
@@ -1051,7 +1071,7 @@ async fn job_output(
10511071
ctx,
10521072
client,
10531073
ssh_auth_sock,
1054-
ssh_key_id.as_ref(),
1074+
ssh_key_id,
10551075
authz => client
10561076
.job_status()
10571077
.authorization(authz)
@@ -1144,7 +1164,7 @@ async fn job_output(
11441164
ctx,
11451165
client,
11461166
ssh_auth_sock,
1147-
ssh_key_id.as_ref(),
1167+
ssh_key_id,
11481168
authz => client
11491169
.job_output()
11501170
.authorization(authz)
@@ -1235,7 +1255,7 @@ async fn job_start_interactive_session(
12351255
ctx,
12361256
client,
12371257
ssh_auth_sock,
1238-
ssh_key_id.as_ref(),
1258+
ssh_key_id,
12391259
authz => client
12401260
.job_start_interactive_session()
12411261
.authorization(&authz)

‎common/src/jobs.rs‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -78,13 +78,11 @@ impl SessionId {
7878
}
7979

8080
pub fn next_job_id(&self, prev_job: &SignedJob) -> Result<JobId, serde_json::Error> {
81-
// TODO: replace with borsh serialize
82-
let prev_job = serde_json::to_string(prev_job)?;
83-
Ok(
84-
id_phrase(U256::from_be_slice(hash(prev_job.as_bytes()).as_bytes()))
85-
.join(WORD_SEPARATOR)
86-
.into(),
87-
)
81+
Ok(id_phrase(U256::from_be_slice(
82+
hash(&prev_job.to_be_signed()).as_bytes(),
83+
))
84+
.join(WORD_SEPARATOR)
85+
.into())
8886
}
8987
}
9088

‎server/src/manager.rs‎

Lines changed: 25 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -395,14 +395,12 @@ impl JobManager {
395395
let wait = params.wait;
396396
let started = {
397397
let job_id = job.job_id().to_owned();
398-
{
399-
let jobs = self.jobs.lock().unwrap();
400-
if jobs.contains_key(&job_id) {
401-
return Err(JobError::InvalidJobId(job_id));
402-
}
403-
if jobs.len() == MAX_JOBS {
404-
return Err(JobError::TooManyJobs(MAX_JOBS));
405-
}
398+
let mut jobs = self.jobs.lock().unwrap();
399+
if jobs.contains_key(&job_id) {
400+
return Err(JobError::InvalidJobId(job_id));
401+
}
402+
if jobs.len() == MAX_JOBS {
403+
return Err(JobError::TooManyJobs(MAX_JOBS));
406404
}
407405

408406
let cert_key_id = job.key_id().to_owned();
@@ -432,10 +430,7 @@ impl JobManager {
432430
session.session_id(),
433431
params,
434432
)?;
435-
self.jobs
436-
.lock()
437-
.unwrap()
438-
.insert(job_id, JobStatus::from(&started));
433+
jobs.insert(job_id, JobStatus::from(&started));
439434
session.job_started(job.into_signed());
440435
started
441436
};
@@ -572,33 +567,30 @@ fn job_ended(
572567
.cloned()
573568
.ok_or_else(|| JobError::InvalidJobId(job_id.to_owned()))?
574569
{
575-
assert!(
576-
jobs.insert(
577-
job_id.to_owned(),
578-
JobStatus::Ended {
579-
job,
580-
session_id,
581-
time_started,
582-
time_ended: time,
583-
status: None,
584-
stdout_len: job_output_len(output_dir, &job_id, Stdout),
585-
stderr_len: job_output_len(output_dir, &job_id, Stderr),
586-
stdout_hash: job_output_hash(output_dir, &job_id, Stdout)?,
587-
stderr_hash: job_output_hash(output_dir, &job_id, Stderr)?,
588-
}
589-
)
590-
.is_some()
591-
);
570+
jobs.insert(
571+
job_id.to_owned(),
572+
JobStatus::Ended {
573+
job,
574+
session_id,
575+
time_started,
576+
time_ended: time,
577+
status: None,
578+
stdout_len: job_output_len(output_dir, &job_id, Stdout),
579+
stderr_len: job_output_len(output_dir, &job_id, Stderr),
580+
stdout_hash: job_output_hash(output_dir, &job_id, Stdout)?,
581+
stderr_hash: job_output_hash(output_dir, &job_id, Stderr)?,
582+
},
583+
)
584+
.ok_or_else(|| JobError::InvalidJobId(job_id))?;
592585
Ok(())
593586
} else {
594587
Err(JobError::InvalidJobId(job_id))
595588
}
596589
}
597590
Ok(ended) => {
598-
assert!(
599-
jobs.insert(ended.job_id().to_owned(), ended.into())
600-
.is_some()
601-
);
591+
let job_id = ended.job_id().to_owned();
592+
jobs.insert(job_id.clone(), ended.into())
593+
.ok_or_else(|| JobError::InvalidJobId(job_id))?;
602594
Ok(())
603595
}
604596
}

0 commit comments

Comments
 (0)