Repository: otiai10/copy
Pull request #91 (2022-10-25 05:36:13)

Imsunv/main
Aknowledgements: @imsunv


Patch:
diff --git a/.gitignore b/.gitignore
index 3b65243..a793485 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,6 @@ coverage.txt
 vendor
 .vagrant
 .idea/
+
+# Test Specific
+test/data/case16/large.file
diff --git a/copy.go b/copy.go
index 77eaf3c..4bdb397 100644
--- a/copy.go
+++ b/copy.go
@@ -83,7 +83,12 @@ func fcopy(src, dest string, info os.FileInfo, opt Options) (err error) {
 
 	var buf []byte = nil
 	var w io.Writer = f
-	// var r io.Reader = s
+	var r io.Reader = s
+
+	if opt.WrapReader != nil {
+		r = opt.WrapReader(s)
+	}
+
 	if opt.CopyBufferSize != 0 {
 		buf = make([]byte, opt.CopyBufferSize)
 		// Disable using `ReadFrom` by io.CopyBuffer.
@@ -91,7 +96,8 @@ func fcopy(src, dest string, info os.FileInfo, opt Options) (err error) {
 		w = struct{ io.Writer }{f}
 		// r = struct{ io.Reader }{s}
 	}
-	if _, err = io.CopyBuffer(w, s, buf); err != nil {
+
+	if _, err = io.CopyBuffer(w, r, buf); err != nil {
 		return err
 	}
 
diff --git a/options.go b/options.go
index 0045813..da08377 100644
--- a/options.go
+++ b/options.go
@@ -1,6 +1,9 @@
 package copy
 
-import "os"
+import (
+	"io"
+	"os"
+)
 
 // Options specifies optional actions on copying.
 type Options struct {
@@ -46,6 +49,11 @@ type Options struct {
 	// See https://golang.org/pkg/io/#CopyBuffer for more information.
 	CopyBufferSize uint
 
+	// If you want to add some limitation on reading src file,
+	// you can wrap the src and provide new reader,
+	// such as `RateLimitReader` in the test case.
+	WrapReader func(src *os.File) io.Reader
+
 	intent struct {
 		src  string
 		dest string
@@ -92,6 +100,7 @@ func getDefaultOptions(src, dest string) Options {
 		Sync:              false,              // Do not sync
 		PreserveTimes:     false,              // Do not preserve the modification time
 		CopyBufferSize:    0,                  // Do not specify, use default bufsize (32*1024)
+		WrapReader:        nil,                // Do not wrap src files, use them as they are.
 		intent: struct {
 			src  string
 			dest string


Test patch:
diff --git a/all_test.go b/all_test.go
index b77a182..5584711 100644
--- a/all_test.go
+++ b/all_test.go
@@ -2,11 +2,13 @@ package copy
 
 import (
 	"errors"
+	"io"
 	"io/ioutil"
 	"os"
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
 	. "github.com/otiai10/mint"
 )
@@ -326,3 +328,45 @@ func TestOptions_PreserveOwner(t *testing.T) {
 	err := Copy("test/data/case13", "test/data.copy/case13", opt)
 	Expect(t, err).ToBe(nil)
 }
+
+func TestOptions_CopyRateLimit(t *testing.T) {
+
+	file, err := os.Create("test/data/case16/large.file")
+	if err != nil {
+		t.Errorf("failed to create test file: %v", err)
+		return
+	}
+
+	size := int64(100 * 1024) // 100 KB
+	if err := file.Truncate(size); err != nil {
+		t.Errorf("failed to truncate test file: %v", err)
+		t.SkipNow()
+		return
+	}
+
+	opt := Options{WrapReader: func(src *os.File) io.Reader {
+		return &SleepyReader{src, 1}
+	}}
+
+	start := time.Now()
+	err = Copy("test/data/case16", "test/data.copy/case16", opt)
+	elasped := time.Since(start)
+	Expect(t, err).ToBe(nil)
+	Expect(t, elasped > 5*time.Second).ToBe(true)
+}
+
+type SleepyReader struct {
+	src *os.File
+	sec time.Duration
+}
+
+func (r *SleepyReader) Read(p []byte) (int, error) {
+	n, e := r.src.Read(p)
+	if e != nil && e != io.EOF {
+		return n, e
+	}
+	if n > 0 {
+		time.Sleep(time.Second * r.sec)
+	}
+	return n, e
+}
diff --git a/test/data/case16/README.md b/test/data/case16/README.md
new file mode 100644
index 0000000..29064a8
--- /dev/null
+++ b/test/data/case16/README.md
@@ -0,0 +1,5 @@
+So if you wanted to limit copy rate with a KB/second value, you could add options like this:
+```go
+opt.CopyRateLimit = 100
+```
+The default value is 0, which means no limit.
\ No newline at end of file
diff --git a/test_setup.go b/test_setup.go
index 0db62fa..f3c2d83 100644
--- a/test_setup.go
+++ b/test_setup.go
@@ -10,6 +10,7 @@ import (
 )
 
 func setup(m *testing.M) {
+	os.RemoveAll("test/data.copy")
 	os.MkdirAll("test/data.copy", os.ModePerm)
 	os.Symlink("test/data/case01", "test/data/case03/case01")
 	os.Chmod("test/data/case07/dir_0555", 0o555)