url-shrinker/server.js

import { createServer } from 'node:http';
import { DatabaseSync } from 'node:sqlite';
import { randomBytes } from 'node:crypto';
import { mkdirSync, readFileSync } from 'node:fs';
import { join, dirname } from 'node:path';
import { fileURLToPath } from 'node:url';

const PORT = Number(process.env.PORT || 3000);
const DATA_DIR = process.env.DATA_DIR || '/data';
const BASE_URL = (process.env.BASE_URL || `http://localhost:${PORT}`).replace(/\/+$/, '');
const CODE_LENGTH = 6;
const ALPHABET = 'ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz23456789';

mkdirSync(DATA_DIR, { recursive: true });
const db = new DatabaseSync(join(DATA_DIR, 'shrinker.db'));
db.exec(`
  PRAGMA journal_mode = WAL;
  CREATE TABLE IF NOT EXISTS links (
    code       TEXT PRIMARY KEY,
    url        TEXT NOT NULL,
    clicks     INTEGER NOT NULL DEFAULT 0,
    created_at TEXT NOT NULL DEFAULT (datetime('now'))
  );
`);

const insertLink = db.prepare('INSERT INTO links (code, url) VALUES (?, ?)');
const getLink = db.prepare('SELECT code, url, clicks, created_at FROM links WHERE code = ?');
const bumpClicks = db.prepare('UPDATE links SET clicks = clicks + 1 WHERE code = ?');
const countLinks = db.prepare('SELECT COUNT(*) AS n, COALESCE(SUM(clicks), 0) AS c FROM links');

const INDEX_HTML = readFileSync(join(dirname(fileURLToPath(import.meta.url)), 'index.html'));

function randomCode() {
  const bytes = randomBytes(CODE_LENGTH);
  let code = '';
  for (let i = 0; i < CODE_LENGTH; i++) code += ALPHABET[bytes[i] % ALPHABET.length];
  return code;
}

function validTarget(raw) {
  let url;
  try {
    url = new URL(raw);
  } catch {
    return null;
  }
  if (url.protocol !== 'http:' && url.protocol !== 'https:') return null;
  return url.href;
}

function sendJson(res, status, body) {
  const payload = JSON.stringify(body);
  res.writeHead(status, {
    'Content-Type': 'application/json; charset=utf-8',
    'Content-Length': Buffer.byteLength(payload),
  });
  res.end(payload);
}

function readBody(req) {
  return new Promise((resolve, reject) => {
    const chunks = [];
    let size = 0;
    req.on('data', (chunk) => {
      size += chunk.length;
      if (size > 16384) {
        reject(new Error('body too large'));
        req.destroy();
        return;
      }
      chunks.push(chunk);
    });
    req.on('end', () => resolve(Buffer.concat(chunks).toString('utf8')));
    req.on('error', reject);
  });
}

const server = createServer(async (req, res) => {
  const { pathname } = new URL(req.url, BASE_URL);

  if ((req.method === 'GET' || req.method === 'HEAD') && pathname === '/') {
    res.writeHead(200, { 'Content-Type': 'text/html; charset=utf-8' });
    res.end(req.method === 'HEAD' ? undefined : INDEX_HTML);
    return;
  }

  if (req.method === 'GET' && pathname === '/healthz') {
    const { n, c } = countLinks.get();
    sendJson(res, 200, { ok: true, links: n, clicks: c });
    return;
  }

  if (req.method === 'POST' && pathname === '/api/shorten') {
    let body;
    try {
      body = JSON.parse(await readBody(req) || '{}');
    } catch {
      sendJson(res, 400, { error: 'invalid JSON body' });
      return;
    }
    const url = validTarget(String(body.url ?? ''));
    if (!url) {
      sendJson(res, 400, { error: 'invalid url — must be absolute http(s)' });
      return;
    }

    let code = null;
    if (body.code !== undefined && body.code !== null && body.code !== '') {
      const custom = String(body.code);
      if (!/^[A-Za-z0-9_-]{3,32}$/.test(custom)) {
        sendJson(res, 400, { error: 'custom code must be 3-32 chars of [A-Za-z0-9_-]' });
        return;
      }
      if (custom === 'api' || custom === 'healthz') {
        sendJson(res, 400, { error: 'that code is reserved' });
        return;
      }
      try {
        insertLink.run(custom, url);
        code = custom;
      } catch {
        sendJson(res, 409, { error: 'code already taken' });
        return;
      }
    } else {
      for (let attempt = 0; attempt < 5 && code === null; attempt++) {
        const candidate = randomCode();
        try {
          insertLink.run(candidate, url);
          code = candidate;
        } catch {
          /* collision, retry */
        }
      }
      if (code === null) {
        sendJson(res, 500, { error: 'could not allocate a code, try again' });
        return;
      }
    }

    sendJson(res, 201, { code, url, short_url: `${BASE_URL}/${code}` });
    return;
  }

  if (req.method === 'GET' && pathname.startsWith('/api/stats/')) {
    const row = getLink.get(pathname.slice('/api/stats/'.length));
    if (!row) {
      sendJson(res, 404, { error: 'not found' });
      return;
    }
    sendJson(res, 200, { ...row, short_url: `${BASE_URL}/${row.code}` });
    return;
  }

  if (req.method === 'GET' && /^\/[A-Za-z0-9_-]{3,32}$/.test(pathname)) {
    const row = getLink.get(pathname.slice(1));
    if (row) {
      bumpClicks.run(row.code);
      res.writeHead(302, { Location: row.url, 'Cache-Control': 'no-store' });
      res.end();
      return;
    }
  }

  sendJson(res, 404, { error: 'not found' });
});

server.listen(PORT, () => {
  console.log(`url-shrinker listening on :${PORT} (base ${BASE_URL}, data ${DATA_DIR})`);
});