well-goknown/vendor/github.com/valyala/fasthttp/peripconn.go

131 lines
2.2 KiB
Go
Raw Permalink Normal View History

package fasthttp
import (
2024-09-26 01:59:44 +00:00
"crypto/tls"
"net"
"sync"
)
type perIPConnCounter struct {
2024-09-26 01:59:44 +00:00
perIPConnPool sync.Pool
perIPTLSConnPool sync.Pool
m map[uint32]int
lock sync.Mutex
}
func (cc *perIPConnCounter) Register(ip uint32) int {
cc.lock.Lock()
if cc.m == nil {
cc.m = make(map[uint32]int)
}
n := cc.m[ip] + 1
cc.m[ip] = n
cc.lock.Unlock()
return n
}
func (cc *perIPConnCounter) Unregister(ip uint32) {
cc.lock.Lock()
defer cc.lock.Unlock()
if cc.m == nil {
// developer safeguard
panic("BUG: perIPConnCounter.Register() wasn't called")
}
n := cc.m[ip] - 1
if n < 0 {
n = 0
}
cc.m[ip] = n
}
type perIPConn struct {
net.Conn
perIPConnCounter *perIPConnCounter
2024-09-26 01:59:44 +00:00
ip uint32
}
type perIPTLSConn struct {
*tls.Conn
perIPConnCounter *perIPConnCounter
ip uint32
}
2024-09-26 01:59:44 +00:00
func acquirePerIPConn(conn net.Conn, ip uint32, counter *perIPConnCounter) net.Conn {
if tlsConn, ok := conn.(*tls.Conn); ok {
v := counter.perIPTLSConnPool.Get()
if v == nil {
return &perIPTLSConn{
perIPConnCounter: counter,
Conn: tlsConn,
ip: ip,
}
}
c := v.(*perIPTLSConn)
c.Conn = tlsConn
c.ip = ip
return c
}
v := counter.perIPConnPool.Get()
if v == nil {
return &perIPConn{
perIPConnCounter: counter,
Conn: conn,
ip: ip,
}
}
c := v.(*perIPConn)
c.Conn = conn
c.ip = ip
return c
}
2024-09-26 01:59:44 +00:00
func (c *perIPConn) Close() error {
err := c.Conn.Close()
c.perIPConnCounter.Unregister(c.ip)
c.Conn = nil
2024-09-26 01:59:44 +00:00
c.perIPConnCounter.perIPConnPool.Put(c)
return err
}
2024-09-26 01:59:44 +00:00
func (c *perIPTLSConn) Close() error {
err := c.Conn.Close()
c.perIPConnCounter.Unregister(c.ip)
2024-09-26 01:59:44 +00:00
c.Conn = nil
c.perIPConnCounter.perIPTLSConnPool.Put(c)
return err
}
func getUint32IP(c net.Conn) uint32 {
return ip2uint32(getConnIP4(c))
}
func getConnIP4(c net.Conn) net.IP {
addr := c.RemoteAddr()
ipAddr, ok := addr.(*net.TCPAddr)
if !ok {
return net.IPv4zero
}
return ipAddr.IP.To4()
}
func ip2uint32(ip net.IP) uint32 {
if len(ip) != 4 {
return 0
}
return uint32(ip[0])<<24 | uint32(ip[1])<<16 | uint32(ip[2])<<8 | uint32(ip[3])
}
func uint322ip(ip uint32) net.IP {
b := make([]byte, 4)
b[0] = byte(ip >> 24)
b[1] = byte(ip >> 16)
b[2] = byte(ip >> 8)
b[3] = byte(ip)
return b
}