Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966use std::net::IpAddr;use std::path::PathBuf;use std::sync::Arc;use std::time::Duration;
use futures::StreamExt;use knot_acl::{KnotAcl, can_push};use knot_index::Resolved;use knot_lfs::TransferOp;use knot_pack::{PackError, PackLimits, RepoLookup};use knot_runtime::{Clock, HttpTransport};use knot_types::{AccountDid, ObjectFormat, OfferedKey, OwnerDid, RepoDid, RepoRkey};use russh::Channel;use russh::server::Msg;use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};use tokio::runtime::Handle;use tokio::sync::mpsc;
use crate::SshState;
const READ_CHUNK: usize = 64 * 1024;const MAX_UPLOAD_REQUEST: usize = 16 * 1024 * 1024;const RECEIVE_BODY_DEADLINE: Duration = Duration::from_secs(1800);const ARCHIVE_REQUEST_DEADLINE: Duration = Duration::from_secs(60);const LFS_PROGRESS_GRACE: Duration = Duration::from_secs(60);const LFS_PROGRESS_FLOOR_BYTES_PER_SEC: u64 = 1024;const LFS_STALL_TIMEOUT: Duration = Duration::from_secs(120);
fn lfs_within_progress_budget(waited: Duration, moved_bytes: u64) -> bool { waited <= LFS_PROGRESS_GRACE + Duration::from_secs(moved_bytes / LFS_PROGRESS_FLOOR_BYTES_PER_SEC)}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]enum Service { Upload, UploadArchive, Receive, Lfs(TransferOp),}
enum ReadError { Io(std::io::Error), Pack(PackError), TooLarge, Truncated, Deadline,}
enum RepoRef { Did(RepoDid), OwnerRkey(OwnerDid, Vec<RepoRkey>), HandleRkey(knot_types::Handle, Vec<RepoRkey>),}
enum ResolvedRef { Did(RepoDid), OwnerRkey(OwnerDid, Vec<RepoRkey>),}
fn parse_exec(command: &[u8]) -> Option<(Service, RepoRef)> { let text = std::str::from_utf8(command).ok()?.trim(); if let Some(rest) = text.strip_prefix("git-lfs-transfer ") { let (path, op_token) = rest.trim().rsplit_once(' ')?; let op = TransferOp::parse(op_token.trim())?; return Some((Service::Lfs(op), parse_repo_path(path)?)); } let (service, rest) = [ ("git-upload-pack ", Service::Upload), ("git upload-pack ", Service::Upload), ("git-upload-archive ", Service::UploadArchive), ("git upload-archive ", Service::UploadArchive), ("git-receive-pack ", Service::Receive), ("git receive-pack ", Service::Receive), ] .into_iter() .find_map(|(prefix, service)| text.strip_prefix(prefix).map(|rest| (service, rest)))?; Some((service, parse_repo_path(rest)?))}
fn parse_repo_path(raw: &str) -> Option<RepoRef> { let path = raw .trim() .trim_matches('\'') .trim_matches('"') .trim_start_matches('/'); match path.split_once('/') { Some((owner, name)) => { let candidates: Vec<RepoRkey> = RepoRkey::clone_path_candidates(name).collect(); if candidates.is_empty() { return None; } match knot_types::OwnerRef::parse(owner)? { knot_types::OwnerRef::Did(owner) => Some(RepoRef::OwnerRkey(owner, candidates)), knot_types::OwnerRef::Handle(handle) => { Some(RepoRef::HandleRkey(handle, candidates)) } } } None => Some(RepoRef::Did(RepoDid::new(path).ok()?)), }}
fn resolve_repo_ref<H: HttpTransport, C: Clock>( state: &Arc<SshState<H, C>>, repo_ref: ResolvedRef,) -> RepoLookup { let candidate = match repo_ref { ResolvedRef::Did(did) => RepoLookup::Hosted(did), ResolvedRef::OwnerRkey(owner, candidates) => RepoLookup::first(candidates, |rkey| { RepoLookup::from_resolved(state.index.resolve_repo(&owner, &rkey), |found| found) }), }; match candidate { RepoLookup::Hosted(did) => { RepoLookup::from_resolved(state.index.owner_of(&did), |_| did.clone()) } undecided => undecided, }}
pub(crate) async fn run_exec<H: HttpTransport, C: Clock>( state: Arc<SshState<H, C>>, key: Option<OfferedKey>, channel: Channel<Msg>, command: &[u8], protocol_v2: bool, peer: Option<IpAddr>,) { let Some((service, repo_ref)) = parse_exec(command) else { fail(channel, &state.catalog.ssh.unsupported_command.text()).await; return; }; let peer_limiter = match &service { Service::Lfs(_) => state .lfs .as_ref() .map_or(&state.peer_slots, |lfs| &lfs.peer_slots), _ => &state.peer_slots, }; let _peer_guard = match peer_limiter.admit(peer, state.atproto.now()) { Ok(guard) => guard, Err(refusal) => { let reason = match refusal { knot_resource::Refusal::RateLimited => "peer request rate exceeded", knot_resource::Refusal::Saturated => "peer concurrency limit reached", }; tracing::warn!(?peer, reason, "ssh exec rejected"); return fail(channel, &state.catalog.ssh.too_many_operations.text()).await; } }; let resolved_ref = match repo_ref { RepoRef::Did(did) => ResolvedRef::Did(did), RepoRef::OwnerRkey(owner, candidates) => ResolvedRef::OwnerRkey(owner, candidates), RepoRef::HandleRkey(owner_handle, candidates) => { match state .atproto .resolve_handle_to_did(&owner_handle) .await .ok() { Some(did) => ResolvedRef::OwnerRkey(did.into(), candidates), None => { fail(channel, &state.catalog.ssh.repo_not_found.text()).await; return; } } } }; let repo_did = match resolve_repo_ref(&state, resolved_ref) { RepoLookup::Hosted(did) => did, RepoLookup::Unhosted => { fail(channel, &state.catalog.ssh.repo_not_found.text()).await; return; } RepoLookup::Unavailable => { fail(channel, &state.catalog.ssh.index_warming.text()).await; return; } }; let layout = state.layout.clone(); let did = repo_did.clone(); let opened = tokio::task::spawn_blocking(move || layout.open(&did).is_ok()) .await .unwrap_or(false); if !opened { fail(channel, &state.catalog.ssh.repo_not_found.text()).await; return; } match service { Service::Upload => serve_upload(state, channel, repo_did, protocol_v2).await, Service::UploadArchive => serve_upload_archive(state, channel, repo_did).await, Service::Receive => serve_receive(state, key, channel, repo_did).await, Service::Lfs(op) => serve_lfs(state, key, channel, repo_did, op).await, }}
async fn serve_lfs<H: HttpTransport, C: Clock>( state: Arc<SshState<H, C>>, key: Option<OfferedKey>, mut channel: Channel<Msg>, repo_did: RepoDid, op: TransferOp,) { let Some(lfs) = state.lfs.clone() else { return fail(channel, &state.catalog.ssh.lfs_disabled.text()).await; }; if op == TransferOp::Upload { let pusher = resolve_pusher(&state, key.as_ref(), &repo_did).await; let allowed = pusher.as_ref().is_some_and(|did| { let acl = KnotAcl::new(&state.admins, state.admission, &state.index); can_push(&acl, did, &repo_did).is_allowed() }); if !allowed { tracing::warn!( repo = repo_did.as_str(), registered = pusher.is_some(), "ssh lfs upload denied" ); let message = match pusher { None => state.catalog.ssh.key_not_registered.text(), Some(_) => state.catalog.ssh.push_denied.text(), }; return fail(channel, &message).await; } } let permit = match Arc::clone(&lfs.slots).acquire_owned().await { Ok(permit) => permit, Err(_) => return fail(channel, &state.catalog.ssh.shutting_down.text()).await, }; let started = std::time::Instant::now();
let (tx, rx) = mpsc::channel::<Vec<u8>>(8); let writer = Box::pin(channel.make_writer()); let handle = lfs.handle.clone(); let runtime = Handle::current(); let did = repo_did.clone(); let catalog = Arc::clone(&state.catalog); let mut engine = tokio::task::spawn_blocking(move || { let _permit = permit; let output = std::io::BufWriter::new(MeteredWrite::new(runtime.clone(), writer)); knot_lfs::serve_transfer( handle.store.as_ref(), handle.admission.as_ref(), &did, op, &catalog.lfs, MpscRead::new(runtime, rx), output, ) }); let joined = { let reader = channel.make_reader(); tokio::select! { joined = &mut engine => joined, () = pump_input(reader, tx) => engine.await, } }; let status = match joined { Ok(Ok(())) => { tracing::info!( repo = repo_did.as_str(), op = match op { TransferOp::Upload => "upload", TransferOp::Download => "download", }, duration_ms = started.elapsed().as_millis() as u64, "ssh lfs transfer finished" ); 0 } Ok(Err(fault)) => { tracing::warn!(repo = repo_did.as_str(), %fault, "ssh lfs transfer failed"); 1 } Err(join) => { tracing::error!(repo = repo_did.as_str(), %join, "ssh lfs transfer task panicked"); 1 } }; finish(channel, status).await;}
async fn pump_input<R: AsyncRead + Unpin>(reader: R, tx: mpsc::Sender<Vec<u8>>) { use futures::TryStreamExt; let _ = tokio_util::io::ReaderStream::with_capacity(reader, READ_CHUNK) .map_err(|_| ()) .try_for_each(|chunk| { let tx = &tx; async move { match chunk.is_empty() { true => Ok(()), false => tx.send(chunk.to_vec()).await.map_err(|_| ()), } } }) .await;}
fn stalled(direction: &'static str) -> std::io::Error { std::io::Error::other(format!("lfs {direction} stalled past the idle timeout"))}
struct MpscRead { runtime: Handle, rx: mpsc::Receiver<Vec<u8>>, buffer: Vec<u8>, offset: usize, waited: Duration, received: u64,}
impl MpscRead { fn new(runtime: Handle, rx: mpsc::Receiver<Vec<u8>>) -> Self { Self { runtime, rx, buffer: Vec::new(), offset: 0, waited: Duration::ZERO, received: 0, } }}
impl std::io::Read for MpscRead { fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> { if self.offset >= self.buffer.len() { let started = std::time::Instant::now(); let rx = &mut self.rx; let received = self .runtime // I know I know, but these aren't runtime workers here .block_on(async { tokio::time::timeout(LFS_STALL_TIMEOUT, rx.recv()).await }); match received { Ok(Some(chunk)) => { self.waited += started.elapsed(); self.received += chunk.len() as u64; if !lfs_within_progress_budget(self.waited, self.received) { return Err(std::io::Error::other( "lfs input trickles below the progress floor", )); } self.buffer = chunk; self.offset = 0; } Ok(None) => return Ok(0), Err(_) => return Err(stalled("input")), } } let take = out.len().min(self.buffer.len() - self.offset); out[..take].copy_from_slice(&self.buffer[self.offset..self.offset + take]); self.offset += take; Ok(take) }}
struct MeteredWrite<W> { runtime: Handle, inner: W, waited: Duration, written: u64,}
impl<W: AsyncWrite + Unpin> MeteredWrite<W> { fn new(runtime: Handle, inner: W) -> Self { Self { runtime, inner, waited: Duration::ZERO, written: 0, } }
fn charge(&mut self, started: std::time::Instant) -> std::io::Result<()> { self.waited += started.elapsed(); match lfs_within_progress_budget(self.waited, self.written) { true => Ok(()), false => Err(std::io::Error::other( "lfs output trickles below the progress floor", )), } }}
impl<W: AsyncWrite + Unpin> std::io::Write for MeteredWrite<W> { fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { let started = std::time::Instant::now(); let inner = &mut self.inner; let wrote = self .runtime .block_on(async { tokio::time::timeout(LFS_STALL_TIMEOUT, inner.write(buf)).await }) .map_err(|_| stalled("output"))??; self.written += wrote as u64; self.charge(started).map(|()| wrote) }
fn flush(&mut self) -> std::io::Result<()> { let started = std::time::Instant::now(); let inner = &mut self.inner; self.runtime .block_on(async { tokio::time::timeout(LFS_STALL_TIMEOUT, inner.flush()).await }) .map_err(|_| stalled("output"))??; self.charge(started) }}
async fn serve_upload_archive<H: HttpTransport, C: Clock>( state: Arc<SshState<H, C>>, mut channel: Channel<Msg>, repo_did: RepoDid,) { let request = { let mut reader = channel.make_reader(); tokio::time::timeout(ARCHIVE_REQUEST_DEADLINE, read_archive_request(&mut reader)).await }; let request = match request { Ok(Ok(request)) => request, Ok(Err(())) => return fail(channel, &state.catalog.ssh.archive_malformed.text()).await, Err(_) => return fail(channel, &state.catalog.ssh.archive_timeout.text()).await, };
let permit = state.slots.pack.acquire().await; let (tx, mut rx) = mpsc::channel::<Vec<u8>>(16); let layout = state.layout.clone(); let did = repo_did.clone(); let handle = tokio::task::spawn_blocking(move || -> Result<(), PackError> { let _permit = permit; let repo = layout.open(&did)?; let mut sink = |chunk: &[u8]| -> std::io::Result<()> { tx.blocking_send(chunk.to_vec()) .map_err(|_| std::io::Error::other("client disconnected")) }; knot_pack::upload_archive_streamed(&repo, &request, &mut sink) });
let mut writer = channel.make_writer(); let mut forward = Ok(()); while let Some(chunk) = rx.recv().await { if writer.write_all(&chunk).await.is_err() { forward = Err(()); break; } } drop(rx); let produced = handle.await; match &produced { Ok(Err(error)) => { tracing::warn!(repo = repo_did.as_str(), %error, "upload-archive failed") } Err(join) => { tracing::error!(repo = repo_did.as_str(), %join, "upload-archive task panicked") } Ok(Ok(())) => {} } match (forward, produced) { (Ok(()), Ok(Ok(()))) if writer.flush().await.is_ok() => finish(channel, 0).await, _ => fail(channel, &state.catalog.ssh.archive_failed.text()).await, }}
async fn read_archive_request<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec<u8>, ()> { let mut buf = Vec::new(); loop { if knot_pack::archive_request_complete(&buf).is_some() { return Ok(buf); } match read_chunk(reader, &mut buf, MAX_UPLOAD_REQUEST).await { Ok(true) => {} Ok(false) | Err(_) => return Err(()), } }}
async fn serve_upload<H: HttpTransport, C: Clock>( state: Arc<SshState<H, C>>, mut channel: Channel<Msg>, repo_did: RepoDid, protocol_v2: bool,) { let advert = { let layout = state.layout.clone(); let did = repo_did.clone(); tokio::task::spawn_blocking(move || -> Result<Vec<u8>, PackError> { let repo = layout.open(&did)?; if protocol_v2 { knot_pack::advertise_upload_ssh(&repo) } else { knot_pack::advertise_upload_v0_ssh(&repo) } }) .await }; let advert = match advert { Ok(Ok(bytes)) => bytes, _ => return fail(channel, &state.catalog.ssh.advertise_failed.text()).await, };
let mut writer = channel.make_writer(); if writer.write_all(&advert).await.is_err() || writer.flush().await.is_err() { return; }
let outcome = { let mut reader = channel.make_reader(); if protocol_v2 { upload_loop_v2(&state, &repo_did, &mut reader, &mut writer).await } else { upload_loop_v0(&state, &repo_did, &mut reader, &mut writer).await } }; let status = match outcome { Ok(()) => 0, Err(()) => 1, }; finish(channel, status).await;}
async fn upload_loop_v2<H, C, R, W>( state: &Arc<SshState<H, C>>, repo_did: &RepoDid, reader: &mut R, writer: &mut W,) -> Result<(), ()>where H: HttpTransport, C: Clock, R: AsyncRead + Unpin, W: AsyncWriteExt + Unpin,{ let mut buf = Vec::new(); let mut framer = knot_pack::UploadFramer::new(); loop { if let Some(len) = framer.advance(&buf) { let request: Vec<u8> = buf.drain(..len).collect(); stream_upload(state, repo_did, request, writer).await?; framer = knot_pack::UploadFramer::new(); continue; } match read_chunk(reader, &mut buf, MAX_UPLOAD_REQUEST).await { Ok(true) => {} Ok(false) => return Ok(()), Err(_) => return Err(()), } }}
async fn upload_loop_v0<H, C, R, W>( state: &Arc<SshState<H, C>>, repo_did: &RepoDid, reader: &mut R, writer: &mut W,) -> Result<(), ()>where H: HttpTransport, C: Clock, R: AsyncRead + Unpin, W: AsyncWriteExt + Unpin,{ let mut buf = Vec::new(); let mut framer = knot_pack::UploadFramer::new(); let mut naks_sent = 0usize; loop { if let Some(len) = framer.advance(&buf) { let request: Vec<u8> = buf.drain(..len).collect(); return stream_upload(state, repo_did, request, writer).await; } let needed = framer.unanswered_flushes(); if naks_sent < needed { let nak = knot_pack::upload_v0_nak(); if writer.write_all(&nak).await.is_err() || writer.flush().await.is_err() { return Err(()); } naks_sent += 1; continue; } match read_chunk(reader, &mut buf, MAX_UPLOAD_REQUEST).await { Ok(true) => {} Ok(false) => return Ok(()), Err(_) => return Err(()), } }}
async fn stream_upload<H, C, W>( state: &Arc<SshState<H, C>>, repo_did: &RepoDid, request: Vec<u8>, writer: &mut W,) -> Result<(), ()>where H: HttpTransport, C: Clock, W: AsyncWriteExt + Unpin,{ let permit = state.slots.pack.acquire().await; let (tx, mut rx) = mpsc::channel::<Vec<u8>>(16); let layout = state.layout.clone(); let did = repo_did.clone(); let catalog = Arc::clone(&state.catalog); let knot = state.hostname.clone(); let handle = tokio::task::spawn_blocking(move || -> Result<(), PackError> { let _permit = permit; let repo = layout.open(&did)?; let mut sink = |chunk: &[u8]| -> std::io::Result<()> { tx.blocking_send(chunk.to_vec()) .map_err(|_| std::io::Error::other("client disconnected")) }; knot_pack::upload_pack_streamed(&repo, &request, &catalog.fetch, &knot, &mut sink) });
let mut forward = Ok(()); while let Some(chunk) = rx.recv().await { if writer.write_all(&chunk).await.is_err() { forward = Err(()); break; } } drop(rx); match (forward, handle.await) { (Ok(()), Ok(Ok(()))) => writer.flush().await.map_err(|_| ()), _ => Err(()), }}
async fn serve_receive<H: HttpTransport, C: Clock>( state: Arc<SshState<H, C>>, key: Option<OfferedKey>, mut channel: Channel<Msg>, repo_did: RepoDid,) { let advert = { let layout = state.layout.clone(); let did = repo_did.clone(); tokio::task::spawn_blocking(move || -> Result<(Vec<u8>, ObjectFormat), PackError> { let repo = layout.open(&did)?; let bytes = knot_pack::advertise_receive_ssh(&repo)?; Ok((bytes, repo.object_format())) }) .await }; let (advert, object_format) = match advert { Ok(Ok(pair)) => pair, _ => return fail(channel, &state.catalog.ssh.advertise_failed.text()).await, };
let mut writer = channel.make_writer(); if writer.write_all(&advert).await.is_err() || writer.flush().await.is_err() { return; }
let pusher = resolve_pusher(&state, key.as_ref(), &repo_did).await; let allowed = |did: &AccountDid| { let acl = KnotAcl::new(&state.admins, state.admission, &state.index); can_push(&acl, did, &repo_did).is_allowed() }; let committer = match pusher { Some(did) if allowed(&did) => did, Some(_) => { tracing::warn!( repo = repo_did.as_str(), registered = true, "ssh push denied" ); return fail(channel, &state.catalog.ssh.push_denied.text()).await; } None => { tracing::warn!( repo = repo_did.as_str(), registered = false, "ssh push denied" ); return fail(channel, &state.catalog.ssh.key_not_registered.text()).await; } };
let _receive_permit = state.slots.receive.acquire().await;
let limits = state.limits; let body = { let mut reader = channel.make_reader(); let dir = state.layout.scratch_dir().to_path_buf(); match tokio::time::timeout( RECEIVE_BODY_DEADLINE, read_receive( &mut reader, dir, state.max_pack_bytes, limits, object_format, ), ) .await { Ok(result) => result, Err(_) => Err(ReadError::Deadline), } }; let body = match body { Ok(body) => body, Err(ReadError::TooLarge) => { return fail(channel, &state.catalog.ssh.push_too_large.text()).await; } Err(ReadError::Deadline) => { return fail(channel, &state.catalog.ssh.receive_deadline.text()).await; } Err(ReadError::Pack(error)) => { tracing::warn!(repo = repo_did.as_str(), %error, "receive framing failed"); return fail(channel, &state.catalog.ssh.malformed_pack.text()).await; } Err(ReadError::Io(error)) => { tracing::warn!(repo = repo_did.as_str(), %error, "receive read error"); return fail(channel, &state.catalog.ssh.receive_read_error.text()).await; } Err(ReadError::Truncated) => { return fail(channel, &state.catalog.ssh.receive_ended_early.text()).await; } }; if body.is_empty() { return finish(channel, 0).await; }
let _pack_permit = state.slots.pack.acquire().await; let landed = knot_receive::land(knot_receive::Push { layout: &state.layout, repo_did: &repo_did, received: body, limits: state.limits, knot_actor: state.knot_actor.clone(), committer, events: Arc::clone(&state.events), index: &state.index, atproto: &state.atproto, resolve_slots: &state.slots.resolve, appview: &state.appview, maintenance: &state.maintenance, hostname: &state.hostname, languages_push_budget: state.languages_push_budget, catalog: Arc::clone(&state.catalog), ci_logs: state.ci_logs.clone(), }) .await; match landed { Ok(framed) => { let _ = writer.write_all(&framed).await; let _ = writer.flush().await; finish(channel, 0).await; } Err(error) => { tracing::warn!(repo = repo_did.as_str(), %error, "receive-pack failed"); fail(channel, &state.catalog.ssh.receive_failed.text()).await; } }}
pub(crate) async fn run_greeting<H: HttpTransport, C: Clock>( state: Arc<SshState<H, C>>, key: Option<OfferedKey>, channel: Channel<Msg>,) { let who = greeting_identity(&state, key.as_ref()).await; let greeting = state.catalog.ssh.greeting.lines(|key| match key { knot_messages::GreetingKey::User => who.clone(), knot_messages::GreetingKey::Knot => state.hostname.as_str().to_string(), }); if greeting.is_empty() { return finish(channel, 0).await; } let body = greeting.join("\r\n"); let _ = channel .extended_data_bytes(1, format!("{body}\r\n").into_bytes()) .await; finish(channel, 0).await;}
async fn greeting_identity<H: HttpTransport, C: Clock>( state: &Arc<SshState<H, C>>, key: Option<&OfferedKey>,) -> String { let Some(did) = key.and_then(|key| state.roster.did_for(key)) else { return "there".to_string(); }; match knot_receive::resolve_handle(&state.atproto, &state.slots.resolve, &did).await { Some(handle) => format!("@{}", handle.as_str()), None => did.as_str().to_string(), }}
async fn resolve_pusher<H: HttpTransport, C: Clock>( state: &Arc<SshState<H, C>>, key: Option<&OfferedKey>, repo: &RepoDid,) -> Option<AccountDid> { let key = key?; let owner = match state.index.owner_of(repo) { Resolved::Ready(Some(owner)) => Some(AccountDid::from(owner)), _ => None, }; { let index = Arc::clone(&state.index); let target = repo.clone(); let _ = tokio::task::spawn_blocking(move || index.ensure_collaborators(&target)).await; } let collaborators = match state.index.collaborators_of(repo) { Resolved::Ready(collaborators) => collaborators, _ => Vec::new(), }; let candidates: Vec<AccountDid> = owner.into_iter().chain(collaborators).collect(); if let Resolved::Ready(Some(cached)) = state.index.owner_of_key(key) && candidates.contains(&cached) { return Some(cached); } let _permit = state.slots.resolve.acquire().await; let matches = futures::stream::iter(candidates).filter_map(|did| async move { let keys = state.atproto.resolve_pubkeys(&did).await.ok()?; keys.iter() .for_each(|resolved| state.index.cache_key(resolved.clone(), &did)); keys.iter().any(|resolved| resolved == key).then_some(did) }); futures::pin_mut!(matches); matches.next().await}
async fn read_chunk<R: AsyncRead + Unpin>( reader: &mut R, buf: &mut Vec<u8>, limit: usize,) -> Result<bool, ReadError> { let mut chunk = [0u8; READ_CHUNK]; let read = reader.read(&mut chunk).await.map_err(ReadError::Io)?; if read == 0 { return Ok(false); } buf.extend_from_slice(&chunk[..read]); if buf.len() > limit { return Err(ReadError::TooLarge); } Ok(true)}
async fn read_receive<R: AsyncRead + Unpin>( reader: &mut R, dir: PathBuf, limit: knot_pack::MaxWireBytes, limits: PackLimits, format: ObjectFormat,) -> Result<knot_pack::ReceivedPack, ReadError> { let (tx, rx) = mpsc::channel::<Vec<u8>>(8); let mut framer = tokio::task::spawn_blocking(move || frame_receive(rx, dir, limit, limits, format)); let mut chunk = [0u8; READ_CHUNK]; let mut io_error = None; loop { tokio::select! { biased; framed = &mut framer => return join_framed(framed, io_error), read = reader.read(&mut chunk) => match read { Ok(0) => break, Ok(read) => { if tx.send(chunk[..read].to_vec()).await.is_err() { break; } } Err(error) => { io_error = Some(error); break; } }, } } drop(tx); join_framed(framer.await, io_error)}
fn join_framed( framed: Result<Result<knot_pack::ReceivedPack, ReadError>, tokio::task::JoinError>, io_error: Option<std::io::Error>,) -> Result<knot_pack::ReceivedPack, ReadError> { match framed { Ok(Ok(body)) => Ok(body), Ok(Err(ReadError::Truncated)) => { Err(io_error.map(ReadError::Io).unwrap_or(ReadError::Truncated)) } Ok(Err(other)) => Err(other), Err(_) => Err(ReadError::Truncated), }}
fn read_error(error: knot_pack::ReceiveReadError) -> ReadError { match error { knot_pack::ReceiveReadError::Io(error) => ReadError::Io(error), knot_pack::ReceiveReadError::Pack(error) => ReadError::Pack(error), knot_pack::ReceiveReadError::TooLarge => ReadError::TooLarge, knot_pack::ReceiveReadError::Truncated => ReadError::Truncated, }}
fn frame_receive( mut rx: mpsc::Receiver<Vec<u8>>, dir: PathBuf, limit: knot_pack::MaxWireBytes, limits: PackLimits, format: ObjectFormat,) -> Result<knot_pack::ReceivedPack, ReadError> { let mut receiver = knot_pack::PackReceiver::new(&dir, limit, limits, format.kind()).map_err(ReadError::Io)?; loop { match rx.blocking_recv() { Some(chunk) => { if receiver.write(&chunk).map_err(read_error)? { return receiver.finish().map_err(read_error); } } None => return receiver.finish().map_err(read_error), } }}
async fn fail(channel: Channel<Msg>, message: &str) { let _ = channel .extended_data_bytes(1, format!("{message}\n").into_bytes()) .await; finish(channel, 1).await;}
async fn finish(channel: Channel<Msg>, status: u32) { let _ = channel.exit_status(status).await; let _ = channel.eof().await; let _ = channel.close().await;}
#[cfg(test)]mod tests { use super::*;
#[test] fn the_lfs_progress_budget_spares_slow_links_and_cuts_trickles() { assert!(lfs_within_progress_budget(Duration::from_secs(59), 0)); assert!(!lfs_within_progress_budget(Duration::from_secs(61), 0)); assert!(lfs_within_progress_budget( Duration::from_secs(50_000), 5 * 1024 * 1024 * 1024 )); assert!(!lfs_within_progress_budget(Duration::from_secs(1_000), 10)); }
#[test] fn the_repo_path_parser_separates_dids_from_handles() { assert!(matches!( parse_repo_path("did:plc:nel/squid"), Some(RepoRef::OwnerRkey(..)) )); assert!(matches!( parse_repo_path("nel.pet/squid"), Some(RepoRef::HandleRkey(..)) )); assert!(matches!( parse_repo_path("did:plc:barnacle"), Some(RepoRef::Did(_)) )); assert!(parse_repo_path("did:nonsense/squid").is_none()); assert!(parse_repo_path("nel.pet").is_none()); }}