monibuca/pkg/util/buf_reader.go

package util

import (
	"io"
	"net"
	"net/textproto"
	"strings"
	"time"
)

const defaultBufSize = 1 << 14

type BufReader struct {
	Allocator *ScalableMemoryAllocator
	buf       MemoryReader
	totalRead int
	BufLen    int
	Mouth     chan []byte
	feedData  func() error
}

func NewBufReaderWithBufLen(reader io.Reader, bufLen int) (r *BufReader) {
	r = &BufReader{
		Allocator: NewScalableMemoryAllocator(bufLen),
		BufLen:    bufLen,
		feedData: func() error {
			buf, err := r.Allocator.Read(reader, r.BufLen)
			if err != nil {
				return err
			}
			n := len(buf)
			r.totalRead += n
			r.buf.Buffers = append(r.buf.Buffers, buf)
			r.buf.Size += n
			r.buf.Length += n
			return nil
		},
	}
	r.buf.Memory = &Memory{}
	//fmt.Println("NewBufReaderWithBufLen", uintptr(unsafe.Pointer(r.allocator)))
	return
}

// NewBufReaderWithTimeout 创建一个具有指定读取超时时间的 BufReader
func NewBufReaderWithTimeout(conn net.Conn, timeout time.Duration) (r *BufReader) {
	r = &BufReader{
		Allocator: NewScalableMemoryAllocator(defaultBufSize),
		BufLen:    defaultBufSize,
		feedData: func() error {
			// 设置读取超时
			if conn != nil && timeout > 0 {
				if err := conn.SetReadDeadline(time.Now().Add(timeout)); err != nil {
					return err
				}
			}
			buf, err := r.Allocator.Read(conn, r.BufLen)
			if err != nil {
				return err
			}
			n := len(buf)
			r.totalRead += n
			r.buf.Buffers = append(r.buf.Buffers, buf)
			r.buf.Size += n
			r.buf.Length += n
			return nil
		},
	}
	r.buf.Memory = &Memory{}
	return
}

func NewBufReaderBuffersChan(feedChan chan net.Buffers) (r *BufReader) {
	r = &BufReader{
		feedData: func() error {
			data, ok := <-feedChan
			if !ok {
				return io.EOF
			}
			for _, buf := range data {
				n := len(buf)
				r.totalRead += n
				r.buf.Size += n
				r.buf.Length += n
			}
			r.buf.Buffers = append(r.buf.Buffers, data...)
			return nil
		},
	}
	r.buf.Memory = &Memory{}
	return
}

func NewBufReaderChan(bufferSize int) (r *BufReader) {
	feedChan := make(chan []byte, bufferSize)
	r = &BufReader{
		Mouth: feedChan,
		feedData: func() error {
			data, ok := <-feedChan
			if !ok {
				return io.EOF
			}
			n := len(data)
			r.totalRead += n
			r.buf.Buffers = append(r.buf.Buffers, data)
			r.buf.Size += n
			r.buf.Length += n
			return nil
		},
	}
	r.buf.Memory = &Memory{}
	return
}

func (r *BufReader) Feed(data []byte) bool {
	select {
	case r.Mouth <- data:
		return true
	default:
		return false
	}
}

func NewBufReader(reader io.Reader) (r *BufReader) {
	return NewBufReaderWithBufLen(reader, defaultBufSize)
}

func (r *BufReader) Recycle() {
	r.buf = MemoryReader{}
	if r.Allocator != nil {
		r.Allocator.Recycle()
	}
	if r.Mouth != nil {
		close(r.Mouth)
	}
}

func (r *BufReader) Buffered() int {
	return r.buf.Length
}

func (r *BufReader) Peek(n int) (buf []byte, err error) {
	defer func(snap MemoryReader) {
		l := r.buf.Length + n
		r.buf = snap
		r.buf.Length = l
	}(
		r.buf)
	for range n {
		if b, err := r.ReadByte(); err != nil {
			return nil, err
		} else {
			buf = append(buf, b)
		}
	}
	return
}

func (r *BufReader) ReadByte() (b byte, err error) {
	for r.buf.Length == 0 {
		if err = r.feedData(); err != nil {
			return
		}
	}
	return r.buf.ReadByte()
}

func (r *BufReader) ReadBE(n int) (num int, err error) {
	for i := range n {
		b, err := r.ReadByte()
		if err != nil {
			return 0, err
		}
		num += int(b) << ((n - i - 1) << 3)
	}
	return
}

func (r *BufReader) ReadLE32(n int) (num uint32, err error) {
	for i := range n {
		b, err := r.ReadByte()
		if err != nil {
			return 0, err
		}
		num += uint32(b) << (i << 3)
	}
	return
}

func (r *BufReader) ReadBE32(n int) (num uint32, err error) {
	for i := range n {
		b, err := r.ReadByte()
		if err != nil {
			return 0, err
		}
		num += uint32(b) << ((n - i - 1) << 3)
	}
	return
}

func (r *BufReader) Skip(n int) (err error) {
	return r.ReadRange(n, nil)
}

func (r *BufReader) ReadRange(n int, yield func([]byte)) (err error) {
	for r.recycleFront(); n > 0 && err == nil; err = r.feedData() {
		if r.buf.Length > 0 {
			if r.buf.Length >= n {
				r.buf.RangeN(n, yield)
				return
			}
			n -= r.buf.Length
			r.buf.Range(yield)
		}
	}
	return
}

func (r *BufReader) Read(to []byte) (n int, err error) {
	n = len(to)
	err = r.ReadNto(n, to)
	return
}

func (r *BufReader) ReadNto(n int, to []byte) (err error) {
	l := 0
	return r.ReadRange(n, func(buf []byte) {
		ll := len(buf)
		copy(to[l:l+ll], buf)
		l += ll
	})
}

func (r *BufReader) ReadString(n int) (s string, err error) {
	err = r.ReadRange(n, func(buf []byte) {
		s += string(buf)
	})
	return
}

func (r *BufReader) ReadBytes(n int) (mem Memory, err error) {
	err = r.ReadRange(n, mem.PushOne)
	return
}

func (r *BufReader) recycleFront() {
	r.buf.ClipFront(r.Allocator.Free)
}

func (r *BufReader) ReadLine() (line string, err error) {
	var lastb, curb byte
	snap, i := r.buf, 0
	for {
		if curb, err = r.ReadByte(); err != nil {
			return "", err
		} else {
			i++
			if l := r.buf.Length; curb == '\n' {
				snap.Length = l + i
				r.buf = snap
				err = r.ReadRange(i, func(buf []byte) {
					line = line + string(buf)
				})
				if lastb == '\r' {
					line = line[:i-2]
				} else {
					line = line[:i-1]
				}
				return
			}
			lastb = curb
		}
	}
}

func (r *BufReader) ReadMIMEHeader() (textproto.MIMEHeader, error) {
	result := make(textproto.MIMEHeader)
	for {
		l, err := r.ReadLine()
		if err != nil {
			return nil, err
		}
		if l == "" {
			break
		}
		key, value, _ := strings.Cut(l, ":")
		result.Add(key, strings.Trim(value, " "))
	}
	return result, nil
}