who_need_help/load/k6/public-and-websocket.js

import { check, sleep } from "k6";
import http from "k6/http";
import exec from "k6/execution";
import { Counter } from "k6/metrics";
import { WebSocket } from "k6/websockets";

function required(name) {
  const value = __ENV[name];

  if (value === undefined || value === "") {
    throw new Error(`${name} is required`);
  }

  return value;
}

function positiveInteger(name) {
  const value = Number(required(name));

  if (!Number.isInteger(value) || value <= 0) {
    throw new Error(`${name} must be a positive integer`);
  }

  return value;
}

function nonNegativeNumber(name) {
  const value = Number(required(name));

  if (!Number.isFinite(value) || value < 0) {
    throw new Error(`${name} must be a non-negative number`);
  }

  return value;
}

const baseUrl = required("BASE_URL").replace(/\/+$/, "");
const publicOrigin = required("PUBLIC_ORIGIN").replace(/\/+$/, "");

if (!/^https?:\/\/[^/?#]+$/.test(baseUrl)) {
  throw new Error("BASE_URL must be an HTTP(S) origin without a path");
}

if (!/^https?:\/\/[^/?#]+$/.test(publicOrigin)) {
  throw new Error("PUBLIC_ORIGIN must be an HTTP(S) origin without a path");
}

const websocketUrl =
  baseUrl.replace(/^http/, "ws") + "/live/websocket?vsn=2.0.0";
const websocketHoldMs = positiveInteger("WS_HOLD_MS");
const websocketConnectTimeoutMs = positiveInteger("WS_CONNECT_TIMEOUT_MS");
const httpThinkSeconds = nonNegativeNumber("HTTP_THINK_SECONDS");
const httpVus = positiveInteger("HTTP_VUS");
const websocketVus = positiveInteger("WS_VUS");
const authenticatedVus = positiveInteger("AUTH_VUS");
const maxVus = httpVus + websocketVus + authenticatedVus;
const authWebsocketTimeoutMs = positiveInteger("AUTH_WS_TIMEOUT_MS");
const authThinkSeconds = nonNegativeNumber("AUTH_THINK_SECONDS");
const fixturePassword = required("FIXTURE_PASSWORD");
const fixtureManifest = JSON.parse(open("/fixtures/fixtures.json"));
const fixtureRunId = required("FIXTURE_RUN_ID");

if (
  fixtureManifest.schema_version !== 1 ||
  fixtureManifest.run_id !== fixtureRunId ||
  !Array.isArray(fixtureManifest.fixtures) ||
  fixtureManifest.fixtures.length < maxVus
) {
  throw new Error("The authenticated-load fixture manifest is invalid or too small");
}

const websocketOpened = new Counter("wnh_websocket_opened");
const websocketErrors = new Counter("wnh_websocket_errors");
const heartbeatReplies = new Counter("wnh_websocket_heartbeat_replies");
const authenticatedLogins = new Counter("wnh_authenticated_logins");
const authenticatedPages = new Counter("wnh_authenticated_pages");
const trackingUpdates = new Counter("wnh_tracking_updates");
const liveViewJoins = new Counter("wnh_liveview_joins");
const liveViewTrackingStarts = new Counter("wnh_liveview_tracking_starts");
const liveViewTrackingStops = new Counter("wnh_liveview_tracking_stops");
const liveViewMessages = new Counter("wnh_liveview_messages");
const authenticatedErrors = new Counter("wnh_authenticated_errors");

let authenticated = false;
let fixture;

export const options = {
  discardResponseBodies: true,
  // The isolated Traefik profile uses its generated local-only certificate.
  // Transport remains TLS; trust verification is scoped to this k6 process.
  insecureSkipTLSVerify: true,
  // Each authenticated VU represents one continuing user session.
  noCookiesReset: true,
  scenarios: {
    public_http: {
      executor: "constant-vus",
      exec: "publicHttp",
      vus: httpVus,
      duration: required("DURATION"),
    },
    phoenix_websocket: {
      executor: "constant-vus",
      exec: "phoenixWebsocket",
      vus: websocketVus,
      duration: required("DURATION"),
    },
    authenticated_mutual_aid: {
      executor: "constant-vus",
      exec: "authenticatedMutualAid",
      vus: authenticatedVus,
      duration: required("DURATION"),
    },
  },
};

export function publicHttp() {
  const responses = http.batch([
    {
      method: "GET",
      url: `${baseUrl}/`,
      tags: { endpoint: "home" },
    },
    {
      method: "GET",
      url: `${baseUrl}/safety`,
      tags: { endpoint: "safety" },
    },
    {
      method: "GET",
      url: `${baseUrl}/healthz/ready`,
      tags: { endpoint: "readiness" },
    },
  ]);

  check(responses[0], { "home returned 200": (response) => response.status === 200 });
  check(responses[1], { "safety returned 200": (response) => response.status === 200 });
  check(responses[2], {
    "readiness returned 200": (response) => response.status === 200,
  });

  sleep(httpThinkSeconds);
}

export function phoenixWebsocket() {
  const socket = new WebSocket(websocketUrl, [], {
    tags: { endpoint: "phoenix_live_socket" },
  });
  const ref = `${__VU}-${__ITER}`;
  let opened = false;
  let closeTimer;
  const connectTimer = setTimeout(() => {
    if (!opened) {
      websocketErrors.add(1);
      socket.close();
    }
  }, websocketConnectTimeoutMs);

  socket.addEventListener("open", () => {
    opened = true;
    clearTimeout(connectTimer);
    websocketOpened.add(1);
    socket.send(JSON.stringify([null, ref, "phoenix", "heartbeat", {}]));
    closeTimer = setTimeout(() => socket.close(), websocketHoldMs);
  });

  socket.addEventListener("message", (event) => {
    let frame;

    try {
      frame = JSON.parse(event.data);
    } catch (_error) {
      return;
    }

    if (
      Array.isArray(frame) &&
      frame[1] === ref &&
      frame[2] === "phoenix" &&
      frame[3] === "phx_reply" &&
      frame[4] &&
      frame[4].status === "ok"
    ) {
      heartbeatReplies.add(1);
    }
  });

  socket.addEventListener("error", () => {
    websocketErrors.add(1);
    clearTimeout(connectTimer);
    socket.close();
  });

  socket.addEventListener("close", () => {
    clearTimeout(connectTimer);

    if (closeTimer !== undefined) {
      clearTimeout(closeTimer);
    }
  });
}

function htmlAttribute(tag, name) {
  const escaped = name.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
  const match = tag.match(new RegExp(`\\s${escaped}=(?:"([^"]*)"|'([^']*)')`));

  if (!match) {
    return null;
  }

  return (match[1] ?? match[2])
    .replaceAll("&amp;", "&")
    .replaceAll("&quot;", '"')
    .replaceAll("&#39;", "'");
}

function csrfToken(body) {
  const metaTags = body.match(/<meta\b[^>]*>/gi) || [];
  const tag = metaTags.find((candidate) => htmlAttribute(candidate, "name") === "csrf-token");
  return tag ? htmlAttribute(tag, "content") : null;
}

function liveViewRoot(body) {
  const tags = body.match(/<[^!/][^>]*>/g) || [];
  const tag = tags.find((candidate) => /\sdata-phx-main(?:\s|=|>)/.test(candidate));

  if (!tag) {
    return null;
  }

  const root = {
    id: htmlAttribute(tag, "id"),
    session: htmlAttribute(tag, "data-phx-session"),
    static: htmlAttribute(tag, "data-phx-static"),
  };

  return root.id && root.session ? root : null;
}

function fixtureForVu() {
  const index = exec.vu.idInInstance - 1;
  const selected = fixtureManifest.fixtures[index];

  if (!selected) {
    throw new Error(`No fixture exists for VU ${exec.vu.idInInstance}`);
  }

  return selected;
}

function authenticateActor(actor) {
  const loginPage = http.get(`${baseUrl}/users/log-in`, {
    responseType: "text",
    tags: { endpoint: "auth_login_page" },
  });
  const csrf = typeof loginPage.body === "string" ? csrfToken(loginPage.body) : null;

  if (loginPage.status !== 200 || !csrf) {
    authenticatedErrors.add(1);
    return false;
  }

  const login = http.post(
    `${baseUrl}/users/log-in`,
    {
      _csrf_token: csrf,
      "user[email]": actor.helper_email,
      "user[password]": fixturePassword,
    },
    {
      redirects: 0,
      tags: { endpoint: "auth_password_login" },
    },
  );
  const jar = http.cookieJar();
  const cookies = jar.cookiesForURL(baseUrl);
  const succeeded =
    login.status === 302 &&
    Array.isArray(cookies._who_need_help_key) &&
    cookies._who_need_help_key.length > 0;

  if (succeeded) {
    authenticatedLogins.add(1);
  } else {
    authenticatedErrors.add(1);
  }

  return succeeded;
}

export function authenticatedMutualAid() {
  if (!fixture) {
    fixture = fixtureForVu();
  }

  if (!authenticated) {
    authenticated = authenticateActor(fixture);
  }

  if (!authenticated) {
    sleep(authThinkSeconds);
    return;
  }

  const requestPath = `/requests/${fixture.request_id}`;
  const page = http.get(`${baseUrl}${requestPath}`, {
    responseType: "text",
    tags: { endpoint: "authenticated_request_page" },
  });
  const body = typeof page.body === "string" ? page.body : "";
  const csrf = csrfToken(body);
  const root = liveViewRoot(body);

  if (page.status !== 200 || !csrf || !root) {
    authenticatedErrors.add(1);
    sleep(authThinkSeconds);
    return;
  }

  authenticatedPages.add(1);

  const socketUrl =
    baseUrl.replace(/^http/, "ws") +
    `/live/websocket?_csrf_token=${encodeURIComponent(csrf)}` +
    "&client_type=browser&vsn=2.0.0";
  const socket = new WebSocket(socketUrl, [], {
    jar: http.cookieJar(),
    headers: { Origin: publicOrigin },
    tags: { endpoint: "authenticated_liveview" },
  });
  const joinRef = `${exec.vu.idInInstance}-${__ITER}-join`;
  const trackingRef = `${exec.vu.idInInstance}-${__ITER}-tracking`;
  const eventRef = `${exec.vu.idInInstance}-${__ITER}-message`;
  const stopRef = `${exec.vu.idInInstance}-${__ITER}-stop`;
  const topic = `lv:${root.id}`;
  let completed = false;
  const timer = setTimeout(() => {
    if (!completed) {
      authenticatedErrors.add(1);
    }
    socket.close();
  }, authWebsocketTimeoutMs);

  socket.addEventListener("open", () => {
    socket.send(
      JSON.stringify([
        joinRef,
        joinRef,
        topic,
        "phx_join",
        {
          url: `${publicOrigin}${requestPath}`,
          params: {
            _csrf_token: csrf,
            client_type: "browser",
            _mounts: 0,
            _mount_attempts: 0,
            _track_static: [],
          },
          session: root.session,
          static: root.static || null,
        },
      ]),
    );
  });

  socket.addEventListener("message", (event) => {
    let frame;

    try {
      frame = JSON.parse(event.data);
    } catch (_error) {
      return;
    }

    if (!Array.isArray(frame) || frame[2] !== topic || frame[3] !== "phx_reply") {
      return;
    }

    if (frame[1] === joinRef && frame[4]?.status === "ok") {
      liveViewJoins.add(1);
      socket.send(
        JSON.stringify([
          joinRef,
          trackingRef,
          topic,
          "event",
          {
            type: "click",
            event: "start-tracking",
            value: {},
            cid: null,
          },
        ]),
      );
    } else if (frame[1] === trackingRef && frame[4]?.status === "ok") {
      liveViewTrackingStarts.add(1);
      const trackingStartResponse = frame[4]?.response ?? null;
      const latitude = fixture.latitude + ((__ITER % 20) + 1) / 1_000_000;
      const longitude = fixture.longitude + ((__ITER % 20) + 1) / 1_000_000;
      const tracking = http.post(
        `${baseUrl}/mobile/tracking/${fixture.assignment_id}/position`,
        {
          latitude: String(latitude),
          longitude: String(longitude),
          accuracy_meters: "5",
          captured_at: new Date().toISOString(),
        },
        {
          headers: { "x-csrf-token": csrf },
          responseType: "text",
          tags: { endpoint: "authenticated_tracking_update" },
        },
      );

      if (tracking.status !== 204) {
        console.error(
          `authenticated tracking update failed: ${JSON.stringify({
            vu_id_in_instance: exec.vu.idInInstance,
            vu_id_in_test: exec.vu.idInTest,
            fixture_index: fixture.index,
            iteration: __ITER,
            status: tracking.status,
            response_body: tracking.body,
            tracking_start_response: trackingStartResponse,
          })}`,
        );
        completed = true;
        authenticatedErrors.add(1);
        clearTimeout(timer);
        socket.close();
        return;
      }

      trackingUpdates.add(1);
      const message =
        `load:${fixtureRunId}:vu-${exec.vu.idInInstance}:iteration-${__ITER}`;
      const formValue = `message%5Bbody%5D=${encodeURIComponent(message)}`;

      socket.send(
        JSON.stringify([
          joinRef,
          eventRef,
          topic,
          "event",
          {
            type: "form",
            event: "send-message",
            value: formValue,
            meta: { _target: "undefined" },
            uploads: {},
            cid: null,
          },
        ]),
      );
    } else if (frame[1] === eventRef && frame[4]?.status === "ok") {
      liveViewMessages.add(1);
      socket.send(
        JSON.stringify([
          joinRef,
          stopRef,
          topic,
          "event",
          {
            type: "click",
            event: "stop-tracking",
            value: {},
            cid: null,
          },
        ]),
      );
    } else if (frame[1] === stopRef && frame[4]?.status === "ok") {
      completed = true;
      liveViewTrackingStops.add(1);
      clearTimeout(timer);
      socket.close();
    } else if (
      frame[1] === joinRef ||
      frame[1] === trackingRef ||
      frame[1] === eventRef ||
      frame[1] === stopRef
    ) {
      completed = true;
      authenticatedErrors.add(1);
      clearTimeout(timer);
      socket.close();
    }
  });

  socket.addEventListener("error", () => {
    if (!completed) {
      completed = true;
      authenticatedErrors.add(1);
    }
    clearTimeout(timer);
    socket.close();
  });

  socket.addEventListener("close", () => clearTimeout(timer));
  sleep(authThinkSeconds);
}