zsocket/conn.go at main · cyper98/zsocket

677 lines (603 loc) · 16.3 KB
package zsocket
	"encoding/binary"
	"encoding/json"
	"sync/atomic"
// Conn is a server-side WebSocket connection managed by zsocket.
// It exposes Read/Write methods, JSON helpers, and graceful Close.
// Concurrency: one reader goroutine internally; writes are serialized via writeMu.
// NOTE: This MVP supports basic fragmentation (accumulate until FIN). Extensions are not implemented.
type Conn struct {
	raw      net.Conn
	br       *bufio.Reader
	bw       *bufio.Writer
	cfg      Config
	subproto string
	isClient bool
	readLimit int64
	compressionMu          sync.RWMutex
	compressor             Compressor
	compressionLevel       int
	writeCompressionEnable bool
	deadlineMu    sync.RWMutex
	readDeadline  time.Time
	writeDeadline time.Time
	// pong keeper
	lastPong atomic.Value // time.Time
	closed   atomic.Bool
	handlerMu    sync.RWMutex
	pingHandler  PingHandler
	pongHandler  PongHandler
	closeHandler CloseHandler
	// Ensure only one write at a time.
	writeMu sync.Mutex
	sendOnce sync.Once
	sendQ    chan queuedMessage
	sendQuit chan struct{}
	// Allow readers to stop.
	closeOnce sync.Once
type queuedMessage struct {
	ctx  context.Context
	mt   MessageType
	data []byte
	res  chan error
// PingHandler handles incoming ping control frames.
type PingHandler func(appData string) error
// PongHandler handles incoming pong control frames.
type PongHandler func(appData string) error
// CloseHandler handles incoming close control frames.
type CloseHandler func(code CloseCode, text string) error
// newConn wraps a hijacked net.Conn into a zsocket Conn with buffers and starts keepalive if configured.
func newConn(c net.Conn, cfg Config, subproto string, isClient bool) (*Conn, error) {
	if cfg.ReadBufferSize <= 0 {
		cfg.ReadBufferSize = DefaultConfig().ReadBufferSize
	if cfg.WriteBufferSize <= 0 {
		cfg.WriteBufferSize = DefaultConfig().WriteBufferSize
	if cfg.ReadLimit < 0 {
		cfg.ReadLimit = 0
	z := &Conn{
		raw:       c,
		br:        bufio.NewReaderSize(c, cfg.ReadBufferSize),
		bw:        bufio.NewWriterSize(c, cfg.WriteBufferSize),
		cfg:       cfg,
		subproto:  subproto,
		isClient:  isClient,
		readLimit: cfg.ReadLimit,
		compressionLevel: func() int {
			if cfg.CompressionLevel == 0 {
				return DefaultConfig().CompressionLevel
			return cfg.CompressionLevel
		sendQuit: make(chan struct{}),
	if cfg.EnableCompression {
		d, err := NewDeflate(z.compressionLevel)
		if err != nil {
			return nil, err
		z.compressor = d
		z.writeCompressionEnable = true
	z.lastPong.Store(time.Now())
	z.pingHandler = func(appData string) error {
		return z.writeControlOpcode(opPong, []byte(appData), time.Time{})
	z.pongHandler = func(string) error {
		z.lastPong.Store(time.Now())
		return nil
	z.closeHandler = func(code CloseCode, text string) error {
		return z.writeControlOpcode(opClose, closePayload(code, text), time.Time{})
	// Start keepalive loop if enabled.
	if cfg.PingInterval > 0 && cfg.PongWait > 0 {
		go z.keepAlive()
	return z, nil
// EnableWriteCompression enables or disables write-side compression.
func (c *Conn) EnableWriteCompression(enable bool) {
	c.compressionMu.Lock()
	c.writeCompressionEnable = enable
	c.compressionMu.Unlock()
// SetCompressionLevel changes the compression level for future messages.
func (c *Conn) SetCompressionLevel(level int) error {
	c.compressionMu.Lock()
	defer c.compressionMu.Unlock()
	if !c.cfg.EnableCompression {
		return fmt.Errorf("zsocket: compression not negotiated")
	d, err := NewDeflate(level)
	if err != nil {
		return err
	c.compressor = d
	c.compressionLevel = level
	return nil
// keepAlive periodically sends ping frames and checks for timely pongs.
// If a pong is not received within PongWait, the connection is closed.
func (c *Conn) keepAlive() {
	ticker := time.NewTicker(c.cfg.PingInterval)
	defer ticker.Stop()
		if c.closed.Load() {
		case <-ticker.C:
			_ = c.writeControlOpcode(opPing, nil, time.Time{}) // best-effort; ignore error
			// check pong freshness
			last := c.lastPong.Load().(time.Time)
			if time.Since(last) > c.cfg.PongWait {
				_ = c.Close(CloseGoingAway, "pong timeout")
// Subprotocol returns the negotiated subprotocol, if any.
func (c *Conn) Subprotocol() string { return c.subproto }
// RemoteAddr returns the peer network address.
func (c *Conn) RemoteAddr() net.Addr { return c.raw.RemoteAddr() }
// LocalAddr returns the local network address.
func (c *Conn) LocalAddr() net.Addr { return c.raw.LocalAddr() }
// NetConn returns the underlying net.Conn.
func (c *Conn) NetConn() net.Conn { return c.raw }
// UnderlyingConn is an alias for NetConn.
func (c *Conn) UnderlyingConn() net.Conn { return c.raw }
// SetReadLimit sets the maximum message payload size in bytes. 0 means unlimited.
func (c *Conn) SetReadLimit(limit int64) {
	if limit < 0 {
		limit = 0
	c.readLimit = limit
// SetReadDeadline sets the deadline for future reads.
func (c *Conn) SetReadDeadline(t time.Time) error {
	c.deadlineMu.Lock()
	c.readDeadline = t
	c.deadlineMu.Unlock()
	return c.raw.SetReadDeadline(t)
// SetWriteDeadline sets the deadline for future writes.
func (c *Conn) SetWriteDeadline(t time.Time) error {
	c.deadlineMu.Lock()
	c.writeDeadline = t
	c.deadlineMu.Unlock()
	return c.raw.SetWriteDeadline(t)
// PingHandler returns the current ping handler.
func (c *Conn) PingHandler() PingHandler {
	c.handlerMu.RLock()
	h := c.pingHandler
	c.handlerMu.RUnlock()
// SetPingHandler sets a custom ping handler. Nil resets the default behavior.
func (c *Conn) SetPingHandler(h PingHandler) {
	if h == nil {
		h = func(appData string) error {
			return c.writeControlOpcode(opPong, []byte(appData), time.Time{})
	c.handlerMu.Lock()
	c.pingHandler = h
	c.handlerMu.Unlock()
// PongHandler returns the current pong handler.
func (c *Conn) PongHandler() PongHandler {
	c.handlerMu.RLock()
	h := c.pongHandler
	c.handlerMu.RUnlock()
// SetPongHandler sets a custom pong handler. Nil resets the default behavior.
func (c *Conn) SetPongHandler(h PongHandler) {
	if h == nil {
		h = func(string) error {
			c.lastPong.Store(time.Now())
			return nil
	c.handlerMu.Lock()
	c.pongHandler = h
	c.handlerMu.Unlock()
// CloseHandler returns the current close handler.
func (c *Conn) CloseHandler() CloseHandler {
	c.handlerMu.RLock()
	h := c.closeHandler
	c.handlerMu.RUnlock()
// SetCloseHandler sets a custom close handler. Nil resets the default behavior.
func (c *Conn) SetCloseHandler(h CloseHandler) {
	if h == nil {
		h = func(code CloseCode, text string) error {
			return c.writeControlOpcode(opClose, closePayload(code, text), time.Time{})
	c.handlerMu.Lock()
	c.closeHandler = h
	c.handlerMu.Unlock()
// Close sends a close frame (best-effort) and closes the connection.
func (c *Conn) Close(code CloseCode, reason string) error {
	var err error
	c.closeOnce.Do(func() {
		c.closed.Store(true)
		close(c.sendQuit)
		// Send close control frame (best-effort).
		_ = c.writeControlOpcode(opClose, closePayload(code, reason), time.Time{})
		err = c.raw.Close()
	return err
// EnableSendQueue configures async write queue with backpressure.
func (c *Conn) EnableSendQueue(size int) {
	if size <= 0 {
		size = 64
	c.sendOnce.Do(func() {
		c.sendQ = make(chan queuedMessage, size)
		go c.sendLoop()
// Send enqueues a message if queue enabled, otherwise writes directly.
func (c *Conn) Send(ctx context.Context, mt MessageType, data []byte) error {
	if c.sendQ == nil {
		return c.Write(ctx, mt, data)
	msg := queuedMessage{ctx: ctx, mt: mt, data: append([]byte(nil), data...), res: make(chan error, 1)}
	case <-c.sendQuit:
		return ErrCloseSent
	case <-ctx.Done():
		return ctx.Err()
	case c.sendQ <- msg:
	case <-c.sendQuit:
		return ErrCloseSent
	case <-ctx.Done():
		return ctx.Err()
	case err := <-msg.res:
		return err
func (c *Conn) sendLoop() {
		case <-c.sendQuit:
		case msg := <-c.sendQ:
			msg.res <- c.Write(msg.ctx, msg.mt, msg.data)
// Read reads the next complete message, returning its type and payload.
// It accumulates fragments until FIN=1. Control frames are handled internally.
func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
	// Context deadline/timeout support on the underlying conn.
	if deadline, ok := ctx.Deadline(); ok {
		_ = c.raw.SetReadDeadline(deadline)
		defer c.restoreReadDeadline()
	} else if d := c.getReadDeadline(); !d.IsZero() {
		_ = c.raw.SetReadDeadline(d)
		msgOpcode  byte
		full       bytes.Buffer
		started    bool
		compressed bool
		fin, rsv1, opcode, payload, err := readFrame(c.br, c.readLimit, !c.isClient, c.cfg.EnableCompression)
		if err != nil {
			return 0, nil, err
		switch opcode {
		case opText, opBinary:
			if started {
				return 0, nil, errFragmentState
			if rsv1 && !c.cfg.EnableCompression {
				return 0, nil, errRSVNonZero
			compressed = rsv1
			msgOpcode = opcode
			full.Write(payload)
			started = true
			if fin {
				b := full.Bytes()
				if compressed {
					b, err = c.decompressData(b)
					if err != nil {
						return 0, nil, err
				return toMsgType(msgOpcode), b, nil
		case opContinuation:
			if !started {
				return 0, nil, errUnexpectedOpcode
			if rsv1 {
				return 0, nil, errRSVNonZero
			full.Write(payload)
			if fin {
				b := full.Bytes()
				if compressed {
					b, err = c.decompressData(b)
					if err != nil {
						return 0, nil, err
				return toMsgType(msgOpcode), b, nil
		case opPing:
			if err := c.PingHandler()(string(payload)); err != nil {
				return 0, nil, err
		case opPong:
			if err := c.PongHandler()(string(payload)); err != nil {
				return 0, nil, err
		case opClose:
			code, reason := parseClosePayload(payload)
			if err := c.CloseHandler()(code, reason); err != nil {
				return 0, nil, err
			return 0, nil, CloseError{Code: code, Text: reason, Reason: reason}
			// Protocol error → close
			_ = c.Close(CloseProtocolError, "unsupported opcode")
			return 0, nil, fmt.Errorf("zsocket: unsupported opcode %d", opcode)
// Write sends a complete message (text or binary) as a single unfragmented frame.
func (c *Conn) Write(ctx context.Context, mt MessageType, data []byte) error {
	if c.closed.Load() {
		return ErrCloseSent
	if mt != TextMessage && mt != BinaryMessage {
		return fmt.Errorf("zsocket: invalid message type %d", mt)
	// Apply write deadline if ctx has one or from config.
	if deadline, ok := ctx.Deadline(); ok {
		_ = c.raw.SetWriteDeadline(deadline)
		defer c.restoreWriteDeadline()
	} else if d := c.getWriteDeadline(); !d.IsZero() {
		_ = c.raw.SetWriteDeadline(d)
	} else if c.cfg.WriteTimeout > 0 {
		_ = c.raw.SetWriteDeadline(time.Now().Add(c.cfg.WriteTimeout))
		defer c.restoreWriteDeadline()
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	var opcode byte
	rsv1 := false
	payload := data
	if mt == TextMessage {
		opcode = opText
		opcode = opBinary
	if c.shouldCompressWrite() {
		var err error
		payload, err = c.compressData(data)
		if err != nil {
			return err
		rsv1 = true
	if c.cfg.WriteBufferPool != nil {
		if pv := c.cfg.WriteBufferPool.Get(); pv != nil {
			if p, ok := pv.(*[]byte); ok && p != nil {
				buf := (*p)[:0]
				buf = append(buf, payload...)
				payload = buf
				defer func() {
					*p = buf
					c.cfg.WriteBufferPool.Put(p)
	if err := writeFrame(c.bw, true, opcode, payload, c.isClient, rsv1); err != nil {
		return err
	return c.bw.Flush()
// WritePreparedMessage writes a prebuilt message efficiently.
func (c *Conn) WritePreparedMessage(pm *PreparedMessage) error {
	if pm == nil {
		return fmt.Errorf("zsocket: nil prepared message")
	return c.Write(context.Background(), pm.messageType, pm.data)
// ReadMessage reads next data message.
func (c *Conn) ReadMessage() (MessageType, []byte, error) {
	return c.Read(context.Background())
// WriteMessage writes one data message.
func (c *Conn) WriteMessage(mt MessageType, data []byte) error {
	return c.Write(context.Background(), mt, data)
// WriteJSON marshals v to JSON and writes as a text message.
func (c *Conn) WriteJSON(ctx context.Context, v any) error {
	b, err := json.Marshal(v)
	if err != nil {
		return err
	return c.Write(ctx, TextMessage, b)
// ReadJSON reads next message and unmarshals JSON into v.
// It accepts both text and binary frames. If binary, it still tries to decode JSON.
func (c *Conn) ReadJSON(ctx context.Context, v any) error {
	mt, b, err := c.Read(ctx)
	if err != nil {
		return err
	if mt != TextMessage && mt != BinaryMessage {
		return fmt.Errorf("zsocket: unexpected message type %d", mt)
	return json.Unmarshal(b, v)
// WriteControl writes a control frame using message type and deadline.
func (c *Conn) WriteControl(mt MessageType, payload []byte, deadline time.Time) error {
	var opcode byte
	switch mt {
	case CloseMessage:
		opcode = opClose
	case PingMessage:
		opcode = opPing
	case PongMessage:
		opcode = opPong
		return fmt.Errorf("zsocket: invalid control message type %d", mt)
	return c.writeControlOpcode(opcode, payload, deadline)
func (c *Conn) writeControlOpcode(opcode byte, payload []byte, deadline time.Time) error {
	if c.closed.Load() {
		return ErrCloseSent
	if !deadline.IsZero() {
		_ = c.raw.SetWriteDeadline(deadline)
		defer c.restoreWriteDeadline()
	} else if d := c.getWriteDeadline(); !d.IsZero() {
		_ = c.raw.SetWriteDeadline(d)
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if err := writeFrame(c.bw, true, opcode, payload, c.isClient, false); err != nil {
		return err
	return c.bw.Flush()
func (c *Conn) shouldCompressWrite() bool {
	c.compressionMu.RLock()
	enabled := c.cfg.EnableCompression && c.writeCompressionEnable && c.compressor != nil
	c.compressionMu.RUnlock()
	return enabled
func (c *Conn) compressData(data []byte) ([]byte, error) {
	c.compressionMu.RLock()
	cmp := c.compressor
	c.compressionMu.RUnlock()
	if cmp == nil {
		return data, nil
	return cmp.Compress(data)
func (c *Conn) decompressData(data []byte) ([]byte, error) {
	c.compressionMu.RLock()
	cmp := c.compressor
	c.compressionMu.RUnlock()
	if cmp == nil {
		return nil, fmt.Errorf("zsocket: compression not available")
	return cmp.Decompress(data)
// NextReader returns the next complete message payload as an io.Reader.
func (c *Conn) NextReader() (MessageType, io.Reader, error) {
	mt, b, err := c.Read(context.Background())
	if err != nil {
		return 0, nil, err
	return mt, bytes.NewReader(b), nil
// NextWriter returns a writer for the next message.
func (c *Conn) NextWriter(mt MessageType) (io.WriteCloser, error) {
	if mt != TextMessage && mt != BinaryMessage {
		return nil, fmt.Errorf("zsocket: invalid message type %d", mt)
	if c.closed.Load() {
		return nil, ErrCloseSent
	return &connWriter{conn: c, mt: mt}, nil
type connWriter struct {
	conn   *Conn
	mt     MessageType
	buf    bytes.Buffer
	closed bool
func (w *connWriter) Write(p []byte) (int, error) {
	if w.closed {
		return 0, ErrCloseSent
	return w.buf.Write(p)
func (w *connWriter) Close() error {
	if w.closed {
		return nil
	w.closed = true
	return w.conn.Write(context.Background(), w.mt, w.buf.Bytes())
func (c *Conn) getReadDeadline() time.Time {
	c.deadlineMu.RLock()
	d := c.readDeadline
	c.deadlineMu.RUnlock()
func (c *Conn) getWriteDeadline() time.Time {
	c.deadlineMu.RLock()
	d := c.writeDeadline
	c.deadlineMu.RUnlock()
func (c *Conn) restoreReadDeadline() {
	_ = c.raw.SetReadDeadline(c.getReadDeadline())
func (c *Conn) restoreWriteDeadline() {
	_ = c.raw.SetWriteDeadline(c.getWriteDeadline())
func toMsgType(opcode byte) MessageType {
	if opcode == opText {
		return TextMessage
	return BinaryMessage
// closePayload builds a close frame payload (2-byte code + optional UTF-8 reason).
func closePayload(code CloseCode, reason string) []byte {
	if code == 0 {
		return nil
	// 2 bytes code + reason
	buf := make([]byte, 2+len(reason))
	binary.BigEndian.PutUint16(buf[:2], uint16(code))
	copy(buf[2:], []byte(reason))
	return buf
// parseClosePayload extracts code & reason from a close frame payload.
func parseClosePayload(p []byte) (CloseCode, string) {
	if len(p) < 2 {
		return CloseNoStatusRcvd, ""
	code := CloseCode(binary.BigEndian.Uint16(p[:2]))
	reason := string(p[2:])
	return code, reason
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

conn.go

Latest commit

History

conn.go

File metadata and controls
