Poor man's Disaster Recovery script

disaster-recovery.ts

#! /usr/bin/env node

/// <reference lib="esnext" />
/// <reference lib="webworker" />

const slackAPIKey = process.env.SLACK_API_KEY;
const slackChannel = process.env.SLACK_CHANNEL;
const cloudflareAPIKey = process.env.CLOUDFLARE_API_KEY;
const cloudflareZoneId = process.env.CLOUDFLARE_ZONE_ID;
const healthcheckPath = process.env.HEALTHCHECK_PATH ?? '/';
const healthcheckProtocol = process.env.HEALTHCHECK_PROTO ?? 'http';
const healthcheckInterval = parseInt(process.env.HEALTHCHECK_INTERVAL ?? '30', 10) * 1000;
const refreshInterval = parseInt(process.env.REFRESH_INTERVAL ?? '3600', 10) * 1000;
const verbosity = parseInt(process.env.VERBOSITY ?? '1', 10);

/*
 * Utilities
 * ====
 */

const log = verbosity >= 1 ? console.log : () => {};
const debug = verbosity >= 2 ? console.log : () => {};
const sleep = (timeout: number) =>
  new Promise<void>(resolve => setTimeout(resolve, timeout));

const wrapPromise = <T,>(promise: Promise<T>): Promise<PromiseSettledResult<T>> =>
  promise.then(value => ({ status: 'fulfilled', value }), reason => ({ status: 'rejected', reason }));

const runWithInterval = (
  fn: () => Promise<void>,
  opts: { interval: number, onFailure?: (err: unknown) => void }
) => {
  let isCanceled = false;
  let lastTimeout: ReturnType<typeof setTimeout> | null = null;
  const run = async () => {
    await fn().catch(err => opts.onFailure?.(err as unknown));
    if (!isCanceled) {
      lastTimeout = setTimeout(run, opts.interval);
    }
  };

  run();

  return () => {
    isCanceled = true;
    if (lastTimeout !== null) {
      clearTimeout(lastTimeout);
    }
  };
};

const exponentialBackoff = <T extends (...args: never[]) => unknown>(
  fn: T,
  opts = { initInterval: 100, maxInterval: 60 * 1000, maxRetry: 5 }
) => async (...args: Parameters<T>): Promise<Awaited<ReturnType<T>>> => {
  let retry = 0;
  let interval = opts.initInterval;

  while (true) {
    try {
      const result = await fn(...args) as Awaited<ReturnType<T>>;
      return result;
    } catch(e) {
      await sleep(interval);
      interval *= 2;
      retry++;

      if (retry > opts.maxRetry) {
        throw e;
      }
    }
  }
};

const createPromisePool = async <T, >(
  generator: Generator<Promise<T>, void>,
  concurrency = 5
) => {
  const output: PromiseSettledResult<T>[] = [];
  const worker = async () => {
    for (const promise of generator) {
      output.push(await wrapPromise(promise));
    }
  };

  await Promise.all(Array.from({ length: concurrency }).map(worker));
  return output;
};

