135 lines
2.8 KiB
Go
135 lines
2.8 KiB
Go
package goproxy
|
|
|
|
import (
|
|
"bufio"
|
|
"io"
|
|
"net"
|
|
"time"
|
|
)
|
|
|
|
// ConnBuffer 连接缓冲区
|
|
// 封装了底层网络连接和缓冲读取器，提供了更方便的读写接口
|
|
type ConnBuffer struct {
|
|
// 底层连接
|
|
conn net.Conn
|
|
// 缓冲读取器
|
|
reader *bufio.Reader
|
|
}
|
|
|
|
// NewConnBuffer 创建连接缓冲区
|
|
func NewConnBuffer(conn net.Conn, reader *bufio.Reader) *ConnBuffer {
|
|
if reader == nil {
|
|
reader = bufio.NewReader(conn)
|
|
}
|
|
return &ConnBuffer{
|
|
conn: conn,
|
|
reader: reader,
|
|
}
|
|
}
|
|
|
|
// Read 从连接读取数据
|
|
func (c *ConnBuffer) Read(b []byte) (int, error) {
|
|
return c.reader.Read(b)
|
|
}
|
|
|
|
// Write 向连接写入数据
|
|
func (c *ConnBuffer) Write(b []byte) (int, error) {
|
|
return c.conn.Write(b)
|
|
}
|
|
|
|
// Close 关闭连接
|
|
func (c *ConnBuffer) Close() error {
|
|
return c.conn.Close()
|
|
}
|
|
|
|
// LocalAddr 获取本地地址
|
|
func (c *ConnBuffer) LocalAddr() net.Addr {
|
|
return c.conn.LocalAddr()
|
|
}
|
|
|
|
// RemoteAddr 获取远程地址
|
|
func (c *ConnBuffer) RemoteAddr() net.Addr {
|
|
return c.conn.RemoteAddr()
|
|
}
|
|
|
|
// SetDeadline 设置读写超时
|
|
func (c *ConnBuffer) SetDeadline(t time.Time) error {
|
|
return c.conn.SetDeadline(t)
|
|
}
|
|
|
|
// SetReadDeadline 设置读取超时
|
|
func (c *ConnBuffer) SetReadDeadline(t time.Time) error {
|
|
return c.conn.SetReadDeadline(t)
|
|
}
|
|
|
|
// SetWriteDeadline 设置写入超时
|
|
func (c *ConnBuffer) SetWriteDeadline(t time.Time) error {
|
|
return c.conn.SetWriteDeadline(t)
|
|
}
|
|
|
|
// BufferReader 获取缓冲读取器
|
|
func (c *ConnBuffer) BufferReader() *bufio.Reader {
|
|
return c.reader
|
|
}
|
|
|
|
// Peek 查看缓冲区中的数据，但不消费
|
|
func (c *ConnBuffer) Peek(n int) ([]byte, error) {
|
|
return c.reader.Peek(n)
|
|
}
|
|
|
|
// ReadByte 读取一个字节
|
|
func (c *ConnBuffer) ReadByte() (byte, error) {
|
|
return c.reader.ReadByte()
|
|
}
|
|
|
|
// UnreadByte 将最后读取的字节放回缓冲区
|
|
func (c *ConnBuffer) UnreadByte() error {
|
|
return c.reader.UnreadByte()
|
|
}
|
|
|
|
// ReadLine 读取一行数据
|
|
func (c *ConnBuffer) ReadLine() (string, error) {
|
|
line, isPrefix, err := c.reader.ReadLine()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// 如果一行数据没有读取完整，继续读取
|
|
if isPrefix {
|
|
var buf []byte
|
|
buf = append(buf, line...)
|
|
for isPrefix && err == nil {
|
|
line, isPrefix, err = c.reader.ReadLine()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
buf = append(buf, line...)
|
|
}
|
|
return string(buf), nil
|
|
}
|
|
|
|
return string(line), nil
|
|
}
|
|
|
|
// ReadN 读取指定字节数的数据
|
|
func (c *ConnBuffer) ReadN(n int) ([]byte, error) {
|
|
buf := make([]byte, n)
|
|
_, err := io.ReadFull(c.reader, buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return buf, nil
|
|
}
|
|
|
|
// Flush 刷新缓冲区
|
|
func (c *ConnBuffer) Flush() error {
|
|
// 由于我们只有读取缓冲区，没有写入缓冲区，所以这里不需要实际操作
|
|
return nil
|
|
}
|
|
|
|
// Reset 重置连接缓冲区
|
|
func (c *ConnBuffer) Reset(conn net.Conn) {
|
|
c.conn = conn
|
|
c.reader = bufio.NewReader(conn)
|
|
}
|