move util dict to parent and add buf method

config/const.go

@@ -21,6 +21,8 @@ const (
 const (
 	// ID 号最高值，高于这个值的ID号为系统内部使用
 	ID_MAX = 65500
-	// 验证ID
-	ID_AUTH = 65502
+
+	ID_AUTH          = 65502 // 验证ID
+	ID_PROXY_CONNECT = 65503 // 代理连接请求
+	ID_PROXY_RESULT  = 65504 // 代理请求执行结果
 )

conn/tcp2/tcp2.go

@@ -13,15 +13,17 @@ import (
 
 	"git.me9.top/git/tinymq/config"
 	"git.me9.top/git/tinymq/conn"
-	"git.me9.top/git/tinymq/conn/util"
+	"git.me9.top/git/tinymq/util"
 )
 
 const PROTO string = "tcp"
 const VERSION uint8 = 2
 
-// 数据包的最大长度
+// 第一级数据包的最大长度，如果超过则自动增加一个地址位到4个字节
 const MAX_LENGTH = 0xFFFF
-const MAX2_LENGTH = 0x1FFFFFFF // 500 M，避免申请过大内存
+
+// 由于都是内部使用，感觉没有必要加这个限制
+// const MAX2_LENGTH = 0x1FFFFFFF // 500 M，避免申请过大内存
 
 type Tcp2 struct {
 	cf       *config.Config
@@ -197,10 +199,7 @@ func Dial(cf *config.Config, addr string, hash string) (conn.Connect, error) {
 
 // 发送数据到网络
 // 如果有加密函数的话会直接修改源数据
-func (c *Tcp2) writeMessage(buf []byte) (err error) {
-	if len(buf) > MAX2_LENGTH {
-		return fmt.Errorf("data length more than %d", MAX2_LENGTH)
-	}
+func (c *Tcp2) WriteRaw(buf []byte) (err error) {
 	if c.cipher != nil {
 		c.cipher.Encrypt(buf, buf)
 	}
@@ -276,11 +275,11 @@ func (c *Tcp2) WriteAuthInfo(channel string, auth []byte) (err error) {
 	index += channelLen
 	copy(buf[index:], auth)
 	buf[start+dlen] = util.CRC8(buf[start : start+dlen])
-	return c.writeMessage(buf)
+	return c.WriteRaw(buf)
 }
 
 // 从连接中读取信息
-func (c *Tcp2) readMessage(deadline int) ([]byte, error) {
+func (c *Tcp2) ReadRaw(deadline int) ([]byte, error) {
 	buf := make([]byte, 2)
 	err := c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
 	if err != nil {
@@ -310,7 +309,7 @@ func (c *Tcp2) readMessage(deadline int) ([]byte, error) {
 			c.cipher.Decrypt(buf, buf)
 		}
 		dlen = binary.BigEndian.Uint32(buf)
-		if dlen < MAX_LENGTH || dlen > MAX2_LENGTH {
+		if dlen < MAX_LENGTH {
 			return nil, errors.New("wrong length in read message")
 		}
 	}
@@ -339,7 +338,7 @@ func (c *Tcp2) ReadAuthInfo() (proto string, version uint8, channel string, auth
 			return
 		}
 	}()
-	msg, err := c.readMessage(c.cf.ReadWait)
+	msg, err := c.ReadRaw(c.cf.ReadWait)
 	if err != nil {
 		return
 	}
@@ -425,7 +424,7 @@ func (c *Tcp2) WriteRequest(id uint16, cmd string, data []byte) error {
 
 		copy(buf[index:], compressedData)
 		buf[start+ddlen] = util.CRC8(buf[start : start+ddlen])
-		return c.writeMessage(buf)
+		return c.WriteRaw(buf)
 	} else {
 		ddlen := 2 + 1 + cmdLen + dlen
 		buf, start := c.writeDataLen(ddlen)
@@ -438,7 +437,7 @@ func (c *Tcp2) WriteRequest(id uint16, cmd string, data []byte) error {
 		index += cmdLen
 		copy(buf[index:], data)
 		buf[start+ddlen] = util.CRC8(buf[start : start+ddlen])
-		return c.writeMessage(buf)
+		return c.WriteRaw(buf)
 	}
 }
 
@@ -469,7 +468,7 @@ func (c *Tcp2) WriteResponse(id uint16, state uint8, data []byte) error {
 
 		copy(buf[index:], compressedData)
 		buf[start+ddlen] = util.CRC8(buf[start : start+ddlen])
-		return c.writeMessage(buf)
+		return c.WriteRaw(buf)
 	} else {
 		ddlen := 2 + 1 + len(data)
 		buf, start := c.writeDataLen(ddlen)
@@ -480,7 +479,7 @@ func (c *Tcp2) WriteResponse(id uint16, state uint8, data []byte) error {
 		index++
 		copy(buf[index:], data)
 		buf[start+ddlen] = util.CRC8(buf[start : start+ddlen])
-		return c.writeMessage(buf)
+		return c.WriteRaw(buf)
 	}
 }
 
@@ -491,12 +490,12 @@ func (c *Tcp2) WritePing(id uint16) error {
 	index := start
 	binary.BigEndian.PutUint16(buf[index:index+2], id)
 	buf[start+dlen] = util.CRC8(buf[start : start+dlen])
-	return c.writeMessage(buf)
+	return c.WriteRaw(buf)
 }
 
 // 获取信息
 func (c *Tcp2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd string, state uint8, data []byte, err error) {
-	msg, err := c.readMessage(deadline)
+	msg, err := c.ReadRaw(deadline)
 	if err != nil {
 		return
 	}
@@ -539,12 +538,12 @@ func (c *Tcp2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd s
 
 // 获取远程的地址
 func (c *Tcp2) RemoteIP() net.IP {
-	return conn.AddrToIP(c.conn.RemoteAddr())
+	return util.AddrToIP(c.conn.RemoteAddr())
 }
 
 // 获取本地的地址
 func (c *Tcp2) LocalIP() net.IP {
-	return conn.AddrToIP(c.conn.LocalAddr())
+	return util.AddrToIP(c.conn.LocalAddr())
 }
 
 func (c *Tcp2) Close() error {

conn/type.go

@@ -12,6 +12,8 @@ const (
 	PingMsg MsgType = iota
 	RequestMsg
 	ResponseMsg
+	ProxyConnectMsg
+	ProxyResultMsg
 )
 
 func (t MsgType) String() string {
@@ -22,6 +24,10 @@ func (t MsgType) String() string {
 		return "RequestMsg"
 	case ResponseMsg:
 		return "ResponseMsg"
+	case ProxyConnectMsg:
+		return "ProxyConnectMsg"
+	case ProxyResultMsg:
+		return "ProxyResultMsg"
 	default:
 		return fmt.Sprintf("Unknown MsgType (%d)", t)
 	}
@@ -43,18 +49,3 @@ type Connect interface {
 
 // 服务请求连接
 type ServerConnectFunc func(conn Connect)
-
-// addrToIP attempts to extract a net.IP from a net.Addr
-func AddrToIP(addr net.Addr) net.IP {
-	switch v := addr.(type) {
-	case *net.TCPAddr:
-		return v.IP
-	case *net.UDPAddr:
-		return v.IP
-	case *net.IPAddr:
-		return v.IP
-	default:
-		// Handle other potential net.Addr implementations or return nil
-		return nil
-	}
-}

conn/ws2/ws2.go

@@ -14,7 +14,7 @@ import (
 
 	"git.me9.top/git/tinymq/config"
 	"git.me9.top/git/tinymq/conn"
-	"git.me9.top/git/tinymq/conn/util"
+	"git.me9.top/git/tinymq/util"
 	"github.com/gorilla/websocket"
 )
 
@@ -467,12 +467,12 @@ func (c *Ws2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd st
 
 // 获取远程的地址
 func (c *Ws2) RemoteIP() net.IP {
-	return conn.AddrToIP(c.conn.RemoteAddr())
+	return util.AddrToIP(c.conn.RemoteAddr())
 }
 
 // 获取本地的地址
 func (c *Ws2) LocalIP() net.IP {
-	return conn.AddrToIP(c.conn.LocalAddr())
+	return util.AddrToIP(c.conn.LocalAddr())
 }
 
 func (c *Ws2) Close() error {

type.go

@@ -49,6 +49,8 @@ type ConnectState byte
 const (
 	Disconnected ConnectState = iota
 	Connected
+	Proxing // 正在进行Proxy请求
+	Proxied // Proxy 连接成功状态
 	Closed
 )
 
@@ -58,6 +60,10 @@ func (t ConnectState) String() string {
 		return "Disconnected"
 	case Connected:
 		return "Connected"
+	case Proxing:
+		return "Proxing"
+	case Proxied:
+		return "Proxied"
 	case Closed:
 		return "Closed"
 	default:

util/buf/alloc.go

@@ -0,0 +1,132 @@
+package buf
+
+// Inspired by https://github.com/xtaci/smux/blob/master/alloc.go
+
+import (
+	"errors"
+	"math/bits"
+	"sync"
+)
+
+var DefaultAllocator = newDefaultAllocator()
+
+type Allocator interface {
+	Get(size int) []byte
+	Put(buf []byte) error
+}
+
+// defaultAllocator for incoming frames, optimized to prevent overwriting after zeroing
+type defaultAllocator struct {
+	buffers [11]sync.Pool
+}
+
+// NewAllocator initiates a []byte allocator for frames less than 65536 bytes,
+// the waste(memory fragmentation) of space allocation is guaranteed to be
+// no more than 50%.
+func newDefaultAllocator() Allocator {
+	return &defaultAllocator{
+		buffers: [...]sync.Pool{ // 64B -> 64K
+			{New: func() any { return new([1 << 6]byte) }},
+			{New: func() any { return new([1 << 7]byte) }},
+			{New: func() any { return new([1 << 8]byte) }},
+			{New: func() any { return new([1 << 9]byte) }},
+			{New: func() any { return new([1 << 10]byte) }},
+			{New: func() any { return new([1 << 11]byte) }},
+			{New: func() any { return new([1 << 12]byte) }},
+			{New: func() any { return new([1 << 13]byte) }},
+			{New: func() any { return new([1 << 14]byte) }},
+			{New: func() any { return new([1 << 15]byte) }},
+			{New: func() any { return new([1 << 16]byte) }},
+		},
+	}
+}
+
+// Get a []byte from pool with most appropriate cap
+func (alloc *defaultAllocator) Get(size int) []byte {
+	if size <= 0 || size > 65536 {
+		return nil
+	}
+
+	var index uint16
+	if size > 64 {
+		index = msb(size)
+		if size != 1<<index {
+			index += 1
+		}
+		index -= 6
+	}
+
+	buffer := alloc.buffers[index].Get()
+	switch index {
+	case 0:
+		return buffer.(*[1 << 6]byte)[:size]
+	case 1:
+		return buffer.(*[1 << 7]byte)[:size]
+	case 2:
+		return buffer.(*[1 << 8]byte)[:size]
+	case 3:
+		return buffer.(*[1 << 9]byte)[:size]
+	case 4:
+		return buffer.(*[1 << 10]byte)[:size]
+	case 5:
+		return buffer.(*[1 << 11]byte)[:size]
+	case 6:
+		return buffer.(*[1 << 12]byte)[:size]
+	case 7:
+		return buffer.(*[1 << 13]byte)[:size]
+	case 8:
+		return buffer.(*[1 << 14]byte)[:size]
+	case 9:
+		return buffer.(*[1 << 15]byte)[:size]
+	case 10:
+		return buffer.(*[1 << 16]byte)[:size]
+	default:
+		panic("invalid pool index")
+	}
+}
+
+// Put returns a []byte to pool for future use,
+// which the cap must be exactly 2^n
+func (alloc *defaultAllocator) Put(buf []byte) error {
+	bits := msb(cap(buf))
+	if cap(buf) < 64 || cap(buf) > 65536 || cap(buf) != 1<<bits {
+		return errors.New("allocator Put() incorrect buffer size")
+	}
+	bits -= 6
+	buf = buf[:cap(buf)]
+
+	//nolint
+	//lint:ignore SA6002 ignore temporarily
+	switch bits {
+	case 0:
+		alloc.buffers[bits].Put((*[1 << 6]byte)(buf))
+	case 1:
+		alloc.buffers[bits].Put((*[1 << 7]byte)(buf))
+	case 2:
+		alloc.buffers[bits].Put((*[1 << 8]byte)(buf))
+	case 3:
+		alloc.buffers[bits].Put((*[1 << 9]byte)(buf))
+	case 4:
+		alloc.buffers[bits].Put((*[1 << 10]byte)(buf))
+	case 5:
+		alloc.buffers[bits].Put((*[1 << 11]byte)(buf))
+	case 6:
+		alloc.buffers[bits].Put((*[1 << 12]byte)(buf))
+	case 7:
+		alloc.buffers[bits].Put((*[1 << 13]byte)(buf))
+	case 8:
+		alloc.buffers[bits].Put((*[1 << 14]byte)(buf))
+	case 9:
+		alloc.buffers[bits].Put((*[1 << 15]byte)(buf))
+	case 10:
+		alloc.buffers[bits].Put((*[1 << 16]byte)(buf))
+	default:
+		panic("invalid pool index")
+	}
+	return nil
+}
+
+// msb return the pos of most significant bit
+func msb(size int) uint16 {
+	return uint16(bits.Len32(uint32(size)) - 1)
+}

util/buf/pool.go

@@ -0,0 +1,12 @@
+package buf
+
+func Get(size int) []byte {
+	if size == 0 {
+		return nil
+	}
+	return DefaultAllocator.Get(size)
+}
+
+func Put(buf []byte) error {
+	return DefaultAllocator.Put(buf)
+}

util/util.go

@@ -0,0 +1,18 @@
+package util
+
+import "net"
+
+// addrToIP attempts to extract a net.IP from a net.Addr
+func AddrToIP(addr net.Addr) net.IP {
+	switch v := addr.(type) {
+	case *net.TCPAddr:
+		return v.IP
+	case *net.UDPAddr:
+		return v.IP
+	case *net.IPAddr:
+		return v.IP
+	default:
+		// Handle other potential net.Addr implementations or return nil
+		return nil
+	}
+}