naza/pkg/connection/connection.go

//package connection 对 net.Conn 接口的二次封装
package connection

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"time"
)

var connectionErr = errors.New("connection: fxxk")

type Connection interface {
	// 包含 interface net.Conn 的所有方法
	net.Conn

	// 额外提供的读方法
	ReadAtLeast(buf []byte, min int) (n int, err error)
	ReadLine() (line []byte, isPrefix bool, err error)

	// 额外提供的写方法
	Printf(fmt string, v ...interface{}) (n int, err error)

	// 带超时的读/写方法
	ReadWithTimeout(b []byte, timeoutMS int) (n int, err error)
	ReadAtLeastWithTimeout(buf []byte, min int, timeoutMS int) (n int, err error)
	ReadLineWithTimeout(timeoutMS int) (line []byte, isPrefix bool, err error)
	WriteWithTimeout(b []byte, timeoutMS int) (n int, err error)
	PrintfWithTimeout(timeoutMS int, fmt string, v ...interface{}) (n int, err error)
}

// 可配置是直接从 net.Conn 上读/写数据，还是中间添加一层buffer缓冲
type Config struct {
	ReadBufSize  int
	WriteBufSize int
}

func New(conn net.Conn, config *Config) Connection {
	var c connection
	c.Conn = conn
	if config != nil {
		if config.ReadBufSize > 0 {
			c.r = bufio.NewReaderSize(conn, config.ReadBufSize)
		}
		if config.WriteBufSize > 0 {
			c.w = bufio.NewWriterSize(conn, config.WriteBufSize)
		}
		c.config = *config
	}
	if c.r == nil {
		c.r = conn
	}
	if c.w == nil {
		c.w = conn
	}

	return &c
}

type connection struct {
	Conn   net.Conn
	r      io.Reader
	w      io.Writer
	config Config
}

func (c *connection) ReadAtLeastWithTimeout(buf []byte, min int, timeoutMS int) (n int, err error) {
	if timeoutMS > 0 {
		c.Conn.SetReadDeadline(time.Now().Add(time.Duration(timeoutMS) * time.Millisecond))
	}
	return io.ReadAtLeast(c.r, buf, min)
}

func (c *connection) ReadLineWithTimeout(timeoutMS int) (line []byte, isPrefix bool, err error) {
	if timeoutMS > 0 {
		c.Conn.SetReadDeadline(time.Now().Add(time.Duration(timeoutMS) * time.Millisecond))
	}
	return c.ReadLine()
}

func (c *connection) ReadWithTimeout(b []byte, timeoutMS int) (n int, err error) {
	if timeoutMS > 0 {
		c.Conn.SetReadDeadline(time.Now().Add(time.Duration(timeoutMS) * time.Millisecond))
	}
	return c.Conn.Read(b)
}

func (c *connection) WriteWithTimeout(b []byte, timeoutMS int) (n int, err error) {
	if timeoutMS > 0 {
		c.Conn.SetWriteDeadline(time.Now().Add(time.Duration(timeoutMS) * time.Millisecond))
	}
	return c.Conn.Write(b)
}

func (c *connection) PrintfWithTimeout(timeoutMS int, format string, v ...interface{}) (n int, err error) {
	if timeoutMS > 0 {
		c.Conn.SetWriteDeadline(time.Now().Add(time.Duration(timeoutMS) * time.Millisecond))
	}
	return fmt.Fprintf(c.Conn, format, v...)
}

func (c *connection) ReadAtLeast(buf []byte, min int) (n int, err error) {
	return io.ReadAtLeast(c.r, buf, min)
}

func (c *connection) ReadLine() (line []byte, isPrefix bool, err error) {
	bufioReader, ok := c.r.(*bufio.Reader)
	if !ok {
		// 目前只有使用了 bufio.Reader 时才能执行 ReadLine 操作
		return nil, false, connectionErr
	}
	return bufioReader.ReadLine()
}

func (c *connection) Printf(format string, v ...interface{}) (n int, err error) {
	return fmt.Fprintf(c.Conn, format, v...)
}

func (c *connection) Read(b []byte) (n int, err error) {
	return c.r.Read(b)
}

func (c *connection) Write(b []byte) (n int, err error) {
	return c.w.Write(b)
}

func (c *connection) Close() error {
	return c.Conn.Close()
}

func (c *connection) LocalAddr() net.Addr {
	return c.Conn.LocalAddr()
}

func (c *connection) RemoteAddr() net.Addr {
	return c.Conn.RemoteAddr()
}

func (c *connection) SetDeadline(t time.Time) error {
	return c.Conn.SetDeadline(t)
}

func (c *connection) SetReadDeadline(t time.Time) error {
	return c.Conn.SetReadDeadline(t)
}
func (c *connection) SetWriteDeadline(t time.Time) error {
	return c.Conn.SetWriteDeadline(t)
}