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