This repository has no description
0

Configure Feed

Select the types of activity you want to include in your feed.

core / knot2 / crates / knot-pack / src / receiver.rs
4.5 kB 162 lines
1use std::io::{Read, Seek, SeekFrom, Write}; 2use std::path::{Path, PathBuf}; 3 4use tempfile::NamedTempFile; 5 6use crate::error::PackError; 7use crate::frame::ReceiveFramer; 8use crate::ids::MaxWireBytes; 9use crate::meter::PackLimits; 10 11#[derive(Debug)] 12pub enum ReceiveReadError { 13 Io(std::io::Error), 14 Pack(PackError), 15 TooLarge, 16 Truncated, 17} 18 19pub struct ReceivedPack { 20 preamble: Vec<u8>, 21 pack: Option<NamedTempFile>, 22 kind: gix::hash::Kind, 23 total_len: usize, 24} 25 26impl ReceivedPack { 27 pub fn preamble(&self) -> &[u8] { 28 &self.preamble 29 } 30 31 pub fn len(&self) -> usize { 32 self.total_len 33 } 34 35 pub fn is_empty(&self) -> bool { 36 self.total_len == 0 37 } 38 39 pub fn open_pack(&self) -> Result<Option<gix_pack::data::File>, PackError> { 40 match &self.pack { 41 Some(tmp) => gix_pack::data::File::at(tmp.path(), self.kind) 42 .map(Some) 43 .map_err(|error| PackError::Pack(error.to_string())), 44 None => Ok(None), 45 } 46 } 47} 48 49pub struct PackReceiver { 50 dir: PathBuf, 51 file: NamedTempFile, 52 framer: ReceiveFramer, 53 written: u64, 54 limit: MaxWireBytes, 55 complete: Option<usize>, 56 kind: gix::hash::Kind, 57} 58 59impl PackReceiver { 60 pub fn new( 61 dir: &Path, 62 limit: MaxWireBytes, 63 limits: PackLimits, 64 kind: gix::hash::Kind, 65 ) -> std::io::Result<Self> { 66 Ok(Self { 67 dir: dir.to_path_buf(), 68 file: NamedTempFile::new_in(dir)?, 69 framer: ReceiveFramer::new(limits, kind), 70 written: 0, 71 limit, 72 complete: None, 73 kind, 74 }) 75 } 76 77 pub fn write(&mut self, chunk: &[u8]) -> Result<bool, ReceiveReadError> { 78 if self.complete.is_some() { 79 return Ok(true); 80 } 81 if self.written as usize + chunk.len() > self.limit.get() { 82 return Err(ReceiveReadError::TooLarge); 83 } 84 self.file.write_all(chunk).map_err(ReceiveReadError::Io)?; 85 self.written += chunk.len() as u64; 86 self.scan() 87 } 88 89 fn scan(&mut self) -> Result<bool, ReceiveReadError> { 90 if self.written == 0 { 91 return Ok(false); 92 } 93 self.file.flush().map_err(ReceiveReadError::Io)?; 94 match self 95 .framer 96 .advance_file(self.file.as_file(), self.written) 97 .map_err(ReceiveReadError::Pack)? 98 { 99 Some(total) => { 100 self.complete = Some(total); 101 Ok(true) 102 } 103 None => Ok(false), 104 } 105 } 106 107 pub fn finish(mut self) -> Result<ReceivedPack, ReceiveReadError> { 108 let total = match self.complete { 109 Some(total) => total, 110 None => { 111 if self.written == 0 { 112 0 113 } else if self.scan()? { 114 self.complete.expect("scan recorded completion") 115 } else { 116 return Err(ReceiveReadError::Truncated); 117 } 118 } 119 }; 120 let pack_start = self.framer.pack_start().unwrap_or(total); 121 let preamble = 122 read_range(self.file.as_file(), 0..pack_start).map_err(ReceiveReadError::Io)?; 123 let pack = match total > pack_start { 124 true => { 125 let mut tmp = NamedTempFile::new_in(&self.dir).map_err(ReceiveReadError::Io)?; 126 copy_range( 127 self.file.as_file(), 128 pack_start as u64..total as u64, 129 tmp.as_file_mut(), 130 ) 131 .map_err(ReceiveReadError::Io)?; 132 tmp.flush().map_err(ReceiveReadError::Io)?; 133 Some(tmp) 134 } 135 false => None, 136 }; 137 Ok(ReceivedPack { 138 preamble, 139 pack, 140 kind: self.kind, 141 total_len: total, 142 }) 143 } 144} 145 146fn read_range(file: &std::fs::File, range: std::ops::Range<usize>) -> std::io::Result<Vec<u8>> { 147 use std::os::unix::fs::FileExt; 148 let mut buf = vec![0u8; range.end.saturating_sub(range.start)]; 149 file.read_exact_at(&mut buf, range.start as u64)?; 150 Ok(buf) 151} 152 153fn copy_range( 154 src: &std::fs::File, 155 range: std::ops::Range<u64>, 156 dst: &mut std::fs::File, 157) -> std::io::Result<()> { 158 let mut reader = src.try_clone()?; 159 reader.seek(SeekFrom::Start(range.start))?; 160 std::io::copy(&mut reader.take(range.end.saturating_sub(range.start)), dst)?; 161 Ok(()) 162}