1use std::fs::{File, OpenOptions};
8use std::io::{self, BufWriter, Read, Write};
9use std::path::{Path, PathBuf};
10use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
11use std::sync::mpsc::{self, SyncSender, TrySendError};
12use std::sync::Arc;
13use std::thread::{self, JoinHandle};
14use std::time::Duration;
15
16use crate::error::{Error, Result};
17use crate::fetch::Progress;
18use crate::objects::HashAlgo;
19use crate::pack_stream_scan::PackStreamScanner;
20use crate::progress::{poll_cancel, Cancel, ProgressPhase, ProgressUnit, ReceivedObjects};
21
22pub enum PackReceiveTarget<'a> {
24 Memory(&'a mut Vec<u8>),
26 File(&'a mut dyn Write),
28}
29
30impl PackReceiveTarget<'_> {
31 fn append(&mut self, data: &[u8]) -> Result<()> {
32 match self {
33 Self::Memory(v) => {
34 v.extend_from_slice(data);
35 Ok(())
36 }
37 Self::File(f) => f.write_all(data).map_err(Error::Io),
38 }
39 }
40}
41
42pub fn append_pack_data(
45 data: &[u8],
46 target: &mut PackReceiveTarget<'_>,
47 pending: &mut Vec<u8>,
48 seen_pack: &mut bool,
49) -> Result<u64> {
50 append_pack_data_counted(data, target, pending, seen_pack, None)
51}
52
53fn append_pack_data_counted(
55 data: &[u8],
56 target: &mut PackReceiveTarget<'_>,
57 pending: &mut Vec<u8>,
58 seen_pack: &mut bool,
59 counter: Option<&mut ObjectCounter>,
60) -> Result<u64> {
61 let chunk = if *seen_pack {
62 data
63 } else {
64 pending.extend_from_slice(data);
65 let Some(pos) = find_pack_magic(pending) else {
66 if pending.len() > 3 {
67 let keep_from = pending.len() - 3;
68 pending.drain(..keep_from);
69 }
70 return Ok(0);
71 };
72 *seen_pack = true;
73 &pending[pos..]
74 };
75 let n = chunk.len() as u64;
76 target.append(chunk)?;
77 if let Some(counter) = counter {
78 counter.push(chunk);
79 }
80 if !pending.is_empty() {
81 pending.clear();
82 }
83 Ok(n)
84}
85
86const COUNTER_BATCH: usize = 256 * 1024;
88const COUNTER_BACKLOG: usize = 64;
92
93#[derive(Default)]
95struct CounterState {
96 header_seen: AtomicBool,
97 total: AtomicU64,
98 received: AtomicU64,
99 deltas: AtomicU64,
100 stop: AtomicBool,
101}
102
103struct ObjectCounter {
106 sender: Option<SyncSender<Vec<u8>>>,
107 batch: Vec<u8>,
108 state: Arc<CounterState>,
109 worker: Option<JoinHandle<()>>,
110 last_reported: Option<ReceivedObjects>,
111}
112
113impl ObjectCounter {
114 fn spawn(hash: HashAlgo) -> Option<Self> {
115 let (sender, receiver) = mpsc::sync_channel::<Vec<u8>>(COUNTER_BACKLOG);
116 let state = Arc::new(CounterState::default());
117 let shared = Arc::clone(&state);
118 let worker = thread::Builder::new()
119 .name("grit-pack-count".to_owned())
120 .spawn(move || {
121 let mut scanner = PackStreamScanner::new(hash);
122 for chunk in receiver {
123 if shared.stop.load(Ordering::Relaxed) || scanner.feed(&chunk).is_err() {
124 return;
125 }
126 if let Some(counts) = scanner.counts() {
127 shared.total.store(counts.total, Ordering::Relaxed);
128 shared.deltas.store(counts.deltas, Ordering::Relaxed);
129 shared.received.store(counts.received, Ordering::Relaxed);
130 shared.header_seen.store(true, Ordering::Release);
131 }
132 if scanner.is_complete() {
133 return;
134 }
135 }
136 })
137 .ok()?;
138 Some(Self {
139 sender: Some(sender),
140 batch: Vec::with_capacity(COUNTER_BATCH),
141 state,
142 worker: Some(worker),
143 last_reported: None,
144 })
145 }
146
147 fn push(&mut self, data: &[u8]) {
148 let Some(sender) = &self.sender else {
149 return;
150 };
151 self.batch.extend_from_slice(data);
152 if self.batch.len() < COUNTER_BATCH {
153 return;
154 }
155 let batch = std::mem::replace(&mut self.batch, Vec::with_capacity(COUNTER_BATCH));
156 match sender.try_send(batch) {
157 Ok(()) => {}
158 Err(TrySendError::Full(_)) => {
159 self.state.stop.store(true, Ordering::Relaxed);
162 self.sender = None;
163 }
164 Err(TrySendError::Disconnected(_)) => self.sender = None,
165 }
166 }
167
168 fn report(&mut self, progress: &mut dyn Progress) {
170 if !self.state.header_seen.load(Ordering::Acquire) {
171 return;
172 }
173 let counts = ReceivedObjects {
174 received: self.state.received.load(Ordering::Relaxed),
175 total: self.state.total.load(Ordering::Relaxed),
176 deltas: self.state.deltas.load(Ordering::Relaxed),
177 };
178 if self.last_reported != Some(counts) {
179 self.last_reported = Some(counts);
180 progress.received_objects(counts);
181 }
182 }
183}
184
185impl Drop for ObjectCounter {
186 fn drop(&mut self) {
187 self.state.stop.store(true, Ordering::Relaxed);
188 self.sender = None;
189 if let Some(worker) = self.worker.take() {
190 let _ = worker.join();
191 }
192 }
193}
194
195fn find_pack_magic(buf: &[u8]) -> Option<usize> {
196 buf.windows(4).position(|w| w == b"PACK")
197}
198
199fn note_received_pack_bytes(progress: &mut dyn Progress, active: &mut bool, nbytes: u64) {
200 if nbytes == 0 {
201 return;
202 }
203 if !*active {
204 progress.start_with_unit(ProgressPhase::ReceivingObjects, ProgressUnit::Bytes, None);
205 *active = true;
206 }
207 progress.inc(nbytes);
208}
209
210fn note_received_pack_bytes_coalesced(
211 progress: &mut dyn Progress,
212 active: &mut bool,
213 pending: &mut u64,
214 nbytes: u64,
215) {
216 if nbytes == 0 {
217 return;
218 }
219 *pending += nbytes;
220 const COALESCE: u64 = 256 * 1024;
221 if *pending >= COALESCE {
222 note_received_pack_bytes(progress, active, *pending);
223 *pending = 0;
224 }
225}
226
227fn flush_coalesced_progress(progress: &mut dyn Progress, active: &mut bool, pending: &mut u64) {
228 if *pending > 0 {
229 note_received_pack_bytes(progress, active, *pending);
230 *pending = 0;
231 }
232}
233
234fn finish_receiving(progress: &mut dyn Progress, active: bool) {
235 if active {
236 progress.finish();
237 }
238}
239
240pub fn read_sideband_pack_to(
243 r: &mut dyn Read,
244 target: &mut PackReceiveTarget<'_>,
245 progress: &mut dyn Progress,
246 cancel: &dyn Cancel,
247) -> Result<()> {
248 crate::trace::trace_span!("grit.pack.receive");
249 let mut pending = Vec::new();
250 let mut payload = Vec::with_capacity(65520);
251 read_sideband_pack_tail(
252 r,
253 target,
254 progress,
255 cancel,
256 &mut pending,
257 &mut false,
258 &mut payload,
259 )?;
260 Ok(())
261}
262
263fn pkt_line_payload_len(len_buf: &[u8; 4]) -> Result<usize> {
266 let mut len = 0usize;
267 for &b in len_buf {
268 let digit = match b {
269 b'0'..=b'9' => (b - b'0') as usize,
270 b'a'..=b'f' => (b - b'a' + 10) as usize,
271 b'A'..=b'F' => (b - b'A' + 10) as usize,
272 _ => return Err(Error::Message("bad pkt length".to_owned())),
273 };
274 len = (len << 4) | digit;
275 }
276 Ok(len)
277}
278
279fn read_interrupted(reader: &mut dyn Read, buf: &mut [u8]) -> Result<usize> {
280 const ATTEMPTS: usize = 3;
281 let mut attempt = 0usize;
282 loop {
283 match reader.read(buf) {
284 Ok(n) => return Ok(n),
285 Err(e) if e.kind() == io::ErrorKind::Interrupted && attempt + 1 < ATTEMPTS => {
286 attempt += 1;
287 }
288 Err(e) => return Err(Error::Io(e)),
289 }
290 }
291}
292
293pub fn copy_reader_to_writer(reader: &mut dyn Read, writer: &mut dyn Write) -> Result<u64> {
295 let mut buf = vec![0u8; 1024 * 1024];
296 let mut total = 0u64;
297 loop {
298 let n = read_interrupted(reader, &mut buf)?;
299 if n == 0 {
300 break;
301 }
302 writer.write_all(&buf[..n]).map_err(Error::Io)?;
303 total += n as u64;
304 }
305 Ok(total)
306}
307
308fn sideband_unexpected_eof() -> Error {
309 Error::Io(io::Error::new(
310 io::ErrorKind::UnexpectedEof,
311 "unexpected end of side-band stream",
312 ))
313}
314
315const SIDEBAND_IO_CHUNK: usize = 1024 * 1024;
316
317struct SidebandStreamBuf {
318 storage: Vec<u8>,
319 pos: usize,
320 end: usize,
321}
322
323impl SidebandStreamBuf {
324 fn new() -> Self {
325 Self {
326 storage: vec![0; SIDEBAND_IO_CHUNK],
327 pos: 0,
328 end: 0,
329 }
330 }
331
332 fn compact(&mut self) {
333 if self.pos == 0 {
334 return;
335 }
336 self.storage.copy_within(self.pos..self.end, 0);
337 self.end -= self.pos;
338 self.pos = 0;
339 }
340
341 fn refill(&mut self, r: &mut dyn Read) -> Result<bool> {
342 self.compact();
343 if self.end >= self.storage.len() {
344 return Ok(false);
345 }
346 let n = read_interrupted(r, &mut self.storage[self.end..])?;
347 self.end += n;
348 Ok(n > 0)
349 }
350
351 fn available(&self) -> usize {
352 self.end.saturating_sub(self.pos)
353 }
354}
355
356pub fn read_sideband_pack_tail(
358 r: &mut dyn Read,
359 target: &mut PackReceiveTarget<'_>,
360 progress: &mut dyn Progress,
361 cancel: &dyn Cancel,
362 pending: &mut Vec<u8>,
363 seen_pack: &mut bool,
364 payload_buf: &mut Vec<u8>,
365) -> Result<()> {
366 read_sideband_pack_tail_buffered(
367 r,
368 target,
369 progress,
370 cancel,
371 pending,
372 seen_pack,
373 payload_buf,
374 &mut SidebandStreamBuf::new(),
375 )
376}
377
378#[allow(clippy::too_many_arguments)]
379fn read_sideband_pack_tail_buffered(
380 r: &mut dyn Read,
381 target: &mut PackReceiveTarget<'_>,
382 progress: &mut dyn Progress,
383 cancel: &dyn Cancel,
384 pending: &mut Vec<u8>,
385 seen_pack: &mut bool,
386 _payload_buf: &mut Vec<u8>,
387 stream: &mut SidebandStreamBuf,
388) -> Result<()> {
389 let mut receiving = false;
390 let mut progress_pending = 0u64;
391 let mut counter = (!*seen_pack && progress.wants_received_objects())
393 .then(|| ObjectCounter::spawn(HashAlgo::Sha1))
394 .flatten();
395 loop {
396 poll_cancel(progress, cancel)?;
397 while stream.available() < 4 {
398 if !stream.refill(r)? {
399 if stream.available() == 0 {
400 flush_coalesced_progress(progress, &mut receiving, &mut progress_pending);
401 finish_receiving(progress, receiving);
402 return Ok(());
403 }
404 break;
405 }
406 }
407 if stream.available() < 4 {
408 flush_coalesced_progress(progress, &mut receiving, &mut progress_pending);
409 finish_receiving(progress, receiving);
410 if stream.available() == 0 {
411 return Ok(());
412 }
413 return Err(sideband_unexpected_eof());
414 }
415 let mut len_buf = [0u8; 4];
416 len_buf.copy_from_slice(&stream.storage[stream.pos..stream.pos + 4]);
417 let len = pkt_line_payload_len(&len_buf)?;
418 match len {
419 0 => {
420 stream.pos += 4;
421 if *seen_pack {
422 flush_coalesced_progress(progress, &mut receiving, &mut progress_pending);
423 finish_receiving(progress, receiving);
424 return Ok(());
425 }
426 continue;
427 }
428 1 | 2 => {
429 stream.pos += 4;
430 continue;
431 }
432 n if n <= 4 => {
433 return Err(Error::Message(format!(
434 "invalid pkt-line length in side-band stream: {n}"
435 )))
436 }
437 _ => {}
438 }
439 let payload_len = len - 4;
440 while stream.available() < 4 + payload_len {
441 if !stream.refill(r)? {
442 return Err(sideband_unexpected_eof());
443 }
444 }
445 let payload = &stream.storage[stream.pos + 4..stream.pos + 4 + payload_len];
446 stream.pos += 4 + payload_len;
447 if payload.is_empty() {
448 continue;
449 }
450 match payload[0] {
451 1 => {
452 let data = &payload[1..];
453 let nbytes = if *seen_pack && counter.is_none() {
454 target.append(data)?;
455 data.len() as u64
456 } else {
457 append_pack_data_counted(data, target, pending, seen_pack, counter.as_mut())?
458 };
459 note_received_pack_bytes_coalesced(
460 progress,
461 &mut receiving,
462 &mut progress_pending,
463 nbytes,
464 );
465 if let (Some(counter), true) = (&mut counter, receiving) {
466 counter.report(progress);
467 }
468 }
469 2 => progress.message(&payload[1..]),
470 3 => {
471 return Err(Error::Message(format!(
472 "remote error: {}",
473 String::from_utf8_lossy(&payload[1..]).trim_end()
474 )))
475 }
476 _ => {
477 let nbytes = append_pack_data_counted(
478 payload,
479 target,
480 pending,
481 seen_pack,
482 counter.as_mut(),
483 )?;
484 note_received_pack_bytes_coalesced(
485 progress,
486 &mut receiving,
487 &mut progress_pending,
488 nbytes,
489 );
490 if let (Some(counter), true) = (&mut counter, receiving) {
491 counter.report(progress);
492 }
493 }
494 }
495 }
496}
497
498#[cfg(test)]
500fn read_sideband_pack_tail_read_exact_loop(
501 r: &mut dyn Read,
502 target: &mut PackReceiveTarget<'_>,
503 progress: &mut dyn Progress,
504 pending: &mut Vec<u8>,
505 seen_pack: &mut bool,
506 payload_buf: &mut Vec<u8>,
507 staging_bytes: &mut u64,
508) -> Result<()> {
509 let mut receiving = false;
510 loop {
511 let mut len_buf = [0u8; 4];
512 match r.read_exact(&mut len_buf) {
513 Ok(()) => {}
514 Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
515 Err(e) => return Err(e.into()),
516 }
517 let len = pkt_line_payload_len(&len_buf)?;
518 match len {
519 0 => {
520 if *seen_pack {
521 break;
522 }
523 continue;
524 }
525 1 | 2 => continue,
526 n if n <= 4 => {
527 return Err(Error::Message(format!(
528 "invalid pkt-line length in side-band stream: {n}"
529 )))
530 }
531 _ => {}
532 }
533 let payload_len = len - 4;
534 if payload_buf.len() < payload_len {
535 payload_buf.resize(payload_len, 0);
536 }
537 payload_buf.truncate(payload_len);
538 *staging_bytes += payload_len as u64;
539 r.read_exact(payload_buf)?;
540 if payload_buf.is_empty() {
541 continue;
542 }
543 match payload_buf[0] {
544 1 => {
545 let nbytes = if *seen_pack {
546 let data = &payload_buf[1..];
547 let n = data.len() as u64;
548 target.append(data)?;
549 n
550 } else {
551 append_pack_data(&payload_buf[1..], target, pending, seen_pack)?
552 };
553 note_received_pack_bytes(progress, &mut receiving, nbytes);
554 }
555 2 => progress.message(&payload_buf[1..]),
556 3 => {
557 return Err(Error::Message(format!(
558 "remote error: {}",
559 String::from_utf8_lossy(&payload_buf[1..]).trim_end()
560 )))
561 }
562 _ => {
563 let nbytes = append_pack_data(payload_buf, target, pending, seen_pack)?;
564 note_received_pack_bytes(progress, &mut receiving, nbytes);
565 }
566 }
567 }
568 finish_receiving(progress, receiving);
569 Ok(())
570}
571
572struct SpoolFileGuard {
574 path: PathBuf,
575 delete_on_drop: bool,
576}
577
578impl SpoolFileGuard {
579 fn new(path: PathBuf) -> Self {
580 Self {
581 path,
582 delete_on_drop: true,
583 }
584 }
585
586 fn path(&self) -> &Path {
587 &self.path
588 }
589
590 fn disarm(&mut self) {
591 self.delete_on_drop = false;
592 }
593}
594
595impl Drop for SpoolFileGuard {
596 fn drop(&mut self) {
597 if self.delete_on_drop {
598 let _ = std::fs::remove_file(&self.path);
599 }
600 }
601}
602
603pub struct TempPackReceive {
605 file: BufWriter<File>,
606 spool: SpoolFileGuard,
607 seen_pack: bool,
608 sideband_payload: Vec<u8>,
609}
610
611fn pack_receive_incoming_dir(objects_dir: &Path) -> PathBuf {
612 objects_dir.join("incoming")
613}
614
615fn ensure_pack_receive_incoming_dir(objects_dir: &Path) -> Result<PathBuf> {
616 let dir = pack_receive_incoming_dir(objects_dir);
617 std::fs::create_dir_all(&dir).map_err(Error::Io)?;
618 #[cfg(unix)]
619 {
620 use std::os::unix::fs::PermissionsExt;
621 std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700))
622 .map_err(Error::Io)?;
623 }
624 Ok(dir)
625}
626
627fn create_private_spool_file(path: &Path) -> Result<File> {
628 #[cfg(unix)]
630 {
631 use std::os::unix::fs::OpenOptionsExt;
632 OpenOptions::new()
633 .write(true)
634 .create_new(true)
635 .mode(0o600)
636 .open(path)
637 .map_err(Error::Io)
638 }
639 #[cfg(not(unix))]
640 {
641 OpenOptions::new()
642 .write(true)
643 .create_new(true)
644 .open(path)
645 .map_err(Error::Io)
646 }
647}
648
649fn sync_spool_for_handoff(file: &File) -> Result<()> {
654 #[cfg(unix)]
655 {
656 file.sync_data().map_err(Error::Io)
657 }
658 #[cfg(not(unix))]
659 {
660 file.sync_all().map_err(Error::Io)
661 }
662}
663
664fn assert_received_pack_durable(path: &Path, expected_len: u64) -> Result<()> {
665 const MAX_ATTEMPTS: u32 = 8;
666 for attempt in 0..MAX_ATTEMPTS {
667 if attempt > 0 {
668 thread::sleep(Duration::from_millis(25_u64.saturating_mul(attempt as u64)));
669 }
670 let meta = match std::fs::metadata(path) {
671 Ok(m) => m,
672 Err(e) if e.kind() == io::ErrorKind::NotFound => continue,
673 Err(e) => return Err(Error::Io(e)),
674 };
675 if !meta.is_file() {
676 return Err(Error::ReceivedPackTempMissing {
677 path: path.to_path_buf(),
678 });
679 }
680 if meta.len() != expected_len {
681 if attempt + 1 < MAX_ATTEMPTS {
682 continue;
683 }
684 return Err(Error::CorruptObject(format!(
685 "received pack size mismatch: expected {expected_len}, got {}",
686 meta.len()
687 )));
688 }
689 let mut header = [0u8; 4];
690 let mut f = match File::open(path) {
691 Ok(f) => f,
692 Err(e) if e.kind() == io::ErrorKind::NotFound => continue,
693 Err(e) => return Err(Error::Io(e)),
694 };
695 match f.read_exact(&mut header) {
696 Ok(()) => {}
697 Err(_) if attempt + 1 < MAX_ATTEMPTS => continue,
698 Err(e) => return Err(Error::Io(e)),
699 }
700 if &header != b"PACK" {
701 return Err(Error::CorruptObject(
702 "received data is not a pack stream".to_owned(),
703 ));
704 }
705 return Ok(());
706 }
707 Err(Error::ReceivedPackTempMissing {
708 path: path.to_path_buf(),
709 })
710}
711
712impl TempPackReceive {
713 pub fn create(objects_dir: &Path) -> Result<Self> {
722 let spool_dir = ensure_pack_receive_incoming_dir(objects_dir)?;
723 Self::create_named(|seq| {
724 let path = spool_dir.join(format!("pack-receive-{}_{seq}.tmp", std::process::id()));
725 create_private_spool_file(&path).map(|file| (path, file))
726 })
727 }
728
729 pub(crate) fn create_in_pack_dir(objects_dir: &Path) -> Result<Self> {
734 let pack_dir = objects_dir.join("pack");
735 std::fs::create_dir_all(&pack_dir).map_err(Error::Io)?;
736 Self::create_named(|seq| {
737 let path = pack_dir.join(format!("tmp_pack_{}_{seq}", std::process::id()));
738 OpenOptions::new()
739 .write(true)
740 .create_new(true)
741 .open(&path)
742 .map(|file| (path, file))
743 .map_err(Error::Io)
744 })
745 }
746
747 fn create_named(open: impl Fn(u64) -> Result<(PathBuf, File)>) -> Result<Self> {
749 static SEQ: AtomicU64 = AtomicU64::new(0);
750 for _ in 0..64 {
751 let seq = SEQ.fetch_add(1, Ordering::Relaxed);
752 match open(seq) {
753 Ok((path, file)) => {
754 return Ok(Self {
755 file: BufWriter::with_capacity(4 * 1024 * 1024, file),
756 spool: SpoolFileGuard::new(path),
757 seen_pack: false,
758 sideband_payload: Vec::with_capacity(65520),
759 });
760 }
761 Err(Error::Io(e)) if e.kind() == io::ErrorKind::AlreadyExists => continue,
762 Err(e) => return Err(e),
763 }
764 }
765 Err(Error::Io(io::Error::new(
766 io::ErrorKind::AlreadyExists,
767 "could not allocate a free pack-receive temp file name",
768 )))
769 }
770
771 pub fn read_sideband(
777 &mut self,
778 reader: &mut dyn Read,
779 progress: &mut dyn Progress,
780 cancel: &dyn Cancel,
781 ) -> Result<()> {
782 let mut target = PackReceiveTarget::File(&mut self.file);
783 let mut pending = Vec::new();
784 let mut seen = self.seen_pack;
785 read_sideband_pack_tail(
786 reader,
787 &mut target,
788 progress,
789 cancel,
790 &mut pending,
791 &mut seen,
792 &mut self.sideband_payload,
793 )?;
794 self.seen_pack = seen;
795 self.file.flush().map_err(Error::Io)?;
796 Ok(())
797 }
798
799 pub fn finish(mut self) -> Result<Option<PathBuf>> {
806 self.file.flush().map_err(Error::Io)?;
807 sync_spool_for_handoff(self.file.get_ref())?;
808 let len = self.file.get_ref().metadata().map_err(Error::Io)?.len();
809 if !self.seen_pack || len == 0 {
810 drop(self.file);
811 return Ok(None);
812 }
813 let path = self.spool.path().to_path_buf();
814 drop(self.file);
815 assert_received_pack_durable(&path, len)?;
816 self.spool.disarm();
817 Ok(Some(path))
818 }
819
820 #[must_use]
822 pub fn path(&self) -> &Path {
823 self.spool.path()
824 }
825}
826
827impl Write for TempPackReceive {
828 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
829 if !buf.is_empty() {
830 self.seen_pack = true;
831 }
832 self.file.write(buf)
833 }
834
835 fn flush(&mut self) -> io::Result<()> {
836 self.file.flush()
837 }
838}
839
840#[cfg(test)]
841mod tests {
842 use super::*;
843 use crate::fetch::Progress;
844 use crate::progress::ProgressPhase;
845 use std::io::Cursor;
846
847 #[derive(Default)]
848 struct ReceiveRecorder {
849 started: bool,
850 total: Option<u64>,
851 increments: Vec<u64>,
852 finished: bool,
853 }
854
855 impl Progress for ReceiveRecorder {}
856
857 impl crate::progress::ProgressSink for ReceiveRecorder {
858 fn start(&mut self, phase: ProgressPhase, total: Option<u64>) {
859 assert_eq!(phase, ProgressPhase::ReceivingObjects);
860 self.started = true;
861 self.total = total;
862 }
863 fn inc(&mut self, units: u64) {
864 self.increments.push(units);
865 }
866 fn finish(&mut self) {
867 self.finished = true;
868 }
869 }
870
871 fn sideband_pkt(channel: u8, data: &[u8]) -> Vec<u8> {
872 let payload_len = 1 + data.len();
873 let total_len = payload_len + 4;
874 let mut out = format!("{total_len:04x}").into_bytes();
875 out.push(channel);
876 out.extend_from_slice(data);
877 out
878 }
879
880 #[test]
881 fn reports_received_bytes_per_packet() {
882 let pack1 = b"PACKaaaa";
883 let pack2 = b"bbbbbbbb";
884 let mut stream = Vec::new();
885 stream.extend_from_slice(&sideband_pkt(1, pack1));
886 stream.extend_from_slice(&sideband_pkt(1, pack2));
887 stream.extend_from_slice(b"0000");
888
889 let mut target_buf = Vec::new();
890 let mut target = PackReceiveTarget::Memory(&mut target_buf);
891 let mut rec = ReceiveRecorder::default();
892 read_sideband_pack_to(
893 &mut Cursor::new(stream),
894 &mut target,
895 &mut rec,
896 &crate::progress::NeverCancel,
897 )
898 .expect("read");
899
900 assert!(rec.started);
901 assert_eq!(rec.total, None);
902 assert_eq!(
903 rec.increments.iter().sum::<u64>(),
904 (pack1.len() + pack2.len()) as u64
905 );
906 assert!(rec.finished);
907 assert_eq!(target_buf, b"PACKaaaabbbbbbbb");
908 }
909
910 fn synthetic_band1_sideband_stream(packets: usize, payload_bytes: usize) -> Vec<u8> {
911 let mut stream = Vec::new();
912 stream.extend_from_slice(&sideband_pkt(1, b"PACK"));
913 let payload = vec![0xcd_u8; payload_bytes];
914 for _ in 0..packets {
915 stream.extend_from_slice(&sideband_pkt(1, &payload));
916 }
917 stream.extend_from_slice(b"0000");
918 stream
919 }
920
921 #[test]
926 fn buffered_sideband_demux_avoids_per_packet_staging_buffer() {
927 const PACKETS: usize = 64;
928 const PAYLOAD: usize = 65_515;
929 let stream = synthetic_band1_sideband_stream(PACKETS, PAYLOAD);
930
931 let mut legacy_out = Vec::new();
932 let mut legacy_target = PackReceiveTarget::Memory(&mut legacy_out);
933 let mut legacy_pending = Vec::new();
934 let mut legacy_seen = false;
935 let mut legacy_scratch = Vec::with_capacity(65520);
936 let mut legacy_staging = 0u64;
937 let mut legacy_rec = ReceiveRecorder::default();
938 read_sideband_pack_tail_read_exact_loop(
939 &mut Cursor::new(stream.clone()),
940 &mut legacy_target,
941 &mut legacy_rec,
942 &mut legacy_pending,
943 &mut legacy_seen,
944 &mut legacy_scratch,
945 &mut legacy_staging,
946 )
947 .expect("legacy demux");
948
949 let mut opt_out = Vec::new();
950 let mut opt_target = PackReceiveTarget::Memory(&mut opt_out);
951 let mut opt_pending = Vec::new();
952 let mut opt_seen = false;
953 let mut opt_scratch = Vec::with_capacity(65520);
954 let mut opt_rec = ReceiveRecorder::default();
955 read_sideband_pack_tail(
956 &mut Cursor::new(stream),
957 &mut opt_target,
958 &mut opt_rec,
959 &crate::progress::NeverCancel,
960 &mut opt_pending,
961 &mut opt_seen,
962 &mut opt_scratch,
963 )
964 .expect("buffered demux");
965
966 assert_eq!(legacy_out, opt_out);
967 let min_legacy_staging = PACKETS as u64 * (PAYLOAD as u64 + 1);
968 assert!(
969 legacy_staging >= min_legacy_staging,
970 "legacy path should stage every pkt payload (>= {min_legacy_staging}), got {legacy_staging}"
971 );
972 assert!(
973 opt_scratch.is_empty(),
974 "buffered demux should not use the staging Vec on band-1-only streams"
975 );
976 }
977
978 #[derive(Default)]
980 struct CountRecorder {
981 counts: Vec<ReceivedObjects>,
982 }
983
984 impl crate::progress::ProgressSink for CountRecorder {}
985
986 impl Progress for CountRecorder {
987 fn wants_received_objects(&self) -> bool {
988 true
989 }
990 fn received_objects(&mut self, counts: ReceivedObjects) {
991 self.counts.push(counts);
992 }
993 }
994
995 #[test]
996 fn counts_objects_while_receiving_when_asked() {
997 use grit_test_support::objects::{HashAlgo as FixtureAlgo, ObjectKind as PK, PackBuilder};
998 let mut pb = PackBuilder::new(FixtureAlgo::Sha1);
999 let mut seed = 0x9e37_79b9_u32;
1001 for _ in 0..12 {
1002 let blob: Vec<u8> = (0..64 * 1024)
1003 .map(|_| {
1004 seed ^= seed << 13;
1005 seed ^= seed >> 17;
1006 seed ^= seed << 5;
1007 seed as u8
1008 })
1009 .collect();
1010 pb.add_full(PK::Blob, &blob);
1011 }
1012 let pack = pb.build().bytes;
1013 let mut wire = Vec::new();
1014 for piece in pack.chunks(65_000) {
1015 wire.extend(sideband_pkt(1, piece));
1016 }
1017 wire.extend_from_slice(b"0000");
1018
1019 let mut out = Vec::new();
1020 let mut target = PackReceiveTarget::Memory(&mut out);
1021 let mut rec = CountRecorder::default();
1022 read_sideband_pack_to(
1023 &mut Cursor::new(wire),
1024 &mut target,
1025 &mut rec,
1026 &crate::progress::NeverCancel,
1027 )
1028 .expect("receive");
1029 assert_eq!(out, pack);
1030 assert!(rec.counts.iter().all(|c| c.total == 12 && c.received <= 12));
1033 assert!(rec
1034 .counts
1035 .windows(2)
1036 .all(|w| w[0].received <= w[1].received));
1037 }
1038
1039 #[test]
1040 fn does_not_count_objects_unless_asked() {
1041 let mut wire = sideband_pkt(1, b"PACK\0\0\0\x02\0\0\0\0");
1042 wire.extend_from_slice(b"0000");
1043 let mut out = Vec::new();
1044 let mut target = PackReceiveTarget::Memory(&mut out);
1045 let mut rec = ReceiveRecorder::default();
1046 read_sideband_pack_to(
1047 &mut Cursor::new(wire),
1048 &mut target,
1049 &mut rec,
1050 &crate::progress::NeverCancel,
1051 )
1052 .expect("receive");
1053 assert_eq!(out.len(), 12);
1054 }
1055
1056 #[test]
1057 fn truncated_sideband_payload_is_unexpected_eof() {
1058 let stream = b"0009\x01PACK0010\x01truncated";
1061 let mut target_buf = Vec::new();
1062 let mut target = PackReceiveTarget::Memory(&mut target_buf);
1063 let mut pending = Vec::new();
1064 let mut seen = false;
1065 let mut scratch = Vec::new();
1066 let err = read_sideband_pack_tail(
1067 &mut Cursor::new(stream),
1068 &mut target,
1069 &mut ReceiveRecorder::default(),
1070 &crate::progress::NeverCancel,
1071 &mut pending,
1072 &mut seen,
1073 &mut scratch,
1074 )
1075 .expect_err("truncated side-band payload");
1076 match err {
1077 Error::Io(e) => assert_eq!(e.kind(), io::ErrorKind::UnexpectedEof),
1078 other => panic!("expected Io UnexpectedEof, got {other:?}"),
1079 }
1080 }
1081
1082 #[test]
1083 fn temp_pack_finish_leaves_durable_path_with_bytes() {
1084 let dir = tempfile::tempdir().expect("tempdir");
1085 let objects = dir.path().join("objects");
1086 let mut recv = TempPackReceive::create(&objects).expect("create");
1087 let header = b"PACK";
1088 use std::io::Write;
1089 recv.write_all(header).expect("write");
1090 let path = recv.finish().expect("finish").expect("some pack");
1091 assert!(path.is_file(), "finish must return an on-disk path");
1092 assert_eq!(std::fs::read(&path).expect("read"), header);
1093 let _ = std::fs::remove_file(path);
1094 }
1095
1096 #[test]
1097 fn temp_pack_finish_removes_empty_file() {
1098 let dir = tempfile::tempdir().expect("tempdir");
1099 let objects = dir.path().join("objects");
1100 let recv = TempPackReceive::create(&objects).expect("create");
1101 let path = recv.path().to_path_buf();
1102 assert!(recv.finish().expect("finish").is_none());
1103 assert!(!path.exists());
1104 }
1105
1106 #[test]
1107 fn temp_pack_spools_under_objects_incoming_not_pack() {
1108 let dir = tempfile::tempdir().expect("tempdir");
1109 let objects = dir.path().join("objects");
1110 std::fs::create_dir_all(objects.join("pack")).expect("pack dir");
1111 let recv = TempPackReceive::create(&objects).expect("create");
1112 let path = recv.path().to_path_buf();
1113 assert!(
1114 path.starts_with(objects.join("incoming")),
1115 "temp pack must live under objects/incoming during receive"
1116 );
1117 assert!(
1118 !path.starts_with(objects.join("pack")),
1119 "temp pack must not live under objects/pack during receive"
1120 );
1121 let _ = std::fs::remove_file(path);
1122 }
1123
1124 #[test]
1125 fn temp_pack_drop_after_partial_receive_removes_spool() {
1126 let dir = tempfile::tempdir().expect("tempdir");
1127 let objects = dir.path().join("objects");
1128 let mut recv = TempPackReceive::create(&objects).expect("create");
1129 let path = recv.path().to_path_buf();
1130 use std::io::Write;
1131 recv.write_all(b"PACKpartial").expect("write");
1132 recv.flush().expect("flush");
1133 assert!(path.is_file());
1134 drop(recv);
1135 assert!(
1136 !path.exists(),
1137 "partial spool must be removed when TempPackReceive drops without finish"
1138 );
1139 }
1140
1141 #[test]
1142 fn temp_pack_drop_after_sideband_error_removes_spool() {
1143 let dir = tempfile::tempdir().expect("tempdir");
1144 let objects = dir.path().join("objects");
1145 let mut recv = TempPackReceive::create(&objects).expect("create");
1146 let path = recv.path().to_path_buf();
1147 let mut stream = sideband_pkt(1, b"PACKaaaa");
1148 stream.extend_from_slice(&sideband_pkt(3, b"network interrupted"));
1149 let mut rec = ReceiveRecorder::default();
1150 let err = recv
1151 .read_sideband(
1152 &mut Cursor::new(stream),
1153 &mut rec,
1154 &crate::progress::NeverCancel,
1155 )
1156 .expect_err("remote side-band error");
1157 assert!(matches!(err, Error::Message(_)));
1158 drop(recv);
1159 assert!(
1160 !path.exists(),
1161 "spool must be removed after failed read_sideband"
1162 );
1163 }
1164
1165 #[cfg(unix)]
1166 #[test]
1167 fn temp_pack_spool_file_is_private_to_owner() {
1168 use std::os::unix::fs::PermissionsExt;
1169 let dir = tempfile::tempdir().expect("tempdir");
1170 let objects = dir.path().join("objects");
1171 let recv = TempPackReceive::create(&objects).expect("create");
1172 let mode = std::fs::metadata(recv.path())
1173 .expect("meta")
1174 .permissions()
1175 .mode()
1176 & 0o777;
1177 assert_eq!(mode, 0o600, "spool file must not be world-readable");
1178 let path = recv.path().to_path_buf();
1179 drop(recv);
1180 let _ = std::fs::remove_file(path);
1181 }
1182
1183 #[test]
1184 fn received_pack_durable_check_missing_file_is_typed() {
1185 let dir = tempfile::tempdir().expect("tempdir");
1186 let path = dir.path().join("removed.pack");
1187 std::fs::write(&path, b"PACK").expect("write");
1188 let len = std::fs::metadata(&path).expect("meta").len();
1189 std::fs::remove_file(&path).expect("remove");
1190 let err = assert_received_pack_durable(&path, len).expect_err("missing pack");
1191 assert!(matches!(err, Error::ReceivedPackTempMissing { .. }));
1192 }
1193}