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 +}