WebSocket Engineer Skill Skill

WebSocket Engineer Skill

Overview

Real-time, bidirectional communication systems using WebSockets and related protocols. Covers protocol fundamentals, server implementation, authentication, scaling, fault tolerance, and browser client patterns.

Protocol Fundamentals

HTTP Upgrade → WebSocket Handshake
GET /ws HTTP/1.1
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: <base64 random>
Sec-WebSocket-Version: 13

HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: <SHA-1 of key + GUID>

After handshake: full-duplex binary framing
Frame types: text (0x1), binary (0x2), close (0x8), ping (0x9), pong (0xA)

Server Implementation — ws (Node.js)

const { WebSocketServer, WebSocket } = require('ws');
const http = require('http');

const server = http.createServer(app);
const wss = new WebSocketServer({ server, path: '/ws' });

// Connection map for targeted sends
const clients = new Map(); // userId → Set<WebSocket>

wss.on('connection', (ws, req) => {
  // 1. Authenticate immediately
  const token = extractToken(req);
  const user = verifyToken(token);
  if (!user) {
    ws.close(4001, 'Unauthorized');
    return;
  }

  // 2. Register client
  if (!clients.has(user.id)) clients.set(user.id, new Set());
  clients.get(user.id).add(ws);
  ws.userId = user.id;

  // 3. Heartbeat setup
  ws.isAlive = true;
  ws.on('pong', () => {
    ws.isAlive = true;
  });

  // 4. Message handling
  ws.on('message', (data, isBinary) => {
    try {
      const msg = isBinary ? data : JSON.parse(data.toString());
      handleMessage(ws, msg);
    } catch (err) {
      ws.send(JSON.stringify({ type: 'error', message: 'Invalid message' }));
    }
  });

  // 5. Cleanup on close
  ws.on('close', (code, reason) => {
    clients.get(ws.userId)?.delete(ws);
    if (clients.get(ws.userId)?.size === 0) clients.delete(ws.userId);
    console.log(`Client ${ws.userId} disconnected: ${code} ${reason}`);
  });

  ws.on('error', err => {
    console.error(`WebSocket error for ${ws.userId}:`, err.message);
  });

  // 6. Send initial state
  ws.send(JSON.stringify({ type: 'connected', userId: user.id }));
});

// Heartbeat interval — terminate dead connections every 30s
const heartbeat = setInterval(() => {
  wss.clients.forEach(ws => {
    if (!ws.isAlive) {
      ws.terminate();
      return;
    }
    ws.isAlive = false;
    ws.ping();
  });
}, 30_000);

wss.on('close', () => clearInterval(heartbeat));

// Targeted send to specific user (all connections)
function sendToUser(userId, payload) {
  const userClients = clients.get(userId);
  if (!userClients) return;
  const data = JSON.stringify(payload);
  userClients.forEach(ws => {
    if (ws.readyState === WebSocket.OPEN) ws.send(data);
  });
}

// Broadcast to all connected clients
function broadcast(payload, excludeUserId = null) {
  const data = JSON.stringify(payload);
  wss.clients.forEach(ws => {
    if (ws.readyState === WebSocket.OPEN && ws.userId !== excludeUserId) {
      ws.send(data);
    }
  });
}

Authentication Patterns

// Token in query string (simple, TLS required)
wss.on('connection', (ws, req) => {
  const url = new URL(req.url, 'http://localhost');
  const token = url.searchParams.get('token');
  // ...
});

// Cookie-based (preferred for browsers)
wss.on('connection', (ws, req) => {
  const cookies = parseCookies(req.headers.cookie);
  const sessionId = cookies['session'];
  // ...
});

// Custom header via Sec-WebSocket-Protocol (workaround for headers limitation)
// Browser: new WebSocket(url, ['v1', token])
// Server: req.headers['sec-websocket-protocol']

// First message auth (for frameworks that don't support upgrade auth)
ws.once('message', data => {
  const { type, token } = JSON.parse(data);
  if (type !== 'auth') {
    ws.close(4001, 'Auth required');
    return;
  }
  const user = verifyToken(token);
  if (!user) {
    ws.close(4001, 'Unauthorized');
    return;
  }
  ws.userId = user.id;
  setupMessageHandler(ws);
});

Horizontal Scaling with Redis

const { createClient } = require('redis');
const { WebSocketServer } = require('ws');

const pub = createClient({ url: process.env.REDIS_URL });
const sub = createClient({ url: process.env.REDIS_URL });
await pub.connect();
await sub.connect();

const CHANNEL = 'ws:broadcast';

