corevm-host 0.1.28

use super::*;

/// File system node.
pub enum Node {
	File(File),
	Dir(Dir),
}

impl Node {
	pub fn open<R: ReadBlock>(block_ref: &BlockRef, block_reader: &mut R) -> Result<Self, Error> {
		let mut reader = NodeReader::new(block_ref, block_reader)?;
		let node = match reader.file().main_block().kind() {
			NodeKind::File => {
				let file = reader.into_file();
				Self::File(file)
			},
			NodeKind::Dir => {
				let dir = Dir::decode(&mut reader).map_err(|_| Error::Io)?;
				Self::Dir(dir)
			},
		};
		Ok(node)
	}
}

// TODO @ivan Resolving a path involves reading and decoding a directory in full, although we only
// need one entry. Ideally we need to be able to check that the entry exists without reading and
// decoding all entries.

/// A directory stored in the block storage.
///
/// Stored as file with a different [`NodeKind`] in the main block.
#[derive(Encode, Decode, Debug)]
pub struct Dir(pub VecMap<FileName, BlockRef>);

impl Dir {
	pub fn open<R: ReadBlock>(block_ref: &BlockRef, block_reader: &mut R) -> Result<Self, Error> {
		let mut reader = NodeReader::new(block_ref, block_reader)?;
		if reader.file().main_block().kind() != NodeKind::Dir {
			return Err(Error::Node);
		}
		let dir = Self::decode(&mut reader).map_err(|_| Error::Io)?;
		Ok(dir)
	}
}

/// A file stored in the block storage.
///
/// Might actually refer to a directory. Check [`MainBlock::kind`] to determine that.
pub struct File {
	main_block: MainBlock,
	current_block: Option<(usize, FileBlock)>,
	position: u64,
}

impl File {
	pub fn main_block(&self) -> &MainBlock {
		&self.main_block
	}

	pub fn position(&self) -> u64 {
		self.position
	}

	pub fn seek(&mut self, position: u64) -> Result<(), Error> {
		if position > self.main_block.file_size {
			return Err(Error::Io);
		}
		self.position = position;
		Ok(())
	}

	pub fn open<R: ReadBlock>(main_block_ref: &BlockRef, reader: &mut R) -> Result<Self, Error> {
		let block = reader.read_block(main_block_ref)?;
		let main_block = MainBlock::decode(block)?;
		Ok(Self { main_block, current_block: None, position: 0 })
	}

	pub fn from_main_block(main_block: MainBlock) -> Self {
		Self { main_block, current_block: None, position: 0 }
	}

	pub fn read<R: ReadBlock>(&mut self, buf: &mut [u8], reader: &mut R) -> Result<usize, Error> {
		let n = (buf.len() as u64).min(self.main_block.file_size - self.position);
		if n == 0 {
			return Ok(0);
		}
		let next_position = self.position + n;
		let n = n as usize;
		let first_block_size = self.main_block.first_block.len() as u64;
		let mut i = self.get_block_index(self.position);
		let mut buf_position = 0;
		if i == usize::MAX {
			// Copy from the first block.
			let a = self.position as usize;
			let b = next_position.min(first_block_size) as usize;
			let m = b - a;
			buf[..m].copy_from_slice(&self.main_block.first_block.0[a..b]);
			self.position += m as u64;
			buf_position += m;
			i = 0;
		}
		// Copy from the rest of the blocks.
		while self.position != next_position {
			let block = match &mut self.current_block {
				Some((block_index, block)) if *block_index == i => block,
				block => {
					let data = reader.read_block(&self.main_block.block_refs[i])?;
					let new_block = FileBlock::new(data)?;
					&mut block.insert((i, new_block)).1
				},
			};
			let a = (self.position - first_block_size - i as u64 * MAX_BLOCK_SIZE as u64) as usize;
			let m = ((next_position - self.position) as usize).min(block.len() - a);
			buf[buf_position..buf_position + m].copy_from_slice(&block.0[a..a + m]);
			self.position += m as u64;
			buf_position += m;
			i += 1;
		}
		Ok(n)
	}

	pub fn read_exact<R: ReadBlock>(
		&mut self,
		buf: &mut [u8],
		reader: &mut R,
	) -> Result<(), Error> {
		let n = self.read(buf, reader)?;
		if n != buf.len() {
			return Err(Error::Io);
		}
		Ok(())
	}

	pub fn read_to_end<R: ReadBlock>(
		&mut self,
		buf: &mut Vec<u8>,
		reader: &mut R,
	) -> Result<usize, Error> {
		let remaining = (self.main_block.file_size - self.position) as usize;
		let old_len = buf.len();
		buf.resize(old_len + remaining, 0_u8);
		self.read_exact(&mut buf[old_len..], reader)?;
		Ok(remaining)
	}

