Real-Time Communication in Go: TCP Sockets, Message Framing, and WebSockets

buf := make([]byte, 1024)
n, err := conn.Read(buf)
msg := string(buf[:n])

package framing

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
)

const maxMessageSize = 16 * 1024 * 1024 // 16 MB safety cap

// WriteMessage writes a length-prefixed message to conn.
func WriteMessage(conn net.Conn, data []byte) error {
	if len(data) > maxMessageSize {
		return fmt.Errorf("message too large: %d bytes", len(data))
	}
	header := make([]byte, 4)
	binary.BigEndian.PutUint32(header, uint32(len(data)))
	if _, err := conn.Write(header); err != nil {
		return fmt.Errorf("write header: %w", err)
	}
	if _, err := conn.Write(data); err != nil {
		return fmt.Errorf("write body: %w", err)
	}
	return nil
}

// ReadMessage reads a length-prefixed message from conn.
// It blocks until a complete message is available.
func ReadMessage(conn net.Conn) ([]byte, error) {
	header := make([]byte, 4)
	if _, err := io.ReadFull(conn, header); err != nil {
		return nil, fmt.Errorf("read header: %w", err)
	}
	size := binary.BigEndian.Uint32(header)
	if size > maxMessageSize {
		return nil, fmt.Errorf("message size %d exceeds limit", size)
	}
	body := make([]byte, size)
	if _, err := io.ReadFull(conn, body); err != nil {
		return nil, fmt.Errorf("read body: %w", err)
	}
	return body, nil
}

package server

import (
	"fmt"
	"log"
	"net"
	"sync"

	"github.com/yourname/chat/framing"
)

type Server struct {
	mu      sync.RWMutex
	clients map[net.Conn]struct{}
	bcast   chan []byte
}

func New() *Server {
	return &Server{
		clients: make(map[net.Conn]struct{}),
		bcast:   make(chan []byte, 256),
	}
}

func (s *Server) Run(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return fmt.Errorf("listen: %w", err)
	}
	defer ln.Close()
	log.Printf("TCP server listening on %s", addr)

	go s.broadcaster()

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("accept: %v", err)
			continue
		}
		s.addClient(conn)
		go s.handleConn(conn)
	}
}

func (s *Server) broadcaster() {
	for msg := range s.bcast {
		s.mu.RLock()
		for conn := range s.clients {
			// Set a write deadline so a slow client cannot block the broadcast.
			_ = conn.SetWriteDeadline(timeoutFrom(2))
			if err := framing.WriteMessage(conn, msg); err != nil {
				log.Printf("broadcast write: %v", err)
			}
		}
		s.mu.RUnlock()
	}
}

func (s *Server) handleConn(conn net.Conn) {
	defer func() {
		s.removeClient(conn)
		conn.Close()
	}()

	for {
		// Read deadline: disconnect idle clients after 60s.
		_ = conn.SetReadDeadline(timeoutFrom(60))

		msg, err := framing.ReadMessage(conn)
		if err != nil {
			log.Printf("read from %s: %v", conn.RemoteAddr(), err)
			return
		}
		s.bcast <- msg
	}
}

func (s *Server) addClient(conn net.Conn) {
	s.mu.Lock()
	s.clients[conn] = struct{}{}
	s.mu.Unlock()
	log.Printf("client connected: %s (total: %d)", conn.RemoteAddr(), s.clientCount())
}

func (s *Server) removeClient(conn net.Conn) {
	s.mu.Lock()
	delete(s.clients, conn)
	s.mu.Unlock()
	log.Printf("client disconnected: %s (total: %d)", conn.RemoteAddr(), s.clientCount())
}

func (s *Server) clientCount() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.clients)
}

package wsserver

import (
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true // Restrict in production: check r.Header.Get("Origin")
	},
}

type Hub struct {
	mu      sync.RWMutex
	clients map[*websocket.Conn]struct{}
	bcast   chan []byte
}

func NewHub() *Hub {
	h := &Hub{
		clients: make(map[*websocket.Conn]struct{}),
		bcast:   make(chan []byte, 256),
	}
	go h.broadcaster()
	return h
}

func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("upgrade: %v", err)
		return
	}
	h.addClient(conn)
	defer func() {
		h.removeClient(conn)
		conn.Close()
	}()

	conn.SetReadLimit(16 * 1024 * 1024)
	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(_ string) error {
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("ws read error: %v", err)
			}
			return
		}
		h.bcast <- msg
	}
}

func (h *Hub) broadcaster() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case msg := <-h.bcast:
			h.mu.RLock()
			for conn := range h.clients {
				conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
				if err := conn.WriteMessage(websocket.TextMessage, msg); err != nil {
					log.Printf("ws write: %v", err)
				}
			}
			h.mu.RUnlock()
		case <-ticker.C:
			// Send ping to all clients to keep connections alive.
			h.mu.RLock()
			for conn := range h.clients {
				conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
				_ = conn.WriteMessage(websocket.PingMessage, nil)
			}
			h.mu.RUnlock()
		}
	}
}

func (h *Hub) addClient(conn *websocket.Conn) {
	h.mu.Lock()
	h.clients[conn] = struct{}{}
	h.mu.Unlock()
}

func (h *Hub) removeClient(conn *websocket.Conn) {
	h.mu.Lock()
	delete(h.clients, conn)
	h.mu.Unlock()
}