// Subscribe to Redis messages and relay to local clients
await sub.subscribe(CHANNEL, message => {
  const payload = JSON.parse(message);
  // Deliver to local connections matching the target
  wss.clients.forEach(ws => {
    if (ws.readyState === WebSocket.OPEN) {
      if (!payload.targetUserId || ws.userId === payload.targetUserId) {
        ws.send(JSON.stringify(payload.data));
      }
    }
  });
});

// Publish cross-instance — all servers receive and relay
async function broadcastViaRedis(data, targetUserId = null) {
  await pub.publish(
    CHANNEL,
    JSON.stringify({ data, targetUserId, serverId: process.env.SERVER_ID })
  );
}

// Rooms with Redis Sets
async function joinRoom(userId, room) {
  await pub.sAdd(`room:${room}`, userId);
}

async function broadcastToRoom(room, data) {
  const members = await pub.sMembers(`room:${room}`);
  for (const userId of members) {
    await pub.publish(`user:${userId}`, JSON.stringify(data));
  }
}

Socket.IO (Higher-Level Abstraction)

const { Server } = require('socket.io');
const { createAdapter } = require('@socket.io/redis-adapter');

const io = new Server(httpServer, {
  cors: { origin: process.env.FRONTEND_URL, credentials: true },
  adapter: createAdapter(pub, sub), // Redis adapter for scaling
  transports: ['websocket', 'polling'], // Prefer WebSocket, fallback to polling
  pingTimeout: 20_000,
  pingInterval: 10_000,
});

// Middleware for auth
io.use(async (socket, next) => {
  const token = socket.handshake.auth.token;
  const user = await verifyToken(token);
  if (!user) return next(new Error('Unauthorized'));
  socket.userId = user.id;
  next();
});

io.on('connection', socket => {
  socket.join(`user:${socket.userId}`); // Auto room per user

  socket.on('join-room', roomId => {
    socket.join(roomId);
    socket.to(roomId).emit('user-joined', { userId: socket.userId });
  });

  socket.on('message', ({ roomId, text }) => {
    // Broadcast to room except sender
    socket.to(roomId).emit('message', { userId: socket.userId, text, ts: Date.now() });
  });

  socket.on('disconnect', reason => {
    console.log(`${socket.userId} disconnected: ${reason}`);
  });
});

// Server-side targeted emit
io.to(`user:${userId}`).emit('notification', { message: 'You have a new message' });
io.to(roomId).emit('event', data);

Browser Client

class RealtimeClient {
  #ws = null;
  #reconnectDelay = 1_000;
  #maxDelay = 30_000;
  #handlers = new Map();

  constructor(url, token) {
    this.url = url;
    this.token = token;
    this.connect();
  }

  connect() {
    this.#ws = new WebSocket(`${this.url}?token=${this.token}`);

    this.#ws.onopen = () => {
      console.log('Connected');
      this.#reconnectDelay = 1_000; // Reset backoff
      this.emit('connected');
    };

    this.#ws.onmessage = ({ data }) => {
      try {
        const msg = JSON.parse(data);
        this.#handlers.get(msg.type)?.forEach(cb => cb(msg));
        this.#handlers.get('*')?.forEach(cb => cb(msg));
      } catch (e) {
        console.error('Parse error', e);
      }
    };

    this.#ws.onclose = event => {
      if (event.code === 4001) {
        this.emit('unauthorized');
        return; // Don't reconnect on auth failure
      }
      console.log(`Disconnected (${event.code}), retrying in ${this.#reconnectDelay}ms`);
      setTimeout(() => this.connect(), this.#reconnectDelay);
      this.#reconnectDelay = Math.min(this.#reconnectDelay * 2, this.#maxDelay);
    };

    this.#ws.onerror = err => {
      console.error('WebSocket error', err);
    };
  }

  on(type, handler) {
    if (!this.#handlers.has(type)) this.#handlers.set(type, new Set());
    this.#handlers.get(type).add(handler);
    return () => this.#handlers.get(type).delete(handler); // Returns unsubscribe fn
  }

  send(type, data = {}) {
    if (this.#ws?.readyState !== WebSocket.OPEN) {
      console.warn('Not connected, dropping message:', type);
      return;
    }
    this.#ws.send(JSON.stringify({ type, ...data }));
  }

  emit(type, data) {
    this.#handlers.get(type)?.forEach(cb => cb(data));
  }

  close() {
    this.#ws?.close(1000, 'Client closing');
  }
}

// Usage
const client = new RealtimeClient('wss://api.example.com/ws', authToken);
const unsubscribe = client.on('message', msg => displayMessage(msg));
client.send('join-room', { roomId: 'general' });
// Cleanup:
unsubscribe();

Server-Sent Events (SSE) — Unidirectional Alternative

