job-search/internal/websocket/tcp.go

package websocket

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/netip"
	"strings"
	"sync"
	"time"

	ws "github.com/coder/websocket"
)

// Plain TCP clients join the hub as if they were WebSocket clients. Instead
// of WebSocket frames, each message is one line of JSON (an Envelope)
// terminated by "\n" in both directions, so a client can be as simple as
// netcat:
//
//	$ nc localhost 9393
//	{"type":"welcome","payload":{"clientId":"..."}}
//	{"type":"ping","id":"1"}
//	{"type":"pong","id":"1","payload":{"time":"..."}}
//
// A trailing "\r" is ignored and blank lines are skipped. A line longer than
// the websocket max_message_size closes the connection, as an oversized frame
// does. When the server closes a connection (shutdown, slow consumer, a line
// too long) it sends an error envelope with the reason first.
//
// Plain TCP has no ping frames, so liveness uses TCP keepalives probing at the
// ping interval; Ping itself writes a blank line, which clients ignore, so a
// dead peer is found when a write fails.
//
// There is no HTTP request, so there is no CSRF or Origin check and no proxy
// handling: the client's address is the TCP peer. Only listen where clients
// are trusted (e.g. localhost), or wrap the listener (e.g. with PROXY protocol
// or TLS) before passing its connections to ServeTCP. ServeTCP doesn't limit
// connections; package tcpserver applies websocket.connection_limit.

// ServeTCP runs a plain TCP connection as a hub client until it closes, as
// ServeHTTP does for a WebSocket upgrade. Cancelling ctx drops the client
// without a message; use Shutdown to close clients gracefully. The
// connection is closed when ServeTCP returns.
func (h *Hub) ServeTCP(ctx context.Context, newConn net.Conn) {
	if tcpSocket, ok := newConn.(*net.TCPConn); ok && h.conf.PingInterval > 0 {
		tcpSocket.SetKeepAliveConfig(net.KeepAliveConfig{ //nolint:errcheck,gosec // best effort
			Enable:   true,
			Idle:     h.conf.PingInterval,
			Interval: h.conf.PingInterval,
			Count:    3,
		})
	}

	ip := ""
	if addrPort, err := netip.ParseAddrPort(newConn.RemoteAddr().String()); err == nil {
		ip = addrPort.Addr().Unmap().String()
	}

	conn := newTCPConn(newConn, h.maxMessageSize(), h.conf.WriteTimeout)

	ctx, cancel := context.WithCancel(ctx)
	client := newClient(h, conn, ip, cancel)

	if !h.register(client) {
		cancel()
		conn.Close(ws.StatusGoingAway, "server shutting down") //nolint:errcheck,gosec // turned away
		return
	}
	defer h.unregister(client)

	h.log.Info("client connected", "client", client.id, "ip", ip, "transport", "tcp")
	client.run(ctx)
	conn.CloseNow() //nolint:errcheck,gosec // finished with; it may already be closed
	h.log.Info("client disconnected", "client", client.id, "ip", ip, "transport", "tcp")
}

// ---------- tcpConn ----------

// tcpConn adapts a net.Conn carrying newline-delimited JSON to the msgConn a
// client uses. Read is only called from the client's read loop; writes may
// come from the write loop, Ping and Close concurrently, so they are
// serialised.
type tcpConn struct {
	conn         net.Conn
	scan         *bufio.Scanner
	closeTimeout time.Duration // caps Close's message; 0 or less for none

	wm sync.Mutex // serialises writes

	closeOnce sync.Once     // records the first close
	closed    chan struct{} // closed by the first Close or CloseNow
	closeErr  ws.CloseError // why the connection was closed
}

// newTCPConn wraps conn. maxLine caps the length of a line read
// (defaultMaxMessageSize if 0 or less) and closeTimeout how long Close spends
// sending its message.
func newTCPConn(
	conn net.Conn, maxLine int, closeTimeout time.Duration,
) *tcpConn {
	if maxLine <= 0 {
		maxLine = defaultMaxMessageSize
	}

	scan := bufio.NewScanner(conn)
	// The buffer must hold the line and its "\n".
	scan.Buffer(make([]byte, 0, min(maxLine+1, 4096)), maxLine+1)

	return &tcpConn{
		conn:         conn,
		scan:         scan,
		closeTimeout: closeTimeout,
		closed:       make(chan struct{}),
	}
}

