conn.go

package gws

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/binary"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/lxzan/gws/internal"
)

// Conn &#32467;&#26500;&#20307;&#34920;&#31034;&#19968;&#20010; WebSocket &#36830;&#25509;
// Conn struct represents a WebSocket connection
type Conn struct {
	// &#20114;&#26021;&#38145;&#65292;&#29992;&#20110;&#20445;&#25252;&#20849;&#20139;&#36164;&#28304;
	// Mutex to protect shared resources
	mu sync.Mutex

	// &#20250;&#35805;&#23384;&#20648;&#65292;&#29992;&#20110;&#23384;&#20648;&#20250;&#35805;&#25968;&#25454;
	// Session storage for storing session data
	ss SessionStorage

	// &#29992;&#20110;&#23384;&#20648;&#38169;&#35823;&#30340;&#21407;&#23376;&#20540;
	// Atomic value for storing errors
	err atomic.Value

	// &#26631;&#35782;&#26159;&#21542;&#20026;&#26381;&#21153;&#22120;&#31471;
	// Indicates if this is a server-side connection
	isServer bool

	// &#23376;&#21327;&#35758;
	// Subprotocol
	subprotocol string

	// &#24213;&#23618;&#32593;&#32476;&#36830;&#25509;
	// Underlying network connection
	conn net.Conn

	// &#37197;&#32622;&#20449;&#24687;
	// Configuration information
	config *Config

	// &#32531;&#20914;&#35835;&#21462;&#22120;
	// Buffered reader
	br *bufio.Reader

	// &#25345;&#32493;&#24103;
	// Continuation frame
	continuationFrame continuationFrame

	// &#24103;&#22836;
	// Frame header
	fh frameHeader

	// &#20107;&#20214;&#22788;&#29702;&#22120;
	// Event handler
	handler Event

	// &#20851;&#38381;&#29366;&#24577;
	// Closed state
	closed uint32

	// &#35835;&#21462;&#38431;&#21015;
	// Read queue
	readQueue channel

	// &#20889;&#20837;&#38431;&#21015;
	// Write queue
	writeQueue workerQueue

	// &#21387;&#32553;&#22120;
	// Deflater
	deflater *deflater

	// &#25968;&#25454;&#21253;&#21457;&#36865;&#31383;&#21475;
	// Data packet send window
	dpsWindow slideWindow

	// &#25968;&#25454;&#21253;&#25509;&#25910;&#31383;&#21475;
	// Data packet receive window
	cpsWindow slideWindow

	// &#27599;&#28040;&#24687;&#21387;&#32553;
	// Per-message deflate
	pd PermessageDeflate
}

// ReadLoop
// &#24490;&#29615;&#35835;&#21462;&#28040;&#24687;. &#22914;&#26524;&#22797;&#29992;&#20102;HTTP Server, &#24314;&#35758;&#24320;&#21551;goroutine, &#38459;&#22622;&#20250;&#23548;&#33268;&#35831;&#27714;&#19978;&#19979;&#25991;&#26080;&#27861;&#34987;GC.
// Read messages in a loop.
// If HTTP Server is reused, it is recommended to enable goroutine, as blocking will prevent the context from being GC.
func (c *Conn) ReadLoop() {
	c.handler.OnOpen(c)

	// &#26080;&#38480;&#24490;&#29615;&#35835;&#21462;&#28040;&#24687;, &#22914;&#26524;&#21457;&#29983;&#38169;&#35823;&#21017;&#35302;&#21457;&#38169;&#35823;&#20107;&#20214;&#24182;&#36864;&#20986;&#24490;&#29615;
	// Infinite loop to read messages, if an error occurs, trigger the error event and exit the loop
	for {
		if err := c.readMessage(); err != nil {
			c.emitError(err)
			break
		}
	}

	err, ok := c.err.Load().(error)
	c.handler.OnClose(c, internal.SelectValue(ok, err, errEmpty))

	// &#22238;&#25910;&#36164;&#28304;
	// Reclaim resources
	if c.isServer {
		c.br.Reset(nil)
		c.config.brPool.Put(c.br)
		c.br = nil
		if c.cpsWindow.enabled {
			c.config.cswPool.Put(c.cpsWindow.dict)
			c.cpsWindow.dict = nil
		}
		if c.dpsWindow.enabled {
			c.config.dswPool.Put(c.dpsWindow.dict)
			c.dpsWindow.dict = nil
		}
	}
}

