From 60c51f9a34da39fde79abbd81f48afd3eb35f7e2 Mon Sep 17 00:00:00 2001
From: Yeoh Joer <yeoh.joer@outlook.com>
Date: Tue, 27 Feb 2024 19:28:13 +0800
Subject: [PATCH] Terminate the program when the maximum ping difference is
 exceeded

---
 bin/smee.js | 12 ++++++------
 index.ts    | 28 ++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/bin/smee.js b/bin/smee.js
index a704acc..c44b9d6 100755
--- a/bin/smee.js
+++ b/bin/smee.js
@@ -10,6 +10,8 @@ program
   .usage('[options]')
   .option('-u, --url <url>', 'URL of the webhook proxy service. Default: https://smee.io/new')
   .option('-t, --target <target>', 'Full URL (including protocol and path) of the target service the events will forwarded to. Default: http://127.0.0.1:PORT/PATH')
+  .option('-h, --healthcheck <interval>', 'Perform health checks based on received pings at specified intervals (in seconds)')
+  .option('-m, --max-ping-difference <seconds>', 'The maximum difference between the last ping and the current time (in seconds) before the client is considered unhealthy. Default: 60', 60)
   .option('-p, --port <n>', 'Local HTTP server port', process.env.PORT || 3000)
   .option('-P, --path <path>', 'URL path to post proxied requests to`', '/')
   .parse(process.argv)
@@ -21,13 +23,11 @@ const {
 } = opts
 
 async function setup () {
-  let source = opts.url
+  const source = opts.url || await Client.createChannel()
+  const healthcheck = Number.parseInt(opts.healthcheck, 10)
+  const maxPingDifference = Number.parseInt(opts.maxPingDifference, 10)
 
-  if (!source) {
-    source = await Client.createChannel()
-  }
-
-  const client = new Client({ source, target })
+  const client = new Client({ source, target, healthcheck, maxPingDifference })
   client.start()
 }
 
diff --git a/index.ts b/index.ts
index 5016592..5bd832f 100644
--- a/index.ts
+++ b/index.ts
@@ -8,6 +8,8 @@ type Severity = "info" | "error";
 interface Options {
   source: string;
   target: string;
+  healthcheck: number;
+  maxPingDifference: number;
   logger?: Pick<Console, Severity>;
   fetch?: any;
 }
@@ -15,6 +17,9 @@ interface Options {
 class Client {
   source: string;
   target: string;
+  healthcheck: number;
+  maxPingDifference: number;
+  lastPing: number;
   fetch: typeof global.fetch;
   logger: Pick<Console, Severity>;
   events!: EventSource;
@@ -22,11 +27,16 @@ class Client {
   constructor({
     source,
     target,
+    healthcheck,
+    maxPingDifference,
     logger = console,
     fetch = global.fetch,
   }: Options) {
     this.source = source;
     this.target = target;
+    this.healthcheck = healthcheck;
+    this.maxPingDifference = maxPingDifference;
+    this.lastPing = Date.now();
     this.logger = logger!;
     this.fetch = fetch;
 
@@ -85,6 +95,11 @@ class Client {
     this.logger.info("Connected", this.events.url);
   }
 
+  onping() {
+    this.logger.info(`Received a ping on ${new Date().toISOString()}`);
+    this.lastPing = Date.now();
+  }
+
   onerror(err: any) {
     this.logger.error(err);
   }
@@ -99,6 +114,19 @@ class Client {
     events.addEventListener("open", this.onopen.bind(this));
     events.addEventListener("error", this.onerror.bind(this));
 
+    if (this.healthcheck) {
+      events.addEventListener("ping", this.onping.bind(this));
+
+      setInterval(() => {
+        const difference = (Date.now() - this.lastPing) / 1000;
+
+        if (difference > this.maxPingDifference) {
+          this.logger.error(`Maximum ping difference exceeded. (Difference: ${difference.toFixed(4)}s, Maximum Allowed: ${this.maxPingDifference}s)`);
+          process.exit(1);
+        }
+      }, this.healthcheck * 1000);
+    }
+
     this.logger.info(`Forwarding ${this.source} to ${this.target}`);
     this.events = events;