unset O_DIRECT flag on last chunk if irregular sized

cre4ture · cre4ture · commit f6266b8e2cb4 · 2024-03-17T19:55:18.000+01:00
diff --git a/src/uu/dd/Cargo.toml b/src/uu/dd/Cargo.toml
@@ -20,10 +20,8 @@ gcd = { workspace = true }
 libc = { workspace = true }
 uucore = { workspace = true, features = ["format", "quoting-style"] }
 
-[target.'cfg(any(target_os = "linux"))'.dependencies]
-nix = { workspace = true, features = ["fs"] }
-
 [target.'cfg(any(target_os = "linux", target_os = "android"))'.dependencies]
+nix = { workspace = true, features = ["fs"] }
 signal-hook = { workspace = true }
 
 [[bin]]
diff --git a/src/uu/dd/src/dd.rs b/src/uu/dd/src/dd.rs
@@ -4,6 +4,7 @@
 // file that was distributed with this source code.
 
 // spell-checker:ignore fname, ftype, tname, fpath, specfile, testfile, unspec, ifile, ofile, outfile, fullblock, urand, fileio, atoe, atoibm, behaviour, bmax, bremain, cflags, creat, ctable, ctty, datastructures, doesnt, etoa, fileout, fname, gnudd, iconvflags, iseek, nocache, noctty, noerror, nofollow, nolinks, nonblock, oconvflags, oseek, outfile, parseargs, rlen, rmax, rremain, rsofar, rstat, sigusr, wlen, wstat seekable oconv canonicalized fadvise Fadvise FADV DONTNEED ESPIPE bufferedoutput
+// spell-checker:ignore GETFL SETFL
 
 mod blocks;
 mod bufferedoutput;
@@ -41,6 +42,8 @@ use std::time::{Duration, Instant};
 
 use clap::{crate_version, Arg, Command};
 use gcd::Gcd;
+#[cfg(any(target_os = "android", target_os = "linux"))]
+use nix::fcntl::{fcntl, FcntlArg, OFlag};
 #[cfg(target_os = "linux")]
 use nix::{
     errno::Errno,
@@ -509,6 +512,20 @@ enum Dest {
 }
 
 impl Dest {
+    fn unset_direct(&mut self) -> io::Result<()> {
+        match self {
+            #[cfg(any(target_os = "linux", target_os = "android"))]
+            Self::File(f, _d) => {
+                let mut mode = OFlag::from_bits_retain(fcntl(f.as_raw_fd(), FcntlArg::F_GETFL)?);
+                mode.remove(OFlag::O_DIRECT);
+                nix::fcntl::fcntl(f.as_raw_fd(), FcntlArg::F_SETFL(mode))?;
+            }
+            _ => {}
+        }
+
+        Ok(())
+    }
+
     fn fsync(&mut self) -> io::Result<()> {
         match self {
             Self::Stdout(stdout) => stdout.flush(),
@@ -774,7 +791,14 @@ impl<'a> Output<'a> {
         let mut writes_partial = 0;
         let mut bytes_total = 0;
 
-        for chunk in buf.chunks(self.settings.obs) {
+        let chunk_size = self.settings.obs;
+        for chunk in buf.chunks(chunk_size) {
+            if (self.settings.oflags.direct) && (chunk.len() < chunk_size) {
+                // in case of direct io, only buffers with chunk_size are accepted.
+                // thus, for writing a (last) buffer with irregular length, we need to switch off the direct io.
+                self.dst.unset_direct()?;
+            }
+
             let wlen = self.dst.write(chunk)?;
             if wlen < self.settings.obs {
                 writes_partial += 1;
@@ -877,7 +901,7 @@ impl<'a> BlockWriter<'a> {
 ///
 /// If there is a problem reading from the input or writing to
 /// this output.
-fn dd_copy(mut i: Input, o: Output) -> std::io::Result<()> {
+fn dd_copy(mut i: Input, o: Output) -> UResult<()> {
     // The read and write statistics.
     //
     // These objects are counters, initialized to zero. After each
@@ -992,11 +1016,18 @@ fn dd_copy(mut i: Input, o: Output) -> std::io::Result<()> {
         // best buffer size for reading based on the number of
         // blocks already read and the number of blocks remaining.
         let loop_bsize = calc_loop_bsize(&i.settings.count, &rstat, &wstat, i.settings.ibs, bsize);
-        let rstat_update = read_helper(&mut i, &mut buf, loop_bsize)?;
+        let rstat_update = read_helper(&mut i, &mut buf, loop_bsize)
+            .map_err_context(|| format!("reading, ls: {loop_bsize}, rbt: {}", rstat.bytes_total))?;
         if rstat_update.is_empty() {
             break;
         }
-        let wstat_update = o.write_blocks(&buf)?;
+        let wstat_update = o.write_blocks(&buf).map_err_context(|| {
+            format!(
+                "writing, ls: {}/{loop_bsize}, wbt: {}",
+                buf.len(),
+                wstat.bytes_total
+            )
+        })?;
 
         // Discard the system file cache for the read portion of
         // the input file.
@@ -1047,7 +1078,7 @@ fn finalize<T>(
     prog_tx: &mpsc::Sender<ProgUpdate>,
     output_thread: thread::JoinHandle<T>,
     truncate: bool,
-) -> std::io::Result<()> {
+) -> UResult<()> {
     // Flush the output in case a partial write has been buffered but
     // not yet written.
     let wstat_update = output.flush()?;
@@ -1292,7 +1323,7 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
         }
         None => Output::new_stdout(&settings)?,
     };
-    dd_copy(i, o).map_err_context(|| "IO error".to_string())
+    dd_copy(i, o)
 }
 
 pub fn uu_app() -> Command {
diff --git a/tests/by-util/test_dd.rs b/tests/by-util/test_dd.rs
@@ -1713,3 +1713,45 @@ fn test_reading_partial_blocks_from_fifo_unbuffered() {
     let expected = b"0+2 records in\n0+2 records out\n4 bytes copied";
     assert!(output.stderr.starts_with(expected));
 }
+
+#[cfg(any(target_os = "linux", target_os = "android"))]
+#[test]
+fn test_reading_and_writing_with_direct_flag_from_and_to_files_with_irregular_size() {
+    let ts = TestScenario::new(util_name!());
+    let at = &ts.fixtures;
+
+    let min_direct_block_size = if cfg!(target_os = "android") {
+        4096
+    } else {
+        512
+    };
+
+    let p = |m: i32, o: i32| (min_direct_block_size * m + o).to_string();
+
+    ts.ccmd("truncate")
+        .args(&["-s", p(1, -1).as_str(), "short"])
+        .succeeds();
+    ts.ccmd("truncate")
+        .args(&["-s", p(16, 0).as_str(), "in"])
+        .succeeds();
+    ts.ccmd("truncate")
+        .args(&["-s", p(16, -1).as_str(), "m1"])
+        .succeeds();
+    ts.ccmd("truncate")
+        .args(&["-s", p(16, 1).as_str(), "p1"])
+        .succeeds();
+
+    ts.ucmd()
+        .arg(format!("bs={}", min_direct_block_size))
+        .args(&["if=in", "oflag=direct", "of=out"])
+        .succeeds();
+
+    for testfile in ["short", "m1", "p1"] {
+        at.remove("out");
+        ts.ucmd()
+            .arg(format!("if={testfile}"))
+            .arg(format!("bs={}", min_direct_block_size))
+            .args(&["iflag=direct", "oflag=direct", "of=out"])
+            .succeeds();
+    }
+}