const createFetchAPI =
  (
    basePath: string,
    { headers: baseHeaders, handleResponse = (res) => res.json() }: {
      headers?: Record<string, string>;
      handleResponse?: (res: Response) => Promise<unknown>,
    }
  ) =>
  <T, >(
    path: string,
    { headers, body, ...opts }: Omit<RequestInit, 'headers' | 'body'> & {
      headers?: Record<string, string>;
      body?: unknown;
    }
  ) =>
    fetch(new URL(path.replace(/^\//, './'), basePath.replace(/\/?$/, '/')), {
      ...opts,
      ...(!!body && ({ body: typeof body === 'object' ? JSON.stringify(body) : (body as BodyInit) })),
      headers: {
        ...(!!body && typeof body === 'object' && { 'Content-Type': 'application/json' }),
        ...baseHeaders,
        ...headers,
      },
    })
    .then(res => handleResponse(res) as Promise<T>);

/*
 * API Requests
 * ====
 */

const fetchCloudflareAPI = createFetchAPI(
  'https://api.cloudflare.com/client/v4',
  {
    headers: { Authorization: `Bearer ${cloudflareAPIKey}` },
    handleResponse: async (res) => {
      const data = (await res.json()) as CloudflareResponse<unknown>;
      if (!data.success) {
        const errors = data.errors.map(err => err.message).join('\n');
        throw new Error(`Failed to fetch Cloudflare API: ${errors}`);
      }

      return data;
    }
  }
);

const fetchSlackAPI = createFetchAPI(
  'https://slack.com/api',
  { headers: { Authorization: `Bearer ${slackAPIKey}` } },
);

const fetchDNSRecords = async function*() {
  let page = 1;

  while (true) {
    const res = await fetchCloudflareAPI<CloudflareListDNSRecordsResponse>(
      `/zones/${cloudflareZoneId}/dns_records?page=${page}`,
      { method: 'GET' },
    );

    yield res.result;
    page++;

    if (page * res.result_info.per_page >= res.result_info.count) {
      return [];
    }
  }
};

const updateDNSRecord = async (state: RecordState, content: string) => {
  await fetchCloudflareAPI<CloudflareUpdateDNSRecordsReponse>(
    `/zones/${cloudflareZoneId}/dns_records/${state.id}`,
    { method: 'PATCH', body: { content } },
  );

  state.active = content;
};

const sendSlackMessage = async (channel: string, text: string) => {
  return fetchSlackAPI(
    '/chat.postMessage',
    { method: 'POST', body: { channel, text } },
  );
};

/*
 * Main Logic
 * ====
 */

const report = async (message: string) => {
  log(message);

  if (slackAPIKey && slackChannel) {
    return sendSlackMessage(slackChannel, `[Disaster Recovery] ${message}`)
      .catch(() => {});
  }
};

const testEndpoint = async (state: RecordState) => {
  for (const endpoint of state.endpoints) {
    const { status } = await wrapPromise(exponentialBackoff(async () => {
      debug(`Fetching ${endpoint}`);

      const response = await fetch(new URL(healthcheckPath, `${healthcheckProtocol}://${endpoint}`));
      if (!response.ok) {
        throw new Error(`Failed to fetch with status ${response.status} ${response.statusText}`);
      }
    }, { initInterval: 250, maxInterval: 500, maxRetry: 2 })());

    if (status === 'fulfilled') {
      return endpoint;
    }
  }

  return null;
};

const testAndUpdateEndpoint = async (state: RecordState) => {
  const firstLivingEndpoints = await testEndpoint(state);
  log(`Living Endpoint for ${state.name} -> ${firstLivingEndpoints}`);

  if (firstLivingEndpoints && state.active !== firstLivingEndpoints) {
    await exponentialBackoff(
      () => updateDNSRecord(state, firstLivingEndpoints),
      { initInterval: 100, maxInterval: 1000, maxRetry: 3 }
    )().then(
      () => report(`Record ${state.name}: Updated to ${state.active}`),
      () => report(`Record ${state.name}: Failed to update ${state.active}`)
    );
  }
};

(() => {
  let records = [] as RecordState[];
  let lastError = 0;

  runWithInterval(exponentialBackoff(async () => {
    records = (await Array.fromAsync(fetchDNSRecords()))
      .flat()
      .map(record => {
        // Only Supports IPv4
        const endpoints = record.comment
          ?.split(/\s+/)
          .find(tag => /^dr:([0-9.,])+$/.test(tag))
          ?.replace(/^dr:/, '')
          .split(',')
          .filter(endpoint => /(?:[0-9]{1,3}\.){3}[0-9]{1,3}/.test(endpoint));

        if (!endpoints?.[1] || record.type !== 'A') {
          return null;
        }

        const { id, name, content: active } = record;
        return { id, name, active, endpoints };
      })
      .filter(<T,>(x: T | null): x is T => !!x);

    debug('Active records:');
    debug(records.map(state => `${state.name} -> ${state.active} (${state.endpoints.join(',')})`).join('\n'));
  }), {
    interval: refreshInterval,
    onFailure: (e) => {
      if (Date.now() - lastError >= 4 * 3600 * 1000) {
        lastError = Date.now();
        report('Failed to fetch records' + (e instanceof Error ? `: ${e.message}` : ''));
      }
    }
  });

  runWithInterval(async () => {
    await createPromisePool((function* () {
      for (const state of records) {
        yield testAndUpdateEndpoint(state);
      }
    })());
  }, { interval: healthcheckInterval });
})();

/*
 * Types
 * ====
 */
declare const process: { env: Record<string, string> };

type RecordState = {
  id: string;
  name: string;
  active: string;
  endpoints: string[];
};

type CloudflareError = { code: number, message: string };
type CloudflareResponse<T> = {
  success: boolean;
  messages: CloudflareError[];
  errors: CloudflareError[];
  result: T;
};

type CloudflarePaginatedResponse<T> = CloudflareResponse<T[]> & {
  result_info: {
    count: number;
    page: number;
    per_page: number;
    total_count: number;
  };
};

type CloudflareDNSRecord = {
  id: string;
  type: string;
  name: string;
  content: string;
  comment: string | null;
};

type CloudflareListDNSRecordsResponse = CloudflarePaginatedResponse<CloudflareDNSRecord>;
type CloudflareUpdateDNSRecordsReponse = CloudflareResponse<unknown>;