From e637ddbdfe21617f0bb9cc6aef4f88dcac3fde8d Mon Sep 17 00:00:00 2001
From: "Hui, Chunyang" <sanqian.hcy@antfin.com>
Date: Fri, 12 Jan 2024 09:24:01 +0000
Subject: [PATCH] Fix mmap file flush exceeding the file length

---
 src/libos/src/vm/vm_area.rs | 38 +++++++++++++++++--
 test/mmap/main.c            | 73 +++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+), 3 deletions(-)

diff --git a/src/libos/src/vm/vm_area.rs b/src/libos/src/vm/vm_area.rs
index 3209e8a6..d1fc7032 100644
--- a/src/libos/src/vm/vm_area.rs
+++ b/src/libos/src/vm/vm_area.rs
@@ -279,7 +279,15 @@ impl VMArea {
 
         if need_flush {
             let file_offset = file_offset.unwrap() + (target_range.start() - self.range.start());
-            file.unwrap().write_at(file_offset, buf);
+            let file_len = file.unwrap().metadata().unwrap().size;
+            if file_offset < file_len {
+                let effective_mem_len = std::cmp::min(target_range.size(), file_len - file_offset);
+                let len = file
+                    .unwrap()
+                    .write_at(file_offset, &buf[..effective_mem_len])
+                    .expect("flush file failure");
+                debug_assert!(len == effective_mem_len);
+            } // else file_offset >= file_len, no need to write file
         }
 
         // reset zeros
@@ -541,16 +549,40 @@ impl VMArea {
         if !cond_fn(file) {
             return;
         }
+        self.flush_file(file, file_offset)
+            .expect("flush memory to file failure");
+    }
+
+    fn flush_file(&self, file: &Arc<dyn File>, file_offset: usize) -> Result<()> {
+        let file_len = file.metadata().unwrap().size;
+        if file_offset >= file_len {
+            return Ok(());
+        }
+
         if self.is_fully_committed() {
-            file.write_at(file_offset, unsafe { self.as_slice() });
+            let effective_mem_len = std::cmp::min(self.range().size(), file_len - file_offset);
+            let len = file.write_at(file_offset, unsafe {
+                &self.as_slice()[..effective_mem_len]
+            })?;
+            debug_assert!(len == effective_mem_len);
         } else {
             let committed = true;
             let vm_range_start = self.range().start();
             for range in self.pages().get_ranges(committed) {
                 let file_offset = file_offset + (range.start() - vm_range_start);
-                file.write_at(file_offset, unsafe { range.as_slice() });
+                if file_offset >= file_len {
+                    break;
+                }
+
+                let effective_mem_len = std::cmp::min(range.size(), file_len - file_offset);
+                let len = file.write_at(file_offset, unsafe {
+                    &range.as_slice()[..effective_mem_len]
+                })?;
+                debug_assert!(len == effective_mem_len);
             }
         }
+
+        Ok(())
     }
 
     pub fn is_shared(&self) -> bool {
diff --git a/test/mmap/main.c b/test/mmap/main.c
index 40bca7eb..f51fe697 100644
--- a/test/mmap/main.c
+++ b/test/mmap/main.c
@@ -1536,6 +1536,78 @@ int test_kernel_space_pf_trigger() {
     return 0;
 }
 
+int test_shared_file_mmap_small_file() {
+    int fd;
+    struct stat sb1, sb2;
+    char *mapped;
+    char write_buf[] = "hello world\n";
+    size_t page_sz;
+    size_t file_sz;
+    size_t new_file_sz;
+
+    if ((fd = open("/root/a.txt", O_RDWR | O_CREAT, S_IRWXU)) < 0) {
+        perror("open");
+        return -1;
+    }
+
+    if (write(fd, write_buf, strlen(write_buf)) < 0) {
+        perror("write");
+        return -1;
+    }
+    if ((fstat(fd, &sb1)) == -1 ) {
+        perror("fstat");
+        return -1;
+    }
+
+    file_sz = sb1.st_size;
+    page_sz = getpagesize();
+
+    // Output the size before mmap
+    printf("before msync file_sz = %ld\n", file_sz);
+    if ((mapped = mmap(NULL, page_sz, PROT_READ | PROT_WRITE, MAP_SHARED, fd,
+                       0)) == (void *) -1) {
+        perror("mmap");
+        return -1;
+    }
+
+    // Appending characters at the end of a file.
+    mapped[file_sz] = '9';
+    // Synchronizing the modified contents of a file
+    if ((msync((void *)mapped, page_sz, MS_SYNC)) == -1) {
+        perror("msync");
+        return -1;
+    }
+
+    if (fstat(fd, &sb2) == -1) {
+        perror("fstat");
+        return -1;
+    }
+
+    new_file_sz = sb2.st_size;
+    // Output the size after mmap
+    printf("msync after file_sz = %ld\n", new_file_sz);
+    if (new_file_sz != file_sz) {
+        THROW_ERROR("The file size doesn't match");
+    }
+
+    if ((munmap((void *)mapped, page_sz)) == -1) {
+        perror("munmap");
+        return -1;
+    }
+
+    if ((fstat(fd, &sb1)) == -1 ) {
+        perror("fstat");
+        return -1;
+    }
+
+    file_sz = sb1.st_size;
+    if (new_file_sz != file_sz) {
+        THROW_ERROR("The file size doesn't match");
+    }
+
+    return 0;
+}
+
 // ============================================================================
 // Test suite main
 // ============================================================================
@@ -1555,6 +1627,7 @@ static test_case_t test_cases[] = {
     TEST_CASE(test_shared_file_mmap_flushing_with_munmap),
     TEST_CASE(test_shared_file_mmap_flushing_with_fdatasync),
     TEST_CASE(test_shared_file_mmap_flushing_with_fsync),
+    TEST_CASE(test_shared_file_mmap_small_file),
     TEST_CASE(test_fixed_mmap_that_does_not_override_any_mmaping),
     TEST_CASE(test_fixed_mmap_that_overrides_existing_mmaping),
     TEST_CASE(test_fixed_mmap_with_non_page_aligned_addr),