Fix copy_file_range usage for files > 2GB

cpuguy83 · cpuguy83 · commit afba265aa60a · 2018-05-17T12:41:36.000-04:00
copy_file_range (or any Linux syscall) will only ever copy 2GB at a
time. In this case there is no error returned from the system call.

This fix uses copy_file_range in a loop until either 0 bytes are copied,
or the desired amount has been copied (st.Size()).

Signed-off-by: Brian Goff &lt;cpuguy83@gmail.com&gt;
diff --git a/fs/copy_linux.go b/fs/copy_linux.go
@@ -49,20 +49,26 @@ func copyFileContent(dst, src *os.File) error {
 		return errors.Wrap(err, "unable to stat source")
 	}
 
-	n, err := unix.CopyFileRange(int(src.Fd()), nil, int(dst.Fd()), nil, int(st.Size()), 0)
-	if err != nil {
-		if err != unix.ENOSYS && err != unix.EXDEV {
-			return errors.Wrap(err, "copy file range failed")
-		}
+	size := st.Size()
+	first := true
+	srcFd := int(src.Fd())
+	dstFd := int(dst.Fd())
 
-		buf := bufferPool.Get().(*[]byte)
-		_, err = io.CopyBuffer(dst, src, *buf)
-		bufferPool.Put(buf)
-		return err
-	}
+	for size > 0 {
+		n, err := unix.CopyFileRange(srcFd, nil, dstFd, nil, int(size), 0)
+		if err != nil {
+			if (err != unix.ENOSYS && err != unix.EXDEV) || !first {
+				return errors.Wrap(err, "copy file range failed")
+			}
+
+			buf := bufferPool.Get().(*[]byte)
+			_, err = io.CopyBuffer(dst, src, *buf)
+			bufferPool.Put(buf)
+			return errors.Wrap(err, "userspace copy failed")
+		}
 
-	if int64(n) != st.Size() {
-		return errors.Wrapf(err, "short copy: %d of %d", int64(n), st.Size())
+		first = false
+		size -= int64(n)
 	}
 
 	return nil
diff --git a/fs/copy_test.go b/fs/copy_test.go
@@ -1,12 +1,12 @@
 package fs
 
 import (
+	_ "crypto/sha256"
+	"io"
 	"io/ioutil"
 	"os"
 	"testing"
 
-	_ "crypto/sha256"
-
 	"github.com/containerd/continuity/fs/fstest"
 	"github.com/pkg/errors"
 )
@@ -46,6 +46,45 @@ func TestCopyDirectoryWithLocalSymlink(t *testing.T) {
 	}
 }
 
+// TestCopyWithLargeFile tests copying a file whose size > 2^32 bytes.
+func TestCopyWithLargeFile(t *testing.T) {
+	dataR, dataW := io.Pipe()
+	var written int64
+	max := int64(3 * 1024 * 1024 * 1024)
+	defer dataR.Close()
+
+	go func() {
+		var (
+			data = make([]byte, 64*1024)
+			err  error
+			n    int
+		)
+		defer func() {
+			dataW.CloseWithError(err)
+		}()
+
+		for written < max {
+			n, err = dataW.Write(data)
+			if err != nil {
+				return
+			}
+			written += int64(n)
+		}
+	}()
+
+	apply := fstest.Apply(
+		fstest.CreateDir("/banana", 0755),
+		fstest.WriteFileStream("/banana/split", dataR, 0644),
+	)
+
+	if err := testCopy(apply); err != nil {
+		t.Fatal(err)
+	}
+	if written < max {
+		t.Fatalf("wrote fewer bytes than expected: %d", written)
+	}
+}
+
 func testCopy(apply fstest.Applier) error {
 	t1, err := ioutil.TempDir("", "test-copy-src-")
 	if err != nil {
diff --git a/fs/fstest/file.go b/fs/fstest/file.go
@@ -1,7 +1,8 @@
 package fstest
 
 import (
-	"io/ioutil"
+	"bytes"
+	"io"
 	"net"
 	"os"
 	"path/filepath"
@@ -22,9 +23,26 @@ func (a applyFn) Apply(root string) error {
 // CreateFile returns a file applier which creates a file as the
 // provided name with the given content and permission.
 func CreateFile(name string, content []byte, perm os.FileMode) Applier {
-	return applyFn(func(root string) error {
+	return WriteFileStream(name, bytes.NewReader(content), perm)
+}
+
+// WriteFileStream returns a file applier which creates a file as the
+// provided name with the given content from the provided i/o stream and permission.
+func WriteFileStream(name string, stream io.Reader, perm os.FileMode) Applier {
+	return applyFn(func(root string) (retErr error) {
 		fullPath := filepath.Join(root, name)
-		if err := ioutil.WriteFile(fullPath, content, perm); err != nil {
+		f, err := os.OpenFile(fullPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
+		if err != nil {
+			return err
+		}
+		defer func() {
+			err := f.Close()
+			if err != nil && retErr == nil {
+				retErr = err
+			}
+		}()
+		_, err = io.Copy(f, stream)
+		if err != nil {
 			return err
 		}
 		return os.Chmod(fullPath, perm)
