From f6d5be04a723eaaca90986ef13c19d2926f345f2 Mon Sep 17 00:00:00 2001
From: Sydonian <794346190@qq.com>
Date: Mon, 5 Aug 2024 09:21:45 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=B7=A5=E5=85=B7=E5=87=BD?=
 =?UTF-8?q?=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 utils/io2/io.go     | 47 +++++++++++++++++++++++++++
 utils/io2/range.go  | 79 +++++++++++++++++++++++++++++++++++++++++++++
 utils/math2/math.go | 16 +++++++++
 3 files changed, 142 insertions(+)
 create mode 100644 utils/io2/range.go

diff --git a/utils/io2/io.go b/utils/io2/io.go
index 89a95e8..4228e38 100644
--- a/utils/io2/io.go
+++ b/utils/io2/io.go
@@ -90,6 +90,41 @@ func AfterReadClosedOnce(closer io.ReadCloser, callback func(closer io.ReadClose
 	}
 }
 
+type afterEOF struct {
+	inner    io.ReadCloser
+	callback func(str io.ReadCloser, err error)
+}
+
+func (hook *afterEOF) Read(buf []byte) (n int, err error) {
+	n, err = hook.inner.Read(buf)
+	if hook.callback != nil {
+		if err == io.EOF {
+			hook.callback(hook.inner, nil)
+			hook.callback = nil
+		} else if err != nil {
+			hook.callback(hook.inner, err)
+			hook.callback = nil
+		}
+	}
+	return n, err
+}
+
+func (hook *afterEOF) Close() error {
+	err := hook.inner.Close()
+	if hook.callback != nil {
+		hook.callback(hook.inner, io.ErrClosedPipe)
+		hook.callback = nil
+	}
+	return err
+}
+
+func AfterEOF(str io.ReadCloser, callback func(str io.ReadCloser, err error)) io.ReadCloser {
+	return &afterEOF{
+		inner:    str,
+		callback: callback,
+	}
+}
+
 type readerWithCloser struct {
 	reader io.Reader
 	closer func(reader io.Reader) error
@@ -152,3 +187,15 @@ func ToReaders(strs []io.ReadCloser) ([]io.Reader, func()) {
 		}
 	}
 }
+
+func DropWithBuf(str io.Reader, buf []byte) error {
+	for {
+		_, err := str.Read(buf)
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+	}
+}
diff --git a/utils/io2/range.go b/utils/io2/range.go
new file mode 100644
index 0000000..a4a0b35
--- /dev/null
+++ b/utils/io2/range.go
@@ -0,0 +1,79 @@
+package io2
+
+import (
+	"io"
+
+	"gitlink.org.cn/cloudream/common/utils/math2"
+)
+
+type rng struct {
+	offset int64
+	length *int64
+	inner  io.ReadCloser
+	err    error
+}
+
+func (r *rng) Read(p []byte) (n int, err error) {
+	if r.err != nil {
+		return 0, r.err
+	}
+
+	if r.offset > 0 {
+		buf := make([]byte, 1024*16)
+		for r.offset > 0 {
+			need := math2.Min(r.offset, int64(len(buf)))
+			rd, err := r.inner.Read(buf[:need])
+			if err != nil {
+				r.err = err
+				return 0, err
+			}
+			r.offset -= int64(rd)
+		}
+	}
+
+	if r.length == nil {
+		return r.inner.Read(p)
+	}
+
+	need := math2.Min(*r.length, int64(len(p)))
+	rd, err := r.inner.Read(p[:need])
+	if err != nil {
+		r.err = err
+		return rd, io.EOF
+	}
+
+	*r.length -= int64(rd)
+	if *r.length == 0 {
+		r.err = io.EOF
+	}
+
+	return rd, nil
+}
+
+func (r *rng) Close() error {
+	r.err = io.ErrClosedPipe
+	return r.inner.Close()
+}
+
+func NewRange(inner io.ReadCloser, offset int64, length *int64) io.ReadCloser {
+	return &rng{
+		offset: offset,
+		length: length,
+		inner:  inner,
+	}
+}
+
+func Ranged(inner io.ReadCloser, offset int64, length int64) io.ReadCloser {
+	return &rng{
+		offset: offset,
+		length: &length,
+		inner:  inner,
+	}
+}
+
+func Offset(inner io.ReadCloser, offset int64) io.ReadCloser {
+	return &rng{
+		offset: offset,
+		inner:  inner,
+	}
+}
diff --git a/utils/math2/math.go b/utils/math2/math.go
index 2ee7c4f..6677d2b 100644
--- a/utils/math2/math.go
+++ b/utils/math2/math.go
@@ -17,3 +17,19 @@ func Min[T constraints.Ordered](v1, v2 T) T {
 
 	return v2
 }
+
+func Ceil[T constraints.Integer](v T, div T) T {
+	return (v + div - 1) / div * div
+}
+
+func Floor[T constraints.Integer](v T, div T) T {
+	return v / div * div
+}
+
+func CeilDiv[T constraints.Integer](v T, div T) T {
+	return (v + div - 1) / div
+}
+
+func FloorDiv[T constraints.Integer](v T, div T) T {
+	return v / div
+}