	fn get_block_index(&self, mut position: u64) -> usize {
		debug_assert!(position < self.main_block.file_size);
		let first_block_size = self.main_block.first_block.len() as u64;
		if position < first_block_size {
			return usize::MAX;
		}
		position -= first_block_size;
		(position / self.main_block.block_size) as usize
	}
}

/// Reads a node (file or directory) from the block storage.
pub struct NodeReader<R: ReadBlock> {
	reader: R,
	file: File,
}

impl<R: ReadBlock> NodeReader<R> {
	pub fn new(main_block_ref: &BlockRef, mut reader: R) -> Result<Self, Error> {
		let file = File::open(main_block_ref, &mut reader)?;
		Ok(Self { reader, file })
	}

	pub fn from_main_block(main_block: MainBlock, reader: R) -> Self {
		let file = File::from_main_block(main_block);
		Self { reader, file }
	}

	pub fn file(&self) -> &File {
		&self.file
	}

	pub fn into_file(self) -> File {
		self.file
	}

	pub fn into_inner(self) -> R {
		self.reader
	}

	pub fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
		self.file.read(buf, &mut self.reader)
	}

	pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), Error> {
		self.file.read_exact(buf, &mut self.reader)
	}

	pub fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize, Error> {
		self.file.read_to_end(buf, &mut self.reader)
	}

	pub fn seek(&mut self, position: u64) -> Result<(), Error> {
		self.file.seek(position)
	}
}

impl<R: ReadBlock> codec::Input for NodeReader<R> {
	fn remaining_len(&mut self) -> Result<Option<usize>, codec::Error> {
		let remaining = self.file.main_block().file_size() - self.file.position();
		Ok(remaining.try_into().ok())
	}

	fn read(&mut self, into: &mut [u8]) -> Result<(), codec::Error> {
		self.read_exact(into).map_err(|_| "I/o error")?;
		Ok(())
	}
}

/// Write a node (file or directory) to the block storage.
pub struct NodeWriter<W: WriteBlock> {
	writer: W,
	buf: Vec<u8>,
	kind: NodeKind,
	file_size: u64,
	block_size: usize,
	first_block_size: usize,
	position: u64,
	block_refs: Vec<BlockRef>,
	first_block: Vec<u8>,
	service_id: ServiceId,
}

impl<W: WriteBlock> NodeWriter<W> {
	pub fn new(
		service_id: ServiceId,
		writer: W,
		file_size: u64,
		block_size: usize,
	) -> Result<Self, Error> {
		Self::do_new(service_id, writer, NodeKind::File, file_size, block_size)
	}

	pub fn new_dir(
		service_id: ServiceId,
		writer: W,
		file_size: u64,
		block_size: usize,
	) -> Result<Self, Error> {
		Self::do_new(service_id, writer, NodeKind::Dir, file_size, block_size)
	}

	fn do_new(
		service_id: ServiceId,
		writer: W,
		kind: NodeKind,
		file_size: u64,
		block_size: usize,
	) -> Result<Self, Error> {
		let first_block_size = {
			let mut num_blocks = file_size.div_ceil(block_size as u64) as usize;
			let mut first_block_size = MAX_BLOCK_SIZE as u64;
			// Two steps of fixed-point iteration.
			for _ in 0..2 {
				let metadata_len =
					main_block_metadata_encoded_len(file_size, block_size as u64, num_blocks)
						.ok_or(Error::Block)? as u64;
				first_block_size = (MAX_BLOCK_SIZE as u64)
					.checked_sub(metadata_len)
					.ok_or(Error::Block)?
					.min(file_size);
				let new_num_blocks =
					(file_size - first_block_size).div_ceil(block_size as u64) as usize;
				if num_blocks == new_num_blocks {
					break;
				}
				num_blocks = new_num_blocks;
			}
			first_block_size as usize
		};
		let buf_capacity = (block_size as u64).min(file_size) as usize;
		Ok(Self {
			writer,
			kind,
			file_size,
			block_size,
			first_block_size,
			position: 0,
			buf: Vec::with_capacity(buf_capacity.max(first_block_size)),
			block_refs: Vec::new(),
			first_block: Vec::new(),
			service_id,
		})
	}