// Read returns the next non-blank line as a text message. Cancelling ctx
// interrupts it. The end of the stream is reported as a normal closure, and a
// read after Close as the close status Close sent, as a *websocket.Conn
// reports them, so the client treats them as a finished close handshake.
func (t *tcpConn) Read(ctx context.Context) (ws.MessageType, []byte, error) {
	stop := context.AfterFunc(ctx, func() {
		t.conn.SetReadDeadline(time.Now()) //nolint:errcheck,gosec // only interrupts the read
	})
	defer stop()

	for t.scan.Scan() {
		line := strings.TrimSuffix(t.scan.Text(), "\r")
		if strings.TrimSpace(line) == "" {
			continue
		}
		return ws.MessageText, []byte(line), nil
	}

	err := t.scan.Err()
	switch {
	case ctx.Err() != nil:
		return 0, nil, ctx.Err()
	case t.isClosed():
		return 0, nil, t.closeErr
	case err == nil || errors.Is(err, io.EOF):
		return 0, nil, ws.CloseError{
			Code:   ws.StatusNormalClosure,
			Reason: "connection closed",
		}
	case errors.Is(err, bufio.ErrTooLong):
		// As a *websocket.Conn closes with 1009 for an oversized message.
		// The client may not see the reason: closing with its unread data
		// still buffered makes the connection reset.
		t.Close(ws.StatusMessageTooBig, "message too long") //nolint:errcheck,gosec // dropped anyway
		return 0, nil, t.closeErr
	}
	return 0, nil, err
}

// Write sends payload as one line. Only text messages are supported.
func (t *tcpConn) Write(
	ctx context.Context, messageType ws.MessageType, payload []byte,
) error {
	if messageType != ws.MessageText {
		return errUnsupportedData
	}
	line := make([]byte, 0, len(payload)+1)
	line = append(append(line, payload...), '\n')
	return t.write(ctx, line)
}

// Ping writes a blank line, which clients ignore, so a dead peer shows up as a
// failed write. TCP keepalives (see ServeTCP) do the rest.
func (t *tcpConn) Ping(ctx context.Context) error {
	return t.write(ctx, []byte("\n"))
}

// Close sends an error envelope giving the reason, then closes the
// connection. Only the first Close or CloseNow does anything; later calls
// return nil. A CloseNow while Close is sending its message cuts it short.
func (t *tcpConn) Close(code ws.StatusCode, reason string) error {
	if !t.markClosed(ws.CloseError{Code: code, Reason: reason}) {
		return nil
	}

	ctx := context.Background()
	if t.closeTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, t.closeTimeout)
		defer cancel()
	}

	message, err := json.Marshal(Envelope{
		Type:  TypeError,
		Error: reason,
	})
	if err == nil {
		t.write(ctx, append(message, '\n')) //nolint:errcheck,gosec // best effort, closing anyway
	}

	if err := t.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
		return err
	}
	return nil
}

// CloseNow closes the connection without a message. It doesn't wait for a
// Close that is sending its message.
func (t *tcpConn) CloseNow() error {
	t.markClosed(ws.CloseError{Code: ws.StatusAbnormalClosure})
	if err := t.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
		return err
	}
	return nil
}

// markClosed records why the connection is being closed and reports whether
// this is the first close. It never blocks, so a Close sending its message
// can't hold up CloseNow.
func (t *tcpConn) markClosed(why ws.CloseError) bool {
	first := false
	t.closeOnce.Do(func() {
		t.closeErr = why
		close(t.closed)
		first = true
	})
	return first
}

// isClosed reports whether Close or CloseNow has been called.
func (t *tcpConn) isClosed() bool {
	select {
	case <-t.closed:
		return true
	default:
		return false
	}
}

// write writes data whole, giving up at ctx's deadline or when ctx is
// cancelled.
func (t *tcpConn) write(ctx context.Context, data []byte) error {
	t.wm.Lock()
	defer t.wm.Unlock()

	deadline, _ := ctx.Deadline() // zero means none
	if err := t.conn.SetWriteDeadline(deadline); err != nil {
		return err
	}

	stop := context.AfterFunc(ctx, func() {
		t.conn.SetWriteDeadline(time.Now()) //nolint:errcheck,gosec // only interrupts the write
	})
	defer stop()

	_, err := t.conn.Write(data)
	if err != nil && ctx.Err() != nil {
		return ctx.Err()
	}

	return err
}