From 7e694fd9b7281234c83b880af76a51f782e8c8c9 Mon Sep 17 00:00:00 2001
From: lyc8503 <admin@lyc8503.site>
Date: Mon, 23 Oct 2023 15:29:27 +0800
Subject: [PATCH] add advanced checks

---
 components/DetailChart.tsx |  3 ++-
 pages/index.tsx            |  2 --
 uptime.config.ts           | 24 +++++++++++++++++++++++-
 uptime.types.ts            |  9 +++++----
 worker/src/index.ts        | 29 +++++++++++++++++++++++++++--
 5 files changed, 57 insertions(+), 10 deletions(-)

diff --git a/components/DetailChart.tsx b/components/DetailChart.tsx
index 5f35d83..869fc6e 100644
--- a/components/DetailChart.tsx
+++ b/components/DetailChart.tsx
@@ -2,6 +2,7 @@ import { Line } from "react-chartjs-2"
 import { Chart as ChartJS, CategoryScale, LinearScale, PointElement, LineElement, Title, Tooltip as ChartTooltip, Legend, TimeScale } from 'chart.js'
 import 'chartjs-adapter-moment'
 import { MonitorState, MonitorTarget } from "@/uptime.types"
+import { iataToCountry } from "@/util/iata"
 
 ChartJS.register(
   CategoryScale,
@@ -47,7 +48,7 @@ export default function DetailChart({ monitor, state }: { monitor: MonitorTarget
         callbacks: {
           label: (item: any) => {
             if (item.parsed.y) {
-              return `${item.parsed.y}ms (${item.raw.loc})`
+              return `${item.parsed.y}ms (${iataToCountry(item.raw.loc)})`
             }
           }
         }
diff --git a/pages/index.tsx b/pages/index.tsx
index b5cd554..cbbf16e 100644
--- a/pages/index.tsx
+++ b/pages/index.tsx
@@ -1,5 +1,3 @@
-'use client'
-
 import Head from 'next/head'
 
 import { Inter } from 'next/font/google'
diff --git a/uptime.config.ts b/uptime.config.ts
index be3d62d..e7f050b 100644
--- a/uptime.config.ts
+++ b/uptime.config.ts
@@ -63,7 +63,8 @@ const config = {
       timeout: 10000,
       headers: {
         "Cookie": "type=tunnel;"
-      }
+      },
+      responseKeyword: 'Hello'
     },
     {
       id: 'broken-test',
@@ -71,6 +72,27 @@ const config = {
       method: 'TCP_PING',
       target: '1.1.1.1:1234'
     },
+    {
+      id: '404test',
+      name: '404-test',
+      method: 'GET',
+      target: 'https://www.baidu.com/404',
+    },
+    {
+      id: '404test2',
+      name: '404-test2',
+      method: 'GET',
+      target: 'https://www.baidu.com/404',
+      expectedCodes: [404, 405]
+    },
+    {
+      id: '404test3',
+      name: '404-test3',
+      method: 'GET',
+      target: 'https://www.baidu.com/404',
+      expectedCodes: [404],
+      responseKeyword: 'Hello'
+    }
   ]
 } 
 
diff --git a/uptime.types.ts b/uptime.types.ts
index d73060b..780e5c4 100644
--- a/uptime.types.ts
+++ b/uptime.types.ts
@@ -22,17 +22,18 @@ type MonitorState = {
   }>
 }
 
-// TODO
 type MonitorTarget = {
   id: string,
   name: string,
   method: string,  // "TCP_PING" or Http Method (e.g. GET, POST, OPTIONS, etc.)
   target: string,  // url for http, hostname:port for tcp
 
-  expectedCode?: number[],
+  // HTTP Code
+  expectedCodes?: number[],
   timeout?: number,
-  headers?: Record<string, string>,
-  body?: string
+  headers?: Record<string, string | undefined>,
+  body?: BodyInit,
+  responseKeyword?: string
 }
 
 
diff --git a/worker/src/index.ts b/worker/src/index.ts
index 82f834b..00154b2 100644
--- a/worker/src/index.ts
+++ b/worker/src/index.ts
@@ -28,6 +28,7 @@ async function getStatus(monitor: MonitorTarget): Promise<{ ping: number; up: bo
 			const reader = socket.readable.getReader()
 			const writer = socket.writable.getWriter()
 			await writer.write(new TextEncoder().encode("PING\n"))
+			await reader.read()  // This now works as a workaround for https://github.com/cloudflare/workerd/issues/1305
 			await socket.close()
 
 			// TODO: should throw an error here but it doesn't?
@@ -48,7 +49,7 @@ async function getStatus(monitor: MonitorTarget): Promise<{ ping: number; up: bo
 		try {
 			const response = await fetchTimeout(monitor.target, monitor.timeout || 10000, {
 				method: monitor.method,
-				headers: monitor.headers,
+				headers: monitor.headers as any,
 				body: monitor.body,
 				cf: {
 					cacheTtlByStatus: {
@@ -59,6 +60,30 @@ async function getStatus(monitor: MonitorTarget): Promise<{ ping: number; up: bo
 			
 			console.log(`${monitor.name} responded with ${response.status}`)
 			status.ping = Date.now() - startTime
+			
+			if (monitor.expectedCodes) {
+				if (!monitor.expectedCodes.includes(response.status)) {
+					status.up = false
+					status.err = `Expected codes: ${JSON.stringify(monitor.expectedCodes)}, Got: ${response.status}`
+					return status
+				}
+			} else {
+				if (response.status < 200 || response.status > 299) {
+					status.up = false
+					status.err = `Expected codes: 2xx, Got: ${response.status}`
+					return status
+				}
+			}
+
+			if (monitor.responseKeyword) {
+				const responseBody = await response.text()
+				if (!responseBody.includes(monitor.responseKeyword)) {
+					status.up = false
+					status.err = "HTTP response doesn't contain the configured keyword"
+					return status
+				}
+			}
+
 			status.up = true
 			status.err = ""
 		} catch (e: any) {
@@ -80,7 +105,7 @@ async function getStatus(monitor: MonitorTarget): Promise<{ ping: number; up: bo
 
 export default {
 	async scheduled(event: ScheduledEvent, env: Env, ctx: ExecutionContext): Promise<void> {
-		const workerLocation = await getWorkerLocation()
+		const workerLocation = await getWorkerLocation() || "ERROR"
 		console.log(`Running scheduled event on ${workerLocation}...`)
 
 		// Read state, set init state if it doesn't exist