	pub fn write_all(&mut self, data: &[u8]) -> Result<(), Error> {
		let next_position = self.position + data.len() as u64;
		if next_position > self.file_size {
			return Err(Error::Io);
		}
		let mut slice = data;
		while !slice.is_empty() {
			let max_block_size =
				if self.first_block.is_empty() { self.first_block_size } else { self.block_size };
			let n = slice.len().min(max_block_size - self.buf.len());
			let (chunk, rest) = slice.split_at(n);
			self.buf.extend_from_slice(chunk);
			if self.buf.len() == max_block_size {
				self.write_block()?;
			}
			slice = rest;
		}
		self.position = next_position;
		Ok(())
	}

	pub fn read_from<R: HostFileRead>(&mut self, reader: &mut R) -> Result<(), Error> {
		let mut remaining = reader.remaining_len()?;
		let next_position = self.position.checked_add(remaining).ok_or(Error::Io)?;
		if next_position > self.file_size {
			return Err(Error::Io);
		}
		while remaining != 0 {
			let max_block_size =
				if self.first_block.is_empty() { self.first_block_size } else { MAX_BLOCK_SIZE };
			let old_len = self.buf.len();
			let n = remaining.min(max_block_size as u64 - old_len as u64) as usize;
			self.buf.resize(old_len + n, 0_u8);
			reader.read_exact(&mut self.buf[old_len..])?;
			if self.buf.len() == max_block_size {
				self.write_block()?;
			}
			remaining -= n as u64;
		}
		self.position = next_position;
		Ok(())
	}

	/// Finish writing the node.
	///
	/// Returns main block hash and the underlying block writer.
	pub fn finish(mut self) -> Result<(BlockRef, W), Error> {
		if !self.buf.is_empty() {
			self.write_block()?;
		}
		if self.position != self.file_size {
			return Err(Error::Io);
		}
		// Write main block.
		let main_block = MainBlock {
			kind: self.kind,
			file_size: self.file_size,
			block_size: self.block_size as u64,
			block_refs: core::mem::take(&mut self.block_refs),
			first_block: FileBlock(core::mem::take(&mut self.first_block).into()),
		};
		debug_assert!(validate_main_block(
			main_block.file_size,
			main_block.block_size,
			&main_block.block_refs,
			&main_block.first_block
		)
		.is_ok());
		self.buf.clear();
		main_block.encode_to(&mut self.buf);
		self.writer.write_block(self.service_id, &self.buf[..])?;
		let main_block_ref =
			BlockRef { service_id: self.service_id, hash: Hash::digest(&self.buf[..]) };
		Ok((main_block_ref, self.writer))
	}

	fn write_block(&mut self) -> Result<(), Error> {
		if self.first_block.is_empty() {
			self.first_block = core::mem::take(&mut self.buf);
			self.buf = Vec::with_capacity((self.block_size as u64).min(self.file_size) as usize);
		} else {
			self.block_refs
				.push(BlockRef { service_id: self.service_id, hash: Hash::digest(&self.buf[..]) });
			self.writer.write_block(self.service_id, &self.buf[..])?;
			self.buf.clear();
		}
		Ok(())
	}
}

// This is a workaround for `codec::Output::write` being infallible.
struct FallibleOutput<'a, W: WriteBlock> {
	writer: &'a mut NodeWriter<W>,
	error: Option<Error>,
}

impl<W: WriteBlock> codec::Output for FallibleOutput<'_, W> {
	fn write(&mut self, bytes: &[u8]) {
		if self.error.is_some() {
			return;
		}
		if let Err(e) = self.writer.write_all(bytes) {
			self.error = Some(e);
		}
	}
}

/// Copy the file from the host file system to the block storage.
pub fn copy_file_in<R: HostFileRead, W: WriteBlock>(
	host_file_reader: &mut R,
	service_id: ServiceId,
	block_writer: &mut W,
	block_size: usize,
) -> Result<BlockRef, Error> {
	let file_size = host_file_reader.remaining_len()?;
	let mut writer = NodeWriter::new(service_id, block_writer, file_size, block_size)?;
	writer.read_from(host_file_reader)?;
	let (main_block_ref, _writer) = writer.finish()?;
	Ok(main_block_ref)
}

/// Copy the file referenced by `main_block_ref` from the block storage to the host file system.
pub fn copy_file_out<R: ReadBlock, W: HostFileWrite + ?Sized>(
	main_block_ref: &BlockRef,
	block_reader: &mut R,
	host_file_writer: &mut W,
) -> Result<(), Error> {
	let mut reader = NodeReader::new(main_block_ref, block_reader)?;
	if reader.file.main_block.kind != NodeKind::File {
		// Not a file.
		return Err(Error::Node);
	}
	do_copy_file_out(&mut reader, host_file_writer)?;
	Ok(())
}