// &#33719;&#21462;&#21387;&#32553;&#23383;&#20856;
// Get compressed dictionary
func (c *Conn) getCpsDict(isBroadcast bool) []byte {
	// &#24191;&#25773;&#27169;&#24335;&#24517;&#39035;&#20445;&#35777;&#27599;&#19968;&#24103;&#37117;&#26159;&#30456;&#21516;&#30340;&#20869;&#23481;, &#25152;&#20197;&#19981;&#20351;&#29992;&#19978;&#19979;&#25991;&#25509;&#31649;&#20248;&#21270;&#21387;&#32553;&#29575;
	// In broadcast mode, each frame must be the same content, so context takeover is not used to optimize compression ratio
	if isBroadcast {
		return nil
	}

	// &#22914;&#26524;&#26159;&#26381;&#21153;&#22120;&#24182;&#19988;&#26381;&#21153;&#22120;&#19978;&#19979;&#25991;&#25509;&#31649;&#21551;&#29992;&#65292;&#36820;&#22238;&#21387;&#32553;&#23383;&#20856;
	// If it is a server and server context takeover is enabled, return the compression dictionary
	if c.isServer && c.pd.ServerContextTakeover {
		return c.cpsWindow.dict
	}

	// &#22914;&#26524;&#26159;&#23458;&#25143;&#31471;&#24182;&#19988;&#23458;&#25143;&#31471;&#19978;&#19979;&#25991;&#25509;&#31649;&#21551;&#29992;&#65292;&#36820;&#22238;&#21387;&#32553;&#23383;&#20856;
	// If client-side and client context takeover is enabled, return the compression dictionary
	if !c.isServer && c.pd.ClientContextTakeover {
		return c.cpsWindow.dict
	}

	return nil
}

// &#33719;&#21462;&#35299;&#21387;&#23383;&#20856;
// Get decompression dictionary
func (c *Conn) getDpsDict() []byte {
	// &#22914;&#26524;&#26159;&#26381;&#21153;&#22120;&#24182;&#19988;&#23458;&#25143;&#31471;&#19978;&#19979;&#25991;&#25509;&#31649;&#21551;&#29992;&#65292;&#36820;&#22238;&#35299;&#21387;&#23383;&#20856;
	// If it is a server and client context takeover is enabled, return the decompression dictionary
	if c.isServer && c.pd.ClientContextTakeover {
		return c.dpsWindow.dict
	}

	// &#22914;&#26524;&#26159;&#23458;&#25143;&#31471;&#24182;&#19988;&#26381;&#21153;&#22120;&#19978;&#19979;&#25991;&#25509;&#31649;&#21551;&#29992;&#65292;&#36820;&#22238;&#35299;&#21387;&#23383;&#20856;
	// If it is a client and server context takeover is enabled, return the decompressed dictionary
	if !c.isServer && c.pd.ServerContextTakeover {
		return c.dpsWindow.dict
	}

	return nil
}

// UTF8&#32534;&#30721;&#26816;&#26597;
// UTF8 encoding check
func (c *Conn) isTextValid(opcode Opcode, payload []byte) bool {
	if c.config.CheckUtf8Enabled {
		return internal.CheckEncoding(uint8(opcode), payload)
	}
	return true
}

// &#26816;&#26597;&#36830;&#25509;&#26159;&#21542;&#24050;&#20851;&#38381;
// Checks if the connection is closed
func (c *Conn) isClosed() bool {
	return atomic.LoadUint32(&c.closed) == 1
}

// &#20851;&#38381;&#36830;&#25509;&#24182;&#23384;&#20648;&#38169;&#35823;&#20449;&#24687;
// Closes the connection and stores the error information
func (c *Conn) close(reason []byte, err error) {
	c.err.Store(err)
	_ = c.doWrite(OpcodeCloseConnection, internal.Bytes(reason))
	_ = c.conn.Close()
}

// &#22788;&#29702;&#38169;&#35823;&#20107;&#20214;
// Handle the error event
func (c *Conn) emitError(err error) {
	if err == nil {
		return
	}

	// &#20351;&#29992;&#21407;&#23376;&#25805;&#20316;&#26816;&#26597;&#24182;&#35774;&#32622;&#36830;&#25509;&#30340;&#20851;&#38381;&#29366;&#24577;
	// Use atomic operation to check and set the closed state of the connection
	if atomic.CompareAndSwapUint32(&c.closed, 0, 1) {
		var responseCode = internal.CloseNormalClosure
		var responseErr error = internal.CloseNormalClosure

		// &#26681;&#25454;&#38169;&#35823;&#31867;&#22411;&#35774;&#32622;&#21709;&#24212;&#20195;&#30721;&#21644;&#21709;&#24212;&#38169;&#35823;
		// Set response code and response error based on the error type
		switch v := err.(type) {
		case internal.StatusCode:
			responseCode = v
		case *internal.Error:
			responseCode = v.Code
			responseErr = v.Err
		default:
			responseErr = err
		}

		var content = responseCode.Bytes()
		content = append(content, err.Error()...)

		// &#22914;&#26524;&#20869;&#23481;&#38271;&#24230;&#36229;&#36807;&#38408;&#20540;&#65292;&#25130;&#26029;&#20869;&#23481;
		// If the content length exceeds the threshold, truncate the content
		if len(content) > internal.ThresholdV1 {
			content = content[:internal.ThresholdV1]
		}

		c.close(content, responseErr)
	}
}

