fixed websockets (#373)

Author: farhanW3

src/server/index.ts

@@ -12,6 +12,7 @@ import { withErrorHandler } from "./middleware/error";
 import { withExpress } from "./middleware/express";
 import { withRequestLogs } from "./middleware/logs";
 import { withOpenApi } from "./middleware/open-api";
+import { withWebSocket } from "./middleware/websocket";
 import { withRoutes } from "./routes";
 import { writeOpenApiToFile } from "./utils/openapi";
 
@@ -51,6 +52,7 @@ const main = async () => {
   await withCors(server);
   await withRequestLogs(server);
   await withErrorHandler(server);
+  await withWebSocket(server);
   await withAuth(server);
   await withExpress(server);
   await withOpenApi(server);

src/server/middleware/websocket.ts

@@ -0,0 +1,23 @@
+import WebSocketPlugin from "@fastify/websocket";
+import { FastifyInstance } from "fastify";
+import { logger } from "../../utils/logger";
+
+export const withWebSocket = async (server: FastifyInstance) => {
+  await server.register(WebSocketPlugin, {
+    errorHandler: function (
+      error,
+      conn /* SocketStream */,
+      req /* FastifyRequest */,
+      reply /* FastifyReply */,
+    ) {
+      logger({
+        service: "websocket",
+        level: "error",
+        message: `Websocket error: ${error}`,
+      });
+      // Do stuff
+      // destroy/close connection
+      conn.destroy(error);
+    },
+  });
+};

src/server/routes/transaction/status.ts

@@ -95,18 +95,11 @@ export async function checkTxStatus(fastify: FastifyInstance) {
       });
     },
     wsHandler: async (connection: SocketStream, request) => {
-      logger({
-        service: "server",
-        level: "info",
-        message: `Websocket route handler`,
-        data: request,
-      });
-
       const { queueId } = request.params;
       // const timeout = await wsTimeout(connection, queueId, request);
 
       logger({
-        service: "server",
+        service: "websocket",
         level: "info",
         message: `Websocket connection established for ${queueId}`,
       });
@@ -127,7 +120,7 @@ export async function checkTxStatus(fastify: FastifyInstance) {
 
       connection.socket.on("error", (error) => {
         logger({
-          service: "server",
+          service: "websocket",
           level: "error",
           message: `Websocket error`,
           error,
@@ -138,7 +131,7 @@ export async function checkTxStatus(fastify: FastifyInstance) {
 
       connection.socket.on("message", async (message, isBinary) => {
         logger({
-          service: "server",
+          service: "websocket",
           level: "info",
           message: `Websocket message received`,
           data: message,
@@ -149,7 +142,7 @@ export async function checkTxStatus(fastify: FastifyInstance) {
 
       connection.socket.on("close", () => {
         logger({
-          service: "server",
+          service: "websocket",
           level: "info",
           message: `Websocket connection closed`,
         });

src/utils/env.ts

@@ -52,9 +52,11 @@ export const env = createEnv({
       .default("debug"),
     LOG_SERVICES: z
       .string()
-      .default("server,worker,cache")
+      .default("server,worker,cache,websocket")
       .transform((s) =>
-        z.array(z.enum(["server", "worker", "cache"])).parse(s.split(",")),
+        z
+          .array(z.enum(["server", "worker", "cache", "websocket"]))
+          .parse(s.split(",")),
       ),
     THIRDWEB_API_SECRET_KEY: z.string().min(1),
     ADMIN_WALLET_ADDRESS: z.string().min(1),