fn do_copy_file_out<R: ReadBlock, W: HostFileWrite + ?Sized>(
	reader: &mut NodeReader<R>,
	writer: &mut W,
) -> Result<(), Error> {
	let mut buf = vec![0_u8; MAX_BLOCK_SIZE];
	loop {
		let n = reader.read(&mut buf[..])?;
		if n == 0 {
			break;
		}
		writer.write_all(&buf[..n])?;
	}
	Ok(())
}

/// Create directory in the block storage.
pub fn create_dir<W: WriteBlock>(
	dir: &Dir,
	service_id: ServiceId,
	block_writer: &mut W,
	block_size: usize,
) -> Result<BlockRef, Error> {
	let file_size = dir.encoded_size() as u64;
	let mut writer = NodeWriter::new_dir(service_id, block_writer, file_size, block_size)?;
	let mut output = FallibleOutput { writer: &mut writer, error: None };
	dir.encode_to(&mut output);
	if let Some(e) = output.error {
		return Err(e);
	}
	let (main_block_ref, _writer) = writer.finish()?;
	Ok(main_block_ref)
}

/// Recurisvely copy the directory from the host file system to the block storage.
///
/// Returns the main block hash of the destination directory.
pub fn copy_dir_in<F: HostFileRead, R: HostDirRead<F>, W: WriteBlock>(
	host_dir_reader: R,
	service_id: ServiceId,
	block_writer: &mut W,
	block_size: usize,
) -> Result<BlockRef, Error> {
	let mut dir_stack = Vec::new();
	let mut queue = VecDeque::new();
	let mut last_dir_ref = BlockRef { service_id: 0, hash: Hash::default() };
	let mut visited_dirs = VecSet::new();
	queue.push_back((host_dir_reader, FileName(Default::default()), usize::MAX));
	while let Some((mut host_dir_reader, dir_name, parent_dir_index)) = queue.pop_front() {
		let mut files = VecMap::new();
		let mut subdirs = false;
		while let Some(entry) = host_dir_reader.next_entry() {
			let entry = entry?;
			match entry.kind {
				NodeKind::File => {
					let mut file = host_dir_reader.open_file(&entry.file_name)?;
					let block_ref = copy_file_in(&mut file, service_id, block_writer, block_size)?;
					files.insert(entry.file_name, block_ref);
				},
				NodeKind::Dir => {
					let (another_dir_reader, dir_id) =
						host_dir_reader.open_dir(&entry.file_name)?;
					let visited = match dir_id {
						Some(dir_id) => !visited_dirs.insert(dir_id),
						None => false,
					};
					if visited {
						return Err(Error::Loop);
					}
					queue.push_back((another_dir_reader, entry.file_name, dir_stack.len()));
					subdirs = true;
				},
			}
		}
		if subdirs {
			dir_stack.push((files, dir_name, parent_dir_index));
			continue;
		}
		// Don't use stack for directories that don't contain other directories.
		let dir = Dir(files);
		last_dir_ref = create_dir(&dir, service_id, block_writer, block_size)?;
		if parent_dir_index == usize::MAX {
			continue;
		}
		dir_stack[parent_dir_index].0.insert(dir_name, last_dir_ref);
	}
	while let Some((files, dir_name, parent_dir_index)) = dir_stack.pop() {
		let dir = Dir(files);
		last_dir_ref = create_dir(&dir, service_id, block_writer, block_size)?;
		if parent_dir_index == usize::MAX {
			continue;
		}
		dir_stack[parent_dir_index].0.insert(dir_name, last_dir_ref);
	}
	Ok(last_dir_ref)
}

/// Recursively copy the directory referenced by `main_block_ref` from the block storage to the
/// host file system.
pub fn copy_dir_out<R: ReadBlock, W: HostDirWrite>(
	main_block_ref: &BlockRef,
	block_reader: &mut R,
	host_dir_writer: W,
) -> Result<(), Error> {
	let reader = NodeReader::new(main_block_ref, &mut *block_reader)?;
	if reader.file.main_block.kind != NodeKind::Dir {
		// Not a directory.
		return Err(Error::Node);
	}
	do_copy_dir_out(reader, host_dir_writer)
}

