src/hero.mjs: Add handlerForConnection()

2 files changed, 56 insertions, 0 deletions

diff --git a/src/hero.mjs b/src/hero.mjs
index 10fa06e..61621af 100644
--- a/src/hero.mjs
+++ b/src/hero.mjs
@@ -309,6 +309,22 @@ export const buildHttpPayload = (code, {
 	"\r\n\r\n" +
 	message;
 
+export const fallback404Handler = (_req, socket) =>
+	socket.end(buildHttpPayload(404)).destroySoon();
+
+export const fallback405Handler = (_req, socket) =>
+	socket.end(buildHttpPayload(405)).destroySoon();
+
+export const handlerForConnection = (table, method, path) => {
+	if (method !== WEBSOCKET_METHOD) {
+		return fallback405Handler;
+	}
+
+	const segments = pathToSegments(path);
+	const handlerFn = u.getIn(table, [WEBSOCKET_KEYWORD, WEBSOCKET_METHOD].concat(segments));
+	return handlerFn || fallback404Handler;
+};
+
 export const makeRequestListener = table => async (req, res) => {
 	const { status, headers, body } = await handleRequest(table, req);
 	res.writeHead(status, headers);
diff --git a/tests/js/hero.mjs b/tests/js/hero.mjs
index f964a88..436535a 100644
--- a/tests/js/hero.mjs
+++ b/tests/js/hero.mjs
@@ -33,6 +33,9 @@ import {
 	handleRequest,
 	emitHeaders,
 	buildHttpPayload,
+	fallback404Handler,
+	fallback405Handler,
+	handlerForConnection,
 	makeRequestListener,
 	actionsFn,
 	lineHandlerFn,
@@ -1217,6 +1220,42 @@ const test_buildHttpPayload = async t => {
 	});
 };
 
+export const test_handlerForConnection = async t => {
+	t.start("handlerForConnection()");
+
+	await t.test("405 handler no matter the table when not GET method", () => {
+		assert.equal(
+			handlerForConnection(null, "POST", null),
+			fallback405Handler,
+		);
+	});
+
+	await t.test("404 handler when there is no route", () => {
+		assert.equal(
+			handlerForConnection({}, "GET", "/the/websocket"),
+			fallback404Handler,
+		);
+	});
+
+	await t.test("the declared handler when it exists", () => {
+		const fn1 = () => {};
+		assert.equal(
+			handlerForConnection({
+				websocket: {
+					GET: {
+						known: {
+							path: {
+								"": fn1,
+							},
+						},
+					},
+				},
+			}, "GET", "/known/path"),
+			fn1,
+		);
+	});
+};
+
 const test_makeRequestListener = async t => {
 	t.start("makeRequestListener()");
 
@@ -1911,6 +1950,7 @@ await runner.runTests([
 	test_handleRequest,
 	test_emitHeaders,
 	test_buildHttpPayload,
+	test_handlerForConnection,
 	test_makeRequestListener,
 	test_actionsFn,
 	test_lineHandlerFn,