// Server — simpler than WebSockets for server → client push
app.get('/events', (req, res) => {
  const user = authenticate(req);
  if (!user) {
    res.sendStatus(401);
    return;
  }

  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');
  res.flushHeaders();

  // Send events
  const send = (event, data) => {
    res.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`);
  };

  send('connected', { userId: user.id });

  // Register for push events
  const unsubscribe = eventBus.on(`user:${user.id}`, data => send('update', data));
  const heartbeat = setInterval(() => res.write(':heartbeat\n\n'), 15_000);

  req.on('close', () => {
    clearInterval(heartbeat);
    unsubscribe();
  });
});

// Browser — SSE client
const source = new EventSource('/events');
source.addEventListener('update', e => {
  const data = JSON.parse(e.data);
  updateUI(data);
});
source.onerror = () => {
  // EventSource auto-reconnects — onerror fires on reconnect attempts too
  console.log('SSE connection lost, browser will reconnect automatically');
};

Message Protocol Design

// Typed message envelope (recommended pattern)
const MessageTypes = {
  // Client → Server
  JOIN_ROOM: 'join_room',
  LEAVE_ROOM: 'leave_room',
  SEND_MESSAGE: 'send_message',
  TYPING_START: 'typing_start',
  TYPING_STOP: 'typing_stop',
  // Server → Client
  ROOM_JOINED: 'room_joined',
  MESSAGE_RECEIVED: 'message_received',
  USER_JOINED: 'user_joined',
  USER_LEFT: 'user_left',
  ERROR: 'error',
};

// Example message schema
const sendMessageSchema = {
  type: 'send_message',
  roomId: 'string (required)',
  text: 'string (max 2000 chars)',
  replyTo: 'string | null (message ID)',
};

// Server validation
function validateMessage(msg) {
  if (!msg.type || typeof msg.type !== 'string') throw new Error('Missing type');
  if (!MessageTypes[msg.type.toUpperCase()]) throw new Error(`Unknown type: ${msg.type}`);
  // Type-specific validation
  switch (msg.type) {
    case MessageTypes.SEND_MESSAGE:
      if (!msg.roomId || !msg.text) throw new Error('Missing roomId or text');
      if (msg.text.length > 2000) throw new Error('Text too long');
      break;
  }
}

Rate Limiting

const rateLimiters = new Map(); // userId → { count, resetAt }

function checkRateLimit(userId, limit = 60, windowMs = 60_000) {
  const now = Date.now();
  const limiter = rateLimiters.get(userId) ?? { count: 0, resetAt: now + windowMs };

  if (now > limiter.resetAt) {
    limiter.count = 0;
    limiter.resetAt = now + windowMs;
  }

  limiter.count++;
  rateLimiters.set(userId, limiter);

  if (limiter.count > limit) {
    throw new Error(`Rate limit exceeded: ${limit} messages per ${windowMs / 1000}s`);
  }
}

Observability

// Metrics to track
const metrics = {
  connections: 0, // Current active connections
  messagesIn: 0, // Messages received per second
  messagesOut: 0, // Messages sent per second
  errors: 0, // Error count
  avgLatency: 0, // Round-trip time (using ping/pong timestamps)
};

// Emit latency metric with custom ping payload
function measureLatency(ws) {
  const start = Date.now();
  const payload = Buffer.alloc(8);
  payload.writeBigInt64BE(BigInt(start));
  ws.ping(payload);
  ws.once('pong', data => {
    const sent = Number(data.readBigInt64BE());
    metrics.avgLatency = Date.now() - sent;
  });
}

Security

Always use wss:// (TLS) in production — plain ws:// is MITM-vulnerable
Authenticate at the HTTP upgrade stage, not per-message
Rate-limit both connection attempts and message frequency
Validate and sanitize ALL incoming message content (treat as untrusted)
Set maxPayload on ws.WebSocketServer to prevent memory exhaustion attacks
Use ws.close(code, reason) for intentional closes; ws.terminate() only for dead connections
Never echo user-controlled data without sanitization (XSS via message relay)
Enable CORS properly for Socket.IO — do not use origin: '*' in production

// Limit message payload size (default is 100MB — too large)
const wss = new WebSocketServer({ server, maxPayload: 64 * 1024 }); // 64KB max

ws npm package: https://github.com/websockets/ws
Socket.IO: https://socket.io/docs/v4/
RFC 6455 (WebSocket Protocol): https://www.rfc-editor.org/rfc/rfc6455
MDN WebSocket API: https://developer.mozilla.org/en-US/docs/Web/API/WebSocket

Agent Skills: WebSocket Engineer Skill

Install this agent skill to your local

Skill Files