1

diff --git a/io/io.go b/io/io.go
index 5edbf9ebe02f4b7d0aa1901f6d54b1c26a9af56a..8b808638fa6b44e06be37f692fbb10f5d39909cb 100644 (file)
--- a/io/io.go
+++ b/io/io.go
@@ -288,6 +288,7 @@ func WithCtxCopy(ctx context.Context, callTree string, copybuf []byte, to time.D
 
 type CopyConfig struct {
        BytePerLoop, MaxLoop, MaxByte, BytePerSec uint64
+       SkipByte                                  int
 }
 
 // close by yourself
@@ -383,12 +384,32 @@ func Copy(r io.Reader, w io.Writer, c CopyConfig) (e error) {
                        }
                }
        }
+
+       if seeker, ok := r.(io.Seeker); ok && c.SkipByte > 0 {
+               _, e = seeker.Seek(int64(c.SkipByte), io.SeekCurrent)
+               if e != nil {
+                       return
+               }
+               c.SkipByte = 0
+       }
+
        buf := make([]byte, c.BytePerLoop)
        readC := uint64(0)
        for {
                if n, err := r.Read(buf); n != 0 {
-                       if _, werr := w.Write(buf[:n]); werr != nil {
-                               return err
+                       if c.SkipByte > 0 {
+                               if n <= int(c.SkipByte) {
+                                       c.SkipByte -= n
+                                       return
+                               } else {
+                                       n, e = w.Write(buf[c.SkipByte:])
+                                       c.SkipByte = 0
+                               }
+                       } else {
+                               n, e = w.Write(buf[:n])
+                       }
+                       if e != nil {
+                               return
                        }
                        if c.BytePerSec != 0 {
                                readC += uint64(n)
@@ -518,7 +539,17 @@ func WriterWithConfig(w io.Writer, c CopyConfig) (wc io.Writer) {
                                time.Sleep(time.Second)
                                readC = 0
                        }
-                       n, e = w.Write(p)
+                       if c.SkipByte > 0 {
+                               if len(p) <= int(c.SkipByte) {
+                                       c.SkipByte -= len(p)
+                                       return
+                               } else {
+                                       n, e = w.Write(p[c.SkipByte:])
+                                       c.SkipByte = 0
+                               }
+                       } else {
+                               n, e = w.Write(p)
+                       }
                        if c.BytePerSec != 0 {
                                readC += uint64(n)
                        }

diff --git a/io/io_test.go b/io/io_test.go
index 864001d6eb7a7ec274bbcd22891c88c035e8e870..e3c02e1fee155b92e9447de06f13a88376dba05a 100644 (file)
--- a/io/io_test.go
+++ b/io/io_test.go
@@ -14,7 +14,7 @@ func Test_CopyIO(t *testing.T) {
 
        var w = &bytes.Buffer{}
 
-       if e := Copy(bytes.NewReader(s), w, CopyConfig{1<<17 + 1, 1, 0, 0}); e != nil || w.Len() != 1<<17+1 || w.Bytes()[1<<17-1] != '1' || w.Bytes()[1<<17] != '2' {
+       if e := Copy(bytes.NewReader(s), w, CopyConfig{1<<17 + 1, 1, 0, 0, 0}); e != nil || w.Len() != 1<<17+1 || w.Bytes()[1<<17-1] != '1' || w.Bytes()[1<<17] != '2' {
                t.Fatal(e)
        }
 }