traefik/server/server.mjs

import WebSocket, { WebSocketServer } from 'ws';
import express from 'express';
import * as uuid from 'uuid';

import { stringify as sortedJson } from './lib/sortedJson.mjs';

let config = {};
const clientIps = new Map();
const httpAnnouncements = new Map();

const renderConfig = () => {
  let newConfig = {
    http: {
      routers: {},
      services: {},
    },
  };
  for (const [
    name,
    {
      clients,
      announcement: { entrypoints = [], middleware = [], ...announcement },
    },
  ] of httpAnnouncements.entries()) {
    newConfig.http.routers[name] = {
      entryPoints: entrypoints,
      middlewares: middleware,
      service: [],
      priority: 1,
      ...announcement,
      service: name,
    };
    newConfig.http.services[name] = {
      loadbalancer: {
        servers: Array.from(clients).map((regId) => {
          const [clientId, port] = regId.split(':');
          return { url: `http://${clientIps.get(clientId)}:${port}` };
        }),
      },
    };
  }
  config = newConfig;
  console.log(JSON.stringify(config, null, 2));
};

const removeClient = (clientId) => {
  const regIdPrefix = `${clientId}:`;
  for (const [name, announcement] of httpAnnouncements.entries()) {
    const newClients = new Set();
    for (const regId of announcement.clients) {
      if (!regId.startsWith(regIdPrefix)) {
        newClients.add(regId);
      }
    }
    if (newClients.size !== 0) {
      announcement.clients = newClients;
    } else {
      httpAnnouncements.delete(name);
    }
  }
  clientIps.delete(clientId);
  renderConfig();
};

const handlers = {
  http: {
    announce(clientId, { name, port, ...announcement }) {
      //
      const str = sortedJson(announcement);
      if (httpAnnouncements.has(name)) {
        const existingAnnouncement = httpAnnouncements.get(name);
        if (existingAnnouncement.str !== str) {
          console.log(
            `rejecting announcement from ${clientId} due to mismatching definition`
          );
          return false;
        }
        existingAnnouncement.clients.add(`${clientId}:${port}`);
        return true;
      }
      httpAnnouncements.set(name, {
        announcement: announcement,
        clients: new Set([`${clientId}:${port}`]),
        str,
      });
      return true;
    },
    retract(clientId, { name, port, ...announcement }) {
      if (!httpAnnouncements.has(name)) {
        console.log(
          `ignoring retraction from ${clientId} due to missing announcement`
        );
        return false;
      }
      const str = sortedJson(announcement);
      const existingAnnouncement = httpAnnouncements.get(name);
      if (existingAnnouncement.str !== str) {
        console.log(
          `ignoring retraction from ${clientId} due to mismatching definition`
        );
        return false;
      }
      existingAnnouncement.clients.delete(`${clientId}:${port}`);
      if (existingAnnouncement.clients.size === 0) {
        httpAnnouncements.delete(name);
      }
      return true;
    },
  },
  'http-middleware-forward-auth': {
    announce() {
      throw new Error('todo');
    },
    retract() {
      throw new Error('todo');
    },
  },
};

const app = express();

app.get('/json', (req, res) => {
  res.json(config);
});

const wss = new WebSocketServer({ server: app });
wss.on('connection', (ws, request) => {
  const clientId = uuid.v4();
  if (request.socket.remoteFamily === 'IPv6') {
    const ip = /^::ffff:(?<ip>((25[0-5]|(2[0-4]|1\d|[1-9]|)\d)\.?\b){4})$/.exec(
      request.socket.remoteAddress
    )?.groups?.ip;
    if (typeof ip !== 'string') {
      ws.close();
      return;
    }

    clientIps.set(clientId, ip);
  } else if (request.socket.remoteFamily === 'IPv4') {
    clientIps.set(clientId, request.socket.remoteAddress);
  } else {
    console.log(`unknown remoteFamily: ${request.socket.remoteAddress}`);
    ws.close();
    return;
  }
  ws.on('error', (err) => {
    removeClient(clientId);
  });

  ws.on('message', (msg) => {
    const { type, payload } = JSON.parse(msg);
    const [handler, method] = type.split('/');
    if (!Object.hasOwn(handlers, handler)) {
      console.log(`no handler for ${type}`);
      return;
    }
    const methods = handlers[handler];
    if (!Object.hasOwn(handlers, handler)) {
      console.log(`no handler for ${type}`);
      return;
    }
    methods[method](clientId, payload);
    // add announcement
    renderConfig();
  });

  ws.on('close', () => {
    removeClient(clientId);
  });
});

const server = app.listen(3000);
server.on('upgrade', (request, socket, head) => {
  wss.handleUpgrade(request, socket, head, (socket) => {
    wss.emit('connection', socket, request);
  });
});