blob: 7e28f5315a39079bc1fa5bd809b656d9234e71cc [file] [log] [blame]
// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
use std::cmp::min;
use std::fs::File;
use std::io::{self, Error, ErrorKind, Seek, SeekFrom};
use std::os::unix::fs::FileExt;
use crate::fallocate;
use crate::FallocateMode;
/// A trait for deallocating space in a file.
pub trait PunchHole {
/// Replace a range of bytes with a hole.
fn punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()>;
}
impl PunchHole for File {
fn punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()> {
fallocate(self, FallocateMode::PunchHole, true, offset, length as u64)
.map_err(|e| io::Error::from_raw_os_error(e.errno()))
}
}
/// A trait for writing zeroes to a stream.
pub trait WriteZeroes {
/// Write up to `length` bytes of zeroes to the stream, returning how many bytes were written.
fn write_zeroes(&mut self, length: usize) -> io::Result<usize>;
/// Write zeroes to the stream until `length` bytes have been written.
///
/// This method will continuously call `write_zeroes` until the requested
/// `length` is satisfied or an error is encountered.
fn write_zeroes_all(&mut self, mut length: usize) -> io::Result<()> {
while length > 0 {
match self.write_zeroes(length) {
Ok(0) => return Err(Error::from(ErrorKind::WriteZero)),
Ok(bytes_written) => {
length = length
.checked_sub(bytes_written)
.ok_or(Error::from(ErrorKind::Other))?
}
Err(e) => {
if e.kind() != ErrorKind::Interrupted {
return Err(e);
}
}
}
}
Ok(())
}
}
/// A trait for writing zeroes to an arbitrary position in a file.
pub trait WriteZeroesAt {
/// Write up to `length` bytes of zeroes starting at `offset`, returning how many bytes were
/// written.
fn write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize>;
/// Write zeroes starting at `offset` until `length` bytes have been written.
///
/// This method will continuously call `write_zeroes_at` until the requested
/// `length` is satisfied or an error is encountered.
fn write_zeroes_all_at(&mut self, mut offset: u64, mut length: usize) -> io::Result<()> {
while length > 0 {
match self.write_zeroes_at(offset, length) {
Ok(0) => return Err(Error::from(ErrorKind::WriteZero)),
Ok(bytes_written) => {
length = length
.checked_sub(bytes_written)
.ok_or(Error::from(ErrorKind::Other))?;
offset = offset
.checked_add(bytes_written as u64)
.ok_or(Error::from(ErrorKind::Other))?;
}
Err(e) => {
if e.kind() != ErrorKind::Interrupted {
return Err(e);
}
}
}
}
Ok(())
}
}
impl WriteZeroesAt for File {
fn write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize> {
// Try to punch a hole first.
if let Ok(()) = self.punch_hole(offset, length as u64) {
return Ok(length);
}
// fall back to write()
// punch_hole() failed; fall back to writing a buffer of zeroes
// until we have written up to length.
let buf_size = min(length, 0x10000);
let buf = vec![0u8; buf_size];
let mut nwritten: usize = 0;
while nwritten < length {
let remaining = length - nwritten;
let write_size = min(remaining, buf_size);
nwritten += self.write_at(&buf[0..write_size], offset + nwritten as u64)?;
}
Ok(length)
}
}
impl<T: WriteZeroesAt + Seek> WriteZeroes for T {
fn write_zeroes(&mut self, length: usize) -> io::Result<usize> {
let offset = self.seek(SeekFrom::Current(0))?;
let nwritten = self.write_zeroes_at(offset, length)?;
// Advance the seek cursor as if we had done a real write().
self.seek(SeekFrom::Current(nwritten as i64))?;
Ok(length)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs::OpenOptions;
use std::io::{Read, Seek, SeekFrom, Write};
use tempfile::TempDir;
#[test]
fn simple_test() {
let tempdir = TempDir::new().unwrap();
let mut path = tempdir.path().to_owned();
path.push("file");
let mut f = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.open(&path)
.unwrap();
f.set_len(16384).unwrap();
// Write buffer of non-zero bytes to offset 1234
let orig_data = [0x55u8; 5678];
f.seek(SeekFrom::Start(1234)).unwrap();
f.write(&orig_data).unwrap();
// Read back the data plus some overlap on each side
let mut readback = [0u8; 16384];
f.seek(SeekFrom::Start(0)).unwrap();
f.read(&mut readback).unwrap();
// Bytes before the write should still be 0
for read in &readback[0..1234] {
assert_eq!(*read, 0);
}
// Bytes that were just written should be 0x55
for read in &readback[1234..(1234 + 5678)] {
assert_eq!(*read, 0x55);
}
// Bytes after the written area should still be 0
for read in &readback[(1234 + 5678)..] {
assert_eq!(*read, 0);
}
// Overwrite some of the data with zeroes
f.seek(SeekFrom::Start(2345)).unwrap();
f.write_zeroes_all(4321).expect("write_zeroes failed");
// Verify seek position after write_zeroes_all()
assert_eq!(f.seek(SeekFrom::Current(0)).unwrap(), 2345 + 4321);
// Read back the data and verify that it is now zero
f.seek(SeekFrom::Start(0)).unwrap();
f.read(&mut readback).unwrap();
// Bytes before the write should still be 0
for read in &readback[0..1234] {
assert_eq!(*read, 0);
}
// Original data should still exist before the write_zeroes region
for read in &readback[1234..2345] {
assert_eq!(*read, 0x55);
}
// The write_zeroes region should now be zero
for read in &readback[2345..(2345 + 4321)] {
assert_eq!(*read, 0);
}
// Original data should still exist after the write_zeroes region
for read in &readback[(2345 + 4321)..(1234 + 5678)] {
assert_eq!(*read, 0x55);
}
// The rest of the file should still be 0
for read in &readback[(1234 + 5678)..] {
assert_eq!(*read, 0);
}
}
#[test]
fn large_write_zeroes() {
let tempdir = TempDir::new().unwrap();
let mut path = tempdir.path().to_owned();
path.push("file");
let mut f = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.open(&path)
.unwrap();
f.set_len(16384).unwrap();
// Write buffer of non-zero bytes
let orig_data = [0x55u8; 0x20000];
f.seek(SeekFrom::Start(0)).unwrap();
f.write(&orig_data).unwrap();
// Overwrite some of the data with zeroes
f.seek(SeekFrom::Start(0)).unwrap();
f.write_zeroes_all(0x10001).expect("write_zeroes failed");
// Verify seek position after write_zeroes_all()
assert_eq!(f.seek(SeekFrom::Current(0)).unwrap(), 0x10001);
// Read back the data and verify that it is now zero
let mut readback = [0u8; 0x20000];
f.seek(SeekFrom::Start(0)).unwrap();
f.read(&mut readback).unwrap();
// The write_zeroes region should now be zero
for read in &readback[0..0x10001] {
assert_eq!(*read, 0);
}
// Original data should still exist after the write_zeroes region
for read in &readback[0x10001..0x20000] {
assert_eq!(*read, 0x55);
}
}
}