1 // Copyright 2018 The Chromium OS Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 use std::cmp::min; 6 use std::fs::File; 7 use std::io::{self, Error, ErrorKind, Seek, SeekFrom}; 8 use std::os::unix::fs::FileExt; 9 10 use crate::fallocate; 11 use crate::FallocateMode; 12 13 /// A trait for deallocating space in a file. 14 pub trait PunchHole { 15 /// Replace a range of bytes with a hole. punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()>16 fn punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()>; 17 } 18 19 impl PunchHole for File { punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()>20 fn punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()> { 21 fallocate(self, FallocateMode::PunchHole, true, offset, length as u64) 22 .map_err(|e| io::Error::from_raw_os_error(e.errno())) 23 } 24 } 25 26 /// A trait for writing zeroes to a stream. 27 pub trait WriteZeroes { 28 /// Write up to `length` bytes of zeroes to the stream, returning how many bytes were written. write_zeroes(&mut self, length: usize) -> io::Result<usize>29 fn write_zeroes(&mut self, length: usize) -> io::Result<usize>; 30 31 /// Write zeroes to the stream until `length` bytes have been written. 32 /// 33 /// This method will continuously call `write_zeroes` until the requested 34 /// `length` is satisfied or an error is encountered. write_zeroes_all(&mut self, mut length: usize) -> io::Result<()>35 fn write_zeroes_all(&mut self, mut length: usize) -> io::Result<()> { 36 while length > 0 { 37 match self.write_zeroes(length) { 38 Ok(0) => return Err(Error::from(ErrorKind::WriteZero)), 39 Ok(bytes_written) => { 40 length = length 41 .checked_sub(bytes_written) 42 .ok_or_else(|| Error::from(ErrorKind::Other))? 43 } 44 Err(e) => { 45 if e.kind() != ErrorKind::Interrupted { 46 return Err(e); 47 } 48 } 49 } 50 } 51 Ok(()) 52 } 53 } 54 55 /// A trait for writing zeroes to an arbitrary position in a file. 56 pub trait WriteZeroesAt { 57 /// Write up to `length` bytes of zeroes starting at `offset`, returning how many bytes were 58 /// written. write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize>59 fn write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize>; 60 61 /// Write zeroes starting at `offset` until `length` bytes have been written. 62 /// 63 /// This method will continuously call `write_zeroes_at` until the requested 64 /// `length` is satisfied or an error is encountered. write_zeroes_all_at(&mut self, mut offset: u64, mut length: usize) -> io::Result<()>65 fn write_zeroes_all_at(&mut self, mut offset: u64, mut length: usize) -> io::Result<()> { 66 while length > 0 { 67 match self.write_zeroes_at(offset, length) { 68 Ok(0) => return Err(Error::from(ErrorKind::WriteZero)), 69 Ok(bytes_written) => { 70 length = length 71 .checked_sub(bytes_written) 72 .ok_or_else(|| Error::from(ErrorKind::Other))?; 73 offset = offset 74 .checked_add(bytes_written as u64) 75 .ok_or_else(|| Error::from(ErrorKind::Other))?; 76 } 77 Err(e) => { 78 if e.kind() != ErrorKind::Interrupted { 79 return Err(e); 80 } 81 } 82 } 83 } 84 Ok(()) 85 } 86 } 87 88 impl WriteZeroesAt for File { write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize>89 fn write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize> { 90 // Try to use fallocate() first. 91 if fallocate(self, FallocateMode::ZeroRange, true, offset, length as u64).is_ok() { 92 return Ok(length); 93 } 94 95 // fall back to write() 96 // fallocate() failed; fall back to writing a buffer of zeroes 97 // until we have written up to length. 98 let buf_size = min(length, 0x10000); 99 let buf = vec![0u8; buf_size]; 100 let mut nwritten: usize = 0; 101 while nwritten < length { 102 let remaining = length - nwritten; 103 let write_size = min(remaining, buf_size); 104 nwritten += self.write_at(&buf[0..write_size], offset + nwritten as u64)?; 105 } 106 Ok(length) 107 } 108 } 109 110 impl<T: WriteZeroesAt + Seek> WriteZeroes for T { write_zeroes(&mut self, length: usize) -> io::Result<usize>111 fn write_zeroes(&mut self, length: usize) -> io::Result<usize> { 112 let offset = self.seek(SeekFrom::Current(0))?; 113 let nwritten = self.write_zeroes_at(offset, length)?; 114 // Advance the seek cursor as if we had done a real write(). 115 self.seek(SeekFrom::Current(nwritten as i64))?; 116 Ok(length) 117 } 118 } 119 120 #[cfg(test)] 121 mod tests { 122 use super::*; 123 use std::io::{Read, Seek, SeekFrom, Write}; 124 use tempfile::tempfile; 125 126 #[test] simple_test()127 fn simple_test() { 128 let mut f = tempfile().unwrap(); 129 f.set_len(16384).unwrap(); 130 131 // Write buffer of non-zero bytes to offset 1234 132 let orig_data = [0x55u8; 5678]; 133 f.seek(SeekFrom::Start(1234)).unwrap(); 134 f.write_all(&orig_data).unwrap(); 135 136 // Read back the data plus some overlap on each side 137 let mut readback = [0u8; 16384]; 138 f.seek(SeekFrom::Start(0)).unwrap(); 139 f.read_exact(&mut readback).unwrap(); 140 // Bytes before the write should still be 0 141 for read in &readback[0..1234] { 142 assert_eq!(*read, 0); 143 } 144 // Bytes that were just written should be 0x55 145 for read in &readback[1234..(1234 + 5678)] { 146 assert_eq!(*read, 0x55); 147 } 148 // Bytes after the written area should still be 0 149 for read in &readback[(1234 + 5678)..] { 150 assert_eq!(*read, 0); 151 } 152 153 // Overwrite some of the data with zeroes 154 f.seek(SeekFrom::Start(2345)).unwrap(); 155 f.write_zeroes_all(4321).expect("write_zeroes failed"); 156 // Verify seek position after write_zeroes_all() 157 assert_eq!(f.seek(SeekFrom::Current(0)).unwrap(), 2345 + 4321); 158 159 // Read back the data and verify that it is now zero 160 f.seek(SeekFrom::Start(0)).unwrap(); 161 f.read_exact(&mut readback).unwrap(); 162 // Bytes before the write should still be 0 163 for read in &readback[0..1234] { 164 assert_eq!(*read, 0); 165 } 166 // Original data should still exist before the write_zeroes region 167 for read in &readback[1234..2345] { 168 assert_eq!(*read, 0x55); 169 } 170 // The write_zeroes region should now be zero 171 for read in &readback[2345..(2345 + 4321)] { 172 assert_eq!(*read, 0); 173 } 174 // Original data should still exist after the write_zeroes region 175 for read in &readback[(2345 + 4321)..(1234 + 5678)] { 176 assert_eq!(*read, 0x55); 177 } 178 // The rest of the file should still be 0 179 for read in &readback[(1234 + 5678)..] { 180 assert_eq!(*read, 0); 181 } 182 } 183 184 #[test] large_write_zeroes()185 fn large_write_zeroes() { 186 let mut f = tempfile().unwrap(); 187 f.set_len(16384).unwrap(); 188 189 // Write buffer of non-zero bytes 190 let orig_data = [0x55u8; 0x20000]; 191 f.seek(SeekFrom::Start(0)).unwrap(); 192 f.write_all(&orig_data).unwrap(); 193 194 // Overwrite some of the data with zeroes 195 f.seek(SeekFrom::Start(0)).unwrap(); 196 f.write_zeroes_all(0x10001).expect("write_zeroes failed"); 197 // Verify seek position after write_zeroes_all() 198 assert_eq!(f.seek(SeekFrom::Current(0)).unwrap(), 0x10001); 199 200 // Read back the data and verify that it is now zero 201 let mut readback = [0u8; 0x20000]; 202 f.seek(SeekFrom::Start(0)).unwrap(); 203 f.read_exact(&mut readback).unwrap(); 204 // The write_zeroes region should now be zero 205 for read in &readback[0..0x10001] { 206 assert_eq!(*read, 0); 207 } 208 // Original data should still exist after the write_zeroes region 209 for read in &readback[0x10001..0x20000] { 210 assert_eq!(*read, 0x55); 211 } 212 } 213 } 214