Update server.js

server.js

@@ -1,48 +1,76 @@
-const http = require("http");
-const fs = require("fs");
-const path = require("path");
+"use strict";
+const http   = require("http");
+const fs     = require("fs");
+const fsp    = require("fs/promises");
+const path   = require("path");
 const crypto = require("crypto");
 const { WebSocketServer } = require("ws");
-const url = require("url");
+const { URL } = require("url");
 
-// ── Paths ──────────────────────────────────────────────────────────────────
-const DATA_DIR = path.join(__dirname, "data");
+// ── Directories ────────────────────────────────────────────────────────────
+const ROOT        = __dirname;
+const PUBLIC_DIR  = path.join(ROOT, "public");
+const DATA_DIR    = path.join(ROOT, "data");
 const ACCOUNTS_DIR = path.join(DATA_DIR, "accounts");
-const PIXELS_DIR = path.join(DATA_DIR, "pixels");
-[DATA_DIR, ACCOUNTS_DIR, PIXELS_DIR].forEach((d) => fs.mkdirSync(d, { recursive: true }));
+const PIXELS_DIR  = path.join(DATA_DIR, "pixels");
 
-// ── Tiny transparent 1×1 PNG (base64) ────────────────────────────────────
-const TRANSPARENT_PNG = Buffer.from(
+for (const d of [DATA_DIR, ACCOUNTS_DIR, PIXELS_DIR, PUBLIC_DIR]) {
+  fs.mkdirSync(d, { recursive: true });
+}
+
+// ── Transparent 1×1 PNG ───────────────────────────────────────────────────
+const PIXEL_PNG = Buffer.from(
   "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==",
   "base64"
 );
 
-// ── Helpers ────────────────────────────────────────────────────────────────
-function uid() { return crypto.randomBytes(16).toString("hex"); }
-function hash(pw) { return crypto.createHash("sha256").update(pw).digest("hex"); }
+// ── MIME types ─────────────────────────────────────────────────────────────
+const MIME = {
+  ".html": "text/html; charset=utf-8",
+  ".css":  "text/css; charset=utf-8",
+  ".js":   "application/javascript; charset=utf-8",
+  ".svg":  "image/svg+xml",
+  ".ico":  "image/x-icon",
+  ".png":  "image/png",
+  ".json": "application/json",
+};
+
+// ── Pure helpers ───────────────────────────────────────────────────────────
+const uid  = () => crypto.randomBytes(16).toString("hex");
+const hash = (pw) => crypto.createHash("sha256").update(pw + "pb-salt").digest("hex");
+
 function readJSON(file) {
-  try { return JSON.parse(fs.readFileSync(file, "utf8")); } catch { return null; }
+  try { return JSON.parse(fs.readFileSync(file, "utf8")); }
+  catch { return null; }
+}
+function writeJSON(file, data) {
+  fs.writeFileSync(file + ".tmp", JSON.stringify(data, null, 2));
+  fs.renameSync(file + ".tmp", file); // atomic write
 }
-function writeJSON(file, data) { fs.writeFileSync(file, JSON.stringify(data, null, 2)); }
 
 // ── Account helpers ────────────────────────────────────────────────────────
-function accountFile(username) { return path.join(ACCOUNTS_DIR, `${username}.json`); }
-function getAccount(username) { return readJSON(accountFile(username)); }
-function saveAccount(username, data) { writeJSON(accountFile(username), data); }
+const accountFile = (u) => path.join(ACCOUNTS_DIR, `${u}.json`);
+const getAccount  = (u) => readJSON(accountFile(u));
+const saveAccount = (u, d) => writeJSON(accountFile(u), d);
 
 function registerUser(username, password) {
-  if (!/^[a-zA-Z0-9_]{3,32}$/.test(username)) return { error: "Invalid username (3-32 alphanumeric/underscore)" };
-  if (fs.existsSync(accountFile(username))) return { error: "Username taken" };
-  const account = { username, passwordHash: hash(password), tokens: [], pixels: [] };
-  saveAccount(username, account);
+  if (!/^[a-zA-Z0-9_]{3,32}$/.test(username))
+    return { error: "Username must be 3–32 characters (letters, numbers, underscore)" };
+  if (!password || password.length < 6)
+    return { error: "Password must be at least 6 characters" };
+  if (fs.existsSync(accountFile(username)))
+    return { error: "Username already taken" };
+  saveAccount(username, { username, passwordHash: hash(password), tokens: [], pixels: [] });
   return { ok: true };
 }
 
 function loginUser(username, password) {
   const acc = getAccount(username);
-  if (!acc || acc.passwordHash !== hash(password)) return { error: "Invalid credentials" };
+  if (!acc || acc.passwordHash !== hash(password))
+    return { error: "Invalid username or password" };
   const token = uid();
-  acc.tokens.push(token);
+  // Cap active sessions at 20 to prevent unbounded growth
+  acc.tokens = [...acc.tokens.slice(-19), token];
   saveAccount(username, acc);
   return { token, username };
 }
@@ -55,144 +83,135 @@ function logoutUser(username, token) {
 }
 
 function verifyToken(username, token) {
+  if (!username || !token) return false;
   const acc = getAccount(username);
-  return acc && acc.tokens.includes(token);
-}
-
-function getUserByToken(token) {
-  const files = fs.readdirSync(ACCOUNTS_DIR);
-  for (const f of files) {
-    const acc = readJSON(path.join(ACCOUNTS_DIR, f));
-    if (acc && acc.tokens.includes(token)) return acc;
-  }
-  return null;
+  return !!(acc && acc.tokens.includes(token));
 }
 
 // ── Pixel helpers ──────────────────────────────────────────────────────────
-function pixelFile(pixelId) { return path.join(PIXELS_DIR, `${pixelId}.json`); }
-function getPixel(pixelId) { return readJSON(pixelFile(pixelId)); }
+const pixelFile = (id) => path.join(PIXELS_DIR, `${id}.json`);
+const getPixel  = (id) => readJSON(pixelFile(id));
 
 function createPixel(username, label) {
   const acc = getAccount(username);
   if (!acc) return { error: "Account not found" };
-  const pixelId = uid();
-  const pixel = {
-    id: pixelId,
+  const id = uid();
+  writeJSON(pixelFile(id), {
+    id,
     owner: username,
-    label: label || "Untitled",
+    label: (label || "Untitled").slice(0, 80),
     createdAt: Date.now(),
     opens: [],
-  };
-  writeJSON(pixelFile(pixelId), pixel);
-  acc.pixels.push(pixelId);
+  });
+  acc.pixels.push(id);
   saveAccount(username, acc);
-  return { pixelId };
+  return { pixelId: id };
 }
 
 function deletePixel(username, pixelId) {
   const pixel = getPixel(pixelId);
   if (!pixel || pixel.owner !== username) return { error: "Not found" };
-  fs.unlinkSync(pixelFile(pixelId));
+  try { fs.unlinkSync(pixelFile(pixelId)); } catch {}
   const acc = getAccount(username);
-  acc.pixels = acc.pixels.filter((p) => p !== pixelId);
-  saveAccount(username, acc);
+  if (acc) {
+    acc.pixels = acc.pixels.filter((p) => p !== pixelId);
+    saveAccount(username, acc);
+  }
   return { ok: true };
 }
 
 function getUserPixels(username) {
   const acc = getAccount(username);
   if (!acc) return [];
-  return acc.pixels
-    .map((id) => getPixel(id))
-    .filter(Boolean)
-    .sort((a, b) => b.createdAt - a.createdAt);
+  return acc.pixels.map(getPixel).filter(Boolean).sort((a, b) => b.createdAt - a.createdAt);
 }
 
-// ── WebSocket connection registry ──────────────────────────────────────────
-// Map: username → Set of { ws, type } where type is "browser" | "sw"
+// ── WebSocket registry ─────────────────────────────────────────────────────
+// username → Set<ws>
 const connections = new Map();
 
-function addConnection(username, ws, type) {
+function addConn(username, ws) {
   if (!connections.has(username)) connections.set(username, new Set());
-  connections.get(username).add({ ws, type });
+  connections.get(username).add(ws);
 }
 
-function removeConnection(username, ws) {
-  const conns = connections.get(username);
-  if (!conns) return;
-  for (const entry of conns) {
-    if (entry.ws === ws) { conns.delete(entry); break; }
-  }
-  if (conns.size === 0) connections.delete(username);
+function removeConn(username, ws) {
+  const s = connections.get(username);
+  if (!s) return;
+  s.delete(ws);
+  if (s.size === 0) connections.delete(username);
 }
 
-function notifyUser(username, payload) {
-  const conns = connections.get(username);
-  if (!conns) return;
+function broadcast(username, payload) {
+  const s = connections.get(username);
+  if (!s) return;
   const msg = JSON.stringify(payload);
-  for (const { ws } of conns) {
-    if (ws.readyState === 1) ws.send(msg);
+  for (const ws of s) {
+    if (ws.readyState === 1 /* OPEN */) ws.send(msg);
   }
 }
 
-// ── HTTP Server ────────────────────────────────────────────────────────────
-const STATIC = { ".html": "text/html", ".css": "text/css", ".js": "application/javascript" };
+// ── Auth header parser ─────────────────────────────────────────────────────
+// Format: "Bearer username:token"  — username is alphanumeric so first : is the split
+function parseAuth(header) {
+  const raw = (header || "").replace(/^Bearer\s+/, "");
+  const idx = raw.indexOf(":");
+  if (idx < 1) return [null, null];
+  return [raw.slice(0, idx), raw.slice(idx + 1)];
+}
 
+// ── HTTP request handler ───────────────────────────────────────────────────
 const server = http.createServer((req, res) => {
-  const parsed = url.parse(req.url, true);
+  // Parse URL safely
+  let parsed;
+  try { parsed = new URL(req.url, "http://x"); }
+  catch { res.writeHead(400); return res.end("Bad request"); }
   const pathname = parsed.pathname;
 
-  // CORS
+  // CORS pre-flight
   res.setHeader("Access-Control-Allow-Origin", "*");
   res.setHeader("Access-Control-Allow-Headers", "Content-Type, Authorization");
   if (req.method === "OPTIONS") { res.writeHead(204); return res.end(); }
 
-  // ── Pixel image endpoint ────────────────────────────────────────────────
+  // ── 1. Pixel tracker ────────────────────────────────────────────────────
   if (pathname.startsWith("/px/")) {
-    const pixelId = pathname.slice(4).replace(/\/$/, "");
+    const pixelId = pathname.slice(4).replace(/\//g, "");
     const pixel = getPixel(pixelId);
     if (!pixel) { res.writeHead(404); return res.end(); }
 
-    const ip = req.headers["x-forwarded-for"] || req.socket.remoteAddress;
-    const ua = req.headers["user-agent"] || "";
-    const openEntry = { ts: Date.now(), ip, ua };
-    pixel.opens.push(openEntry);
+    const ip    = (req.headers["x-forwarded-for"] || req.socket.remoteAddress || "").split(",")[0].trim();
+    const ua    = req.headers["user-agent"] || "";
+    const entry = { ts: Date.now(), ip, ua };
+    pixel.opens.push(entry);
     writeJSON(pixelFile(pixelId), pixel);
 
-    // Notify all connected sessions of the owner
-    notifyUser(pixel.owner, {
-      type: "open",
-      pixelId: pixel.id,
-      label: pixel.label,
-      ts: openEntry.ts,
-      ip,
-      ua,
-    });
+    broadcast(pixel.owner, { type: "open", pixelId: pixel.id, label: pixel.label, ...entry });
 
     res.writeHead(200, {
-      "Content-Type": "image/png",
-      "Content-Length": TRANSPARENT_PNG.length,
-      "Cache-Control": "no-store, no-cache, must-revalidate",
-      Pragma: "no-cache",
-      Expires: "0",
+      "Content-Type":   "image/png",
+      "Content-Length": PIXEL_PNG.length,
+      "Cache-Control":  "no-store, no-cache, must-revalidate, proxy-revalidate",
+      "Pragma":         "no-cache",
+      "Expires":        "0",
     });
-    return res.end(TRANSPARENT_PNG);
+    return res.end(PIXEL_PNG);
   }
 
-  // ── REST API ────────────────────────────────────────────────────────────
+  // ── 2. REST API ─────────────────────────────────────────────────────────
   if (pathname.startsWith("/api/")) {
     let body = "";
-    req.on("data", (d) => (body += d));
+    req.on("data", (chunk) => { if (body.length < 65536) body += chunk; });
     req.on("end", () => {
       let data = {};
-      try { data = body ? JSON.parse(body) : {}; } catch {}
+      try { if (body) data = JSON.parse(body); } catch {
+        res.writeHead(400, { "Content-Type": "application/json" });
+        return res.end(JSON.stringify({ error: "Invalid JSON" }));
+      }
 
-      // Auth helper
-      const authHeader = req.headers["authorization"] || "";
-      const [authUser, authToken] = (authHeader.replace("Bearer ", "")).split(":");
+      const [authUser, authToken] = parseAuth(req.headers["authorization"]);
 
       function authed() {
-        if (!authUser || !authToken || !verifyToken(authUser, authToken)) {
+        if (!verifyToken(authUser, authToken)) {
           res.writeHead(401, { "Content-Type": "application/json" });
           res.end(JSON.stringify({ error: "Unauthorized" }));
           return false;
@@ -201,18 +220,18 @@ const server = http.createServer((req, res) => {
       }
 
       function json(status, obj) {
-        res.writeHead(status, { "Content-Type": "application/json" });
-        res.end(JSON.stringify(obj));
+        const body = JSON.stringify(obj);
+        res.writeHead(status, {
+          "Content-Type":   "application/json",
+          "Content-Length": Buffer.byteLength(body),
+        });
+        res.end(body);
       }
 
       const route = `${req.method} ${pathname}`;
 
-      if (route === "POST /api/register") {
-        return json(200, registerUser(data.username, data.password));
-      }
-      if (route === "POST /api/login") {
-        return json(200, loginUser(data.username, data.password));
-      }
+      if (route === "POST /api/register") return json(200, registerUser(data.username, data.password));
+      if (route === "POST /api/login")    return json(200, loginUser(data.username, data.password));
       if (route === "POST /api/logout") {
         if (!authed()) return;
         logoutUser(authUser, authToken);
@@ -230,74 +249,92 @@ const server = http.createServer((req, res) => {
         if (!authed()) return;
         return json(200, deletePixel(authUser, data.pixelId));
       }
-
       return json(404, { error: "Not found" });
     });
     return;
   }
 
-  // ── Service Worker (must be at root scope, served before generic static) ──
+  // ── 3. Service worker (root scope, no-cache) ────────────────────────────
   if (pathname === "/sw.js") {
-    const swPath = path.join(__dirname, "sw.js");
-    if (fs.existsSync(swPath)) {
+    const swPath = path.join(ROOT, "sw.js");
+    if (!fs.existsSync(swPath)) { res.writeHead(404); return res.end(); }
+    const content = fs.readFileSync(swPath);
+    res.writeHead(200, {
+      "Content-Type":          "application/javascript; charset=utf-8",
+      "Content-Length":        content.length,
+      "Service-Worker-Allowed": "/",
+      "Cache-Control":         "no-store",
+    });
+    return res.end(content);
+  }
+
+  // ── 4. Static files from /public (path-traversal safe) ──────────────────
+  {
+    const rel      = pathname === "/" ? "/index.html" : pathname;
+    const resolved = path.resolve(PUBLIC_DIR, "." + rel);
+
+    // Reject anything that escapes /public
+    if (!resolved.startsWith(PUBLIC_DIR + path.sep) && resolved !== PUBLIC_DIR) {
+      res.writeHead(403); return res.end("Forbidden");
+    }
+
+    const ext = path.extname(resolved).toLowerCase();
+    if (MIME[ext] && fs.existsSync(resolved)) {
+      const content = fs.readFileSync(resolved);
       res.writeHead(200, {
-        "Content-Type": "application/javascript",
-        "Service-Worker-Allowed": "/",
-        "Cache-Control": "no-store",
+        "Content-Type":   MIME[ext],
+        "Content-Length": content.length,
+        "Cache-Control":  "public, max-age=300",
       });
-      return res.end(fs.readFileSync(swPath));
+      return res.end(content);
     }
   }
 
-  // ── Static files (served from /public) ──────────────────────────────────
-  let filePath = pathname === "/" ? "/index.html" : pathname;
-  filePath = path.join(__dirname, "public", filePath);
-  const ext = path.extname(filePath);
-  if (STATIC[ext] && fs.existsSync(filePath)) {
-    res.writeHead(200, { "Content-Type": STATIC[ext] });
-    return res.end(fs.readFileSync(filePath));
-  }
-
-  res.writeHead(404);
-  res.end("Not found");
+  res.writeHead(404); res.end("Not found");
 });
 
-// ── WebSocket Server ───────────────────────────────────────────────────────
+// ── WebSocket server ───────────────────────────────────────────────────────
 const wss = new WebSocketServer({ server });
 
 wss.on("connection", (ws, req) => {
-  const parsed = url.parse(req.url, true);
-  // token param is "username:rawtoken"
-  const rawParam = parsed.query.token || "";
-  const colonIdx = rawParam.indexOf(":");
-  const wsUsername = colonIdx > -1 ? rawParam.slice(0, colonIdx) : "";
-  const wsToken    = colonIdx > -1 ? rawParam.slice(colonIdx + 1) : rawParam;
-  const type = parsed.query.type || "browser"; // "browser" | "sw"
-
-  if (!wsUsername || !verifyToken(wsUsername, wsToken)) {
+  let parsed;
+  try { parsed = new URL(req.url, "http://x"); }
+  catch { return ws.close(4000, "Bad URL"); }
+
+  // token param: "username:rawtoken"
+  const rawToken = parsed.searchParams.get("token") || "";
+  const colonIdx = rawToken.indexOf(":");
+  const wsUsername = colonIdx > 0 ? rawToken.slice(0, colonIdx) : "";
+  const wsToken    = colonIdx > 0 ? rawToken.slice(colonIdx + 1) : "";
+
+  if (!verifyToken(wsUsername, wsToken)) {
     ws.close(4001, "Unauthorized");
     return;
   }
-  const user = getAccount(wsUsername);
 
-  const username = user.username; // same as wsUsername
-  addConnection(username, ws, type);
+  addConn(wsUsername, ws);
+  ws.send(JSON.stringify({ type: "connected", username: wsUsername }));
 
-  // Confirm connection
-  ws.send(JSON.stringify({ type: "connected", username }));
+  // Server-side ping to detect dead connections (every 30s)
+  const keepalive = setInterval(() => {
+    if (ws.readyState === 1) ws.ping();
+    else { clearInterval(keepalive); }
+  }, 30000);
 
   ws.on("message", (raw) => {
     try {
       const msg = JSON.parse(raw);
-      // Client heartbeat ping — reply with pong so client knows connection is live
+      // Client heartbeat — reply so client can verify the connection is live
       if (msg.type === "ping") ws.send(JSON.stringify({ type: "pong" }));
     } catch {}
   });
 
-  ws.on("close", () => removeConnection(username, ws));
-  ws.on("error", () => removeConnection(username, ws));
+  ws.on("pong", () => {}); // keep-alive acknowledged
+
+  ws.on("close", () => { clearInterval(keepalive); removeConn(wsUsername, ws); });
+  ws.on("error", (e) => { console.error("[WS] error", e.message); clearInterval(keepalive); removeConn(wsUsername, ws); });
 });
 
 // ── Start ──────────────────────────────────────────────────────────────────
-const PORT = process.env.PORT || 7860;
-server.listen(PORT, () => console.log(`PixelBeacon running on :${PORT}`));
+const PORT = parseInt(process.env.PORT || "7860", 10);
+server.listen(PORT, "0.0.0.0", () => console.log(`PixelBeacon listening on :${PORT}`));