fn do_copy_dir_out<R: ReadBlock, W: HostDirWrite>(
	mut node_reader: NodeReader<R>,
	host_dir_writer: W,
) -> Result<(), Error> {
	let mut queue = VecDeque::new();
	let dir = Dir::decode(&mut node_reader).map_err(|_| Error::Io)?;
	let mut block_reader = node_reader.into_inner();
	queue.push_back((dir, host_dir_writer));
	while let Some((dir, mut host_dir_writer)) = queue.pop_front() {
		for (file_name, hash) in dir.0.iter() {
			let mut reader = NodeReader::new(hash, &mut block_reader)?;
			match reader.file.main_block.kind {
				NodeKind::File => {
					let mut file = host_dir_writer.create_file(file_name)?;
					do_copy_file_out(&mut reader, &mut file)?;
				},
				NodeKind::Dir => {
					let another_dir = Dir::decode(&mut reader).map_err(|_| Error::Io)?;
					let another_dir_writer = host_dir_writer.create_dir(file_name)?;
					queue.push_back((another_dir, another_dir_writer));
				},
			}
		}
	}
	Ok(())
}

/// Recursively traverses file system nodes.
pub struct NodeIter<R: ReadBlock> {
	queue: VecDeque<BlockRef>,
	reader: R,
}

impl<R: ReadBlock> NodeIter<R> {
	pub fn new(main_block_ref: BlockRef, reader: R) -> Self {
		let mut queue = VecDeque::new();
		queue.push_back(main_block_ref);
		Self { queue, reader }
	}
}

impl<R: ReadBlock> Iterator for NodeIter<R> {
	type Item = Result<(BlockRef, File), Error>;

	fn next(&mut self) -> Option<Self::Item> {
		macro_rules! check {
			($body: expr) => {
				match $body {
					Ok(ret) => ret,
					Err(e) => return Some(Err(e)),
				}
			};
		}
		let hash = self.queue.pop_front()?;
		let mut reader = check!(NodeReader::new(&hash, &mut self.reader));
		match reader.file.main_block.kind {
			NodeKind::File => Some(Ok((hash, reader.into_file()))),
			NodeKind::Dir => {
				let dir = check!(Dir::decode(&mut reader).map_err(|_| Error::Io));
				for (_name, hash) in dir.0.into_iter() {
					self.queue.push_back(hash);
				}
				check!(reader.seek(0));
				Some(Ok((hash, reader.into_file())))
			},
		}
	}
}

/// Copy the file or directory (recursively) referenced by `main_block_ref` from the block storage
/// to the host file system.
pub fn copy_out<R: ReadBlock, W: HostWrite>(
	main_block_ref: &BlockRef,
	block_reader: &mut R,
	host_writer: W,
) -> Result<(), Error> {
	let mut reader = NodeReader::new(main_block_ref, block_reader)?;
	match reader.file.main_block.kind {
		NodeKind::File => {
			let mut file_writer = host_writer.into_file_writer()?;
			do_copy_file_out(&mut reader, &mut file_writer)?
		},
		NodeKind::Dir => do_copy_dir_out(reader, host_writer.into_dir_writer()?)?,
	}
	Ok(())
}

/// Fully read the file referenced by `main_block_ref`.
pub fn read<R: ReadBlock>(
	main_block_ref: &BlockRef,
	block_reader: &mut R,
) -> Result<Vec<u8>, Error> {
	let mut reader = NodeReader::new(main_block_ref, block_reader)?;
	let mut buf = Vec::with_capacity(reader.file().main_block().file_size() as usize);
	reader.read_to_end(&mut buf)?;
	Ok(buf)
}

/// Resolve `path` into main block reference.
///
/// `root_dir_ref` refers to the root directory, `current_dir` is the current working directory
/// path. Panics if `current_dir` is not absolute.
pub fn resolve_path<R: ReadBlock>(
	root_dir_ref: BlockRef,
	current_dir: &CStr,
	path: &CStr,
	block_reader: &mut R,
) -> Result<BlockRef, Error> {
	assert!(
		current_dir.to_bytes().is_empty() || matches!(current_dir.to_bytes(), [b'/', ..]),
		"Current directory path must be absolute: {current_dir:?}"
	);
	let mut dir_stack = Vec::new();
	let mut pending_ref = root_dir_ref;
	let cwd_components = match path.to_bytes() {
		[b'/', ..] => c"",
		_ => current_dir,
	}
	.to_bytes()
	.split(|b| *b == b'/');
	let path_components = path.to_bytes().split(|b| *b == b'/');
	for component in cwd_components.chain(path_components) {
		match component {
			b"." | b"" => {},
			b".." =>
				if let Some(dir_ref) = dir_stack.pop() {
					pending_ref = dir_ref;
				},
			name => {
				let dir_ref = pending_ref;
				let dir = Dir::open(&dir_ref, block_reader)?;
				pending_ref = dir.0.get(name).cloned().ok_or(Error::Path)?;
				dir_stack.push(dir_ref);
			},
		}
	}
	Ok(pending_ref)
}