From 0bdbc677dcdb166bb6aa6dae47242ffeffe4e673 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Sat, 16 Nov 2024 23:50:08 +0800
Subject: [PATCH] 1

---
 io/io.go      | 37 ++++++++++++++++++++++++++++++++++---
 io/io_test.go |  2 +-
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/io/io.go b/io/io.go
index 5edbf9e..8b80863 100644
--- 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 864001d..e3c02e1 100644
--- 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)
 	}
 }
-- 
2.39.2