From 479a1258c761d952fa6381c5ecaa2dac1e6bfa36 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Wed, 16 Aug 2023 00:14:26 +0800
Subject: [PATCH] add

---
 file/FileWR.go | 22 ++++++++++++++++++----
 web/Web.go     | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/file/FileWR.go b/file/FileWR.go
index a08d630..89e18cd 100644
--- a/file/FileWR.go
+++ b/file/FileWR.go
@@ -77,7 +77,7 @@ func (t *File) CopyTo(to *File, byteInSec int64, tryLock bool) error {
 	}
 	defer to.l.Unlock()
 
-	return transferIO(t.read(), to.write(), byteInSec)
+	return transferIO(t.read(), to.write(), byteInSec, -1)
 }
 
 func (t *File) CopyToIoWriter(to io.Writer, byteInSec int64, tryLock bool) error {
@@ -91,7 +91,21 @@ func (t *File) CopyToIoWriter(to io.Writer, byteInSec int64, tryLock bool) error
 	}
 	defer t.l.RUnlock()
 
-	return transferIO(t.read(), to, byteInSec)
+	return transferIO(t.read(), to, byteInSec, -1)
+}
+
+func (t *File) CopyToIoWriterUntil(to io.Writer, byteInSec, totalSec int64, tryLock bool) error {
+	t.getRWCloser()
+	if t.Config.AutoClose {
+		defer t.Close()
+	}
+
+	if !t.l.TryRLock() {
+		return ErrFailToLock
+	}
+	defer t.l.RUnlock()
+
+	return transferIO(t.read(), to, byteInSec, totalSec)
 }
 
 func (t *File) Write(data []byte, tryLock bool) (int, error) {
@@ -414,12 +428,12 @@ func newPath(path string, mode fs.FileMode) {
 	}
 }
 
-func transferIO(r io.Reader, w io.Writer, byteInSec int64) (e error) {
+func transferIO(r io.Reader, w io.Writer, byteInSec, totalSec int64) (e error) {
 	if byteInSec > 0 {
 		ticker := time.NewTicker(time.Second)
 		defer ticker.Stop()
 
-		for buf := make([]byte, byteInSec); true; {
+		for buf := make([]byte, byteInSec); totalSec < 0 || totalSec > 0; totalSec -= 1 {
 			if n, err := r.Read(buf); n != 0 {
 				if _, werr := w.Write(buf[:n]); werr != nil {
 					return err
diff --git a/web/Web.go b/web/Web.go
index 4513e64..79aced2 100644
--- a/web/Web.go
+++ b/web/Web.go
@@ -348,6 +348,39 @@ func (t *Cache) Cache(key string, aliveDur time.Duration, w http.ResponseWriter)
 	return res
 }
 
+type withflush struct {
+	f  func()
+	h  func() http.Header
+	w  func([]byte) (int, error)
+	wh func(int)
+}
+
+func (t withflush) Header() http.Header {
+	return t.h()
+}
+func (t withflush) Write(b []byte) (i int, e error) {
+	i, e = t.w(b)
+	if t.f != nil {
+		t.f()
+	}
+	return
+}
+func (t withflush) WriteHeader(i int) {
+	t.wh(i)
+}
+
+func WithFlush(w http.ResponseWriter) http.ResponseWriter {
+	if Flusher, ok := w.(http.Flusher); ok {
+		return withflush{
+			f:  Flusher.Flush,
+			h:  w.Header,
+			w:  w.Write,
+			wh: w.WriteHeader,
+		}
+	}
+	return w
+}
+
 func WithStatusCode(w http.ResponseWriter, code int) {
 	w.WriteHeader(code)
 	_, _ = w.Write([]byte(http.StatusText(code)))
-- 
2.39.2