// &#22788;&#29702;&#20851;&#38381;&#20107;&#20214;
// Handles the close event
func (c *Conn) emitClose(buf *bytes.Buffer) error {
	var responseCode = internal.CloseNormalClosure
	var realCode = internal.CloseNormalClosure.Uint16()
	switch buf.Len() {
	case 0:
		responseCode = 0
		realCode = 0
	case 1:
		responseCode = internal.CloseProtocolError
		realCode = uint16(buf.Bytes()[0])
		buf.Reset()
	default:
		var b [2]byte
		_, _ = buf.Read(b[0:])
		realCode = binary.BigEndian.Uint16(b[0:])
		switch realCode {
		case 1004, 1005, 1006, 1014, 1015:
			responseCode = internal.CloseProtocolError
		default:
			if realCode < 1000 || realCode >= 5000 || (realCode >= 1016 && realCode < 3000) {
				responseCode = internal.CloseProtocolError
			} else if realCode < 1016 {
				responseCode = internal.CloseNormalClosure
			} else {
				responseCode = internal.StatusCode(realCode)
			}
		}
		if !c.isTextValid(OpcodeCloseConnection, buf.Bytes()) {
			responseCode = internal.CloseUnsupportedData
		}
	}
	if atomic.CompareAndSwapUint32(&c.closed, 0, 1) {
		c.close(responseCode.Bytes(), &CloseError{Code: realCode, Reason: buf.Bytes()})
	}
	return internal.CloseNormalClosure
}

// SetDeadline &#35774;&#32622;&#36830;&#25509;&#30340;&#25130;&#27490;&#26102;&#38388;
// Sets the deadline for the connection
func (c *Conn) SetDeadline(t time.Time) error {
	err := c.conn.SetDeadline(t)
	c.emitError(err)
	return err
}

// SetReadDeadline &#35774;&#32622;&#35835;&#21462;&#25805;&#20316;&#30340;&#25130;&#27490;&#26102;&#38388;
// Sets the deadline for read operations
func (c *Conn) SetReadDeadline(t time.Time) error {
	err := c.conn.SetReadDeadline(t)
	c.emitError(err)
	return err
}

// SetWriteDeadline &#35774;&#32622;&#20889;&#20837;&#25805;&#20316;&#30340;&#25130;&#27490;&#26102;&#38388;
// Sets the deadline for write operations
func (c *Conn) SetWriteDeadline(t time.Time) error {
	err := c.conn.SetWriteDeadline(t)
	c.emitError(err)
	return err
}

// LocalAddr &#36820;&#22238;&#26412;&#22320;&#32593;&#32476;&#22320;&#22336;
// Returns the local network address
func (c *Conn) LocalAddr() net.Addr {
	return c.conn.LocalAddr()
}

// RemoteAddr &#36820;&#22238;&#36828;&#31243;&#32593;&#32476;&#22320;&#22336;
// Returns the remote network address
func (c *Conn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// NetConn
// &#33719;&#21462;&#24213;&#23618;&#30340; TCP/TLS/KCP &#31561;&#36830;&#25509;
// Gets the underlying TCP/TLS/KCP... connection
func (c *Conn) NetConn() net.Conn {
	return c.conn
}

// SetNoDelay
// &#25511;&#21046;&#25805;&#20316;&#31995;&#32479;&#26159;&#21542;&#24212;&#35813;&#24310;&#36831;&#25968;&#25454;&#21253;&#20256;&#36755;&#20197;&#26399;&#26395;&#21457;&#36865;&#26356;&#23569;&#30340;&#25968;&#25454;&#21253;&#65288;Nagle &#31639;&#27861;&#65289;&#12290;
// &#40664;&#35748;&#20540;&#26159; true&#65288;&#26080;&#24310;&#36831;&#65289;&#65292;&#36825;&#24847;&#21619;&#30528;&#25968;&#25454;&#22312; Write &#20043;&#21518;&#23613;&#24555;&#21457;&#36865;&#12290;
// Controls whether the operating system should delay packet transmission in hopes of sending fewer packets (Nagle's algorithm).
// The default is true (no delay), meaning that data is sent as soon as possible after a Write.
func (c *Conn) SetNoDelay(noDelay bool) error {
	switch v := c.conn.(type) {
	case *net.TCPConn:
		return v.SetNoDelay(noDelay)

	case *tls.Conn:
		if netConn, ok := v.NetConn().(*net.TCPConn); ok {
			return netConn.SetNoDelay(noDelay)
		}
	}
	return nil
}

// SubProtocol &#33719;&#21462;&#21327;&#21830;&#30340;&#23376;&#21327;&#35758;
// Gets the negotiated sub-protocol
func (c *Conn) SubProtocol() string { return c.subprotocol }

// Session &#33719;&#21462;&#20250;&#35805;&#23384;&#20648;
// Gets the session storage
func (c *Conn) Session() SessionStorage { return c.ss }