Fix/web socket (#89)

Added error and loading screen and improved web socket security
code100x · Sep 13, 2024 · 7eef477 · 7eef477
1 parent f3d84aa
commit 7eef477
Show file tree

Hide file tree

Showing 8 changed files with 142 additions and 54 deletions.
diff --git a/next-app/app/creator/[creatorId]/page.tsx b/next-app/app/creator/[creatorId]/page.tsx
@@ -5,6 +5,8 @@ import { useSocket } from "@/context/socket-context";
 import { useSession } from "next-auth/react";
 import { useEffect } from "react";
 import jwt from "jsonwebtoken";
+import ErrorScreen from "@/components/ErrorScreen";
+import LoadingScreen from "@/components/LoadingScreen";
 
 export default function Creator({
   params: { creatorId },
@@ -13,18 +15,17 @@ export default function Creator({
     creatorId: string;
   };
 }) {
-  const { socket, user } = useSocket();
-  const session = useSession();
+  const { socket, user, connectionError, loading, setUser } = useSocket();
   useRedirect();
 
   useEffect(() => {
-    if (user) {
+    if (user && !user.token) {
       const token = jwt.sign(
         {
           creatorId: creatorId,
           userId: user.id,
         },
-        process.env.NEXT_PUBLIC_SECRET ?? "secret",
+        process.env.NEXT_PUBLIC_SECRET ?? "secret"
       );
 
       socket?.send(
@@ -33,14 +34,25 @@ export default function Creator({
           data: {
             token,
           },
-        }),
+        })
       );
+
+      setUser({ ...user, token });
     }
   }, [user]);
 
-  if (!session.data) {
-    return <h1>Please Log in....</h1>;
+  if (connectionError) {
+    return <ErrorScreen>Cannot connect to socket server</ErrorScreen>;
+  }
+
+  if (loading) {
+    return <LoadingScreen />;
   }
+
+  if (!user) {
+    return <ErrorScreen>Please Log in....</ErrorScreen>;
+  }
+
   return <StreamView creatorId={creatorId} playVideo={false} />;
 }
 

diff --git a/next-app/app/dashboard/page.tsx b/next-app/app/dashboard/page.tsx
@@ -4,13 +4,15 @@ import { useSocket } from "@/context/socket-context";
 import useRedirect from "../../hooks/useRedirect";
 import jwt from "jsonwebtoken";
 import StreamView from "../../components/StreamView";
+import ErrorScreen from "@/components/ErrorScreen";
+import LoadingScreen from "@/components/LoadingScreen";
 
 export default function Component() {
-  const { socket, user, connectionError } = useSocket();
+  const { socket, user, loading, setUser, connectionError } = useSocket();
   useRedirect();
 
   useEffect(() => {
-    if (user) {
+    if (user && !user.token) {
       const token = jwt.sign(
         {
           creatorId: user?.id,
@@ -19,7 +21,7 @@ export default function Component() {
         process.env.NEXT_PUBLIC_SECRET || "",
         {
           expiresIn: "24h",
-        },
+        }
       );
 
       socket?.send(
@@ -28,17 +30,22 @@ export default function Component() {
           data: {
             token,
           },
-        }),
+        })
       );
+      setUser({ ...user, token });
     }
   }, [user]);
 
   if (connectionError) {
-    return <h1>Cannot connect to socket server</h1>;
+    return <ErrorScreen>Cannot connect to socket server</ErrorScreen>;
+  }
+
+  if (loading) {
+    return <LoadingScreen />;
   }
 
   if (!user) {
-    return <h1>Please Log in....</h1>;
+    return <ErrorScreen>Please Log in....</ErrorScreen>;
   }
 
   return <StreamView creatorId={user.id} playVideo={true} />;

diff --git a/next-app/components/ErrorScreen.tsx b/next-app/components/ErrorScreen.tsx
@@ -0,0 +1,9 @@
+import React, { PropsWithChildren } from "react";
+
+export default function ErrorScreen({ children }: PropsWithChildren) {
+  return (
+    <div className="w-screen h-screen flex justify-center items-center bg-[rgb(10,10,10)] text-gray-200">
+      {children}
+    </div>
+  );
+}
diff --git a/next-app/components/LoadingScreen.tsx b/next-app/components/LoadingScreen.tsx
@@ -0,0 +1,14 @@
+import React from "react";
+
+export default function LoadingScreen() {
+  return (
+    <div className=" bg-[rgb(10,10,10)] text-gray-200">
+      <div className="flex space-x-2 justify-center items-center w-screen h-screen dark:invert">
+        <span className="sr-only">Loading...</span>
+        <div className="h-8 w-8 bg-white rounded-full animate-bounce [animation-delay:-0.3s]"></div>
+        <div className="h-8 w-8 bg-white rounded-full animate-bounce [animation-delay:-0.15s]"></div>
+        <div className="h-8 w-8 bg-white rounded-full animate-bounce"></div>
+      </div>
+    </div>
+  );
+}
diff --git a/next-app/context/socket-context.tsx b/next-app/context/socket-context.tsx
@@ -1,6 +1,8 @@
 import { useSession } from "next-auth/react";
 import {
+  Dispatch,
   PropsWithChildren,
+  SetStateAction,
   createContext,
   useContext,
   useEffect,
@@ -9,37 +11,45 @@ import {
 
 type SocketContextType = {
   socket: null | WebSocket;
-  user: null | { id: string };
+  user: null | { id: string; token?: string };
   connectionError: boolean;
+  setUser: Dispatch<SetStateAction<{ id: string; token?: string } | null>>;
+  loading: boolean;
 };
 
 const SocketContext = createContext<SocketContextType>({
   socket: null,
   user: null,
   connectionError: false,
+  setUser: () => {},
+  loading: true,
 });
 
 export const SocketContextProvider = ({ children }: PropsWithChildren) => {
   const [socket, setSocket] = useState<WebSocket | null>(null);
-  const [user, setUser] = useState<{ id: string } | null>(null);
+  const [user, setUser] = useState<{ id: string; token?: string } | null>(null);
   const [connectionError, setConnectionError] = useState<boolean>(false);
+  const [loading, setLoading] = useState(true);
   const session = useSession();
 
   useEffect(() => {
     if (!socket && session.data?.user.id) {
       const ws = new WebSocket(process.env.NEXT_PUBLIC_WSS_URL as string);
       ws.onopen = () => {
         setSocket(ws);
-        setUser(session.data?.user);
+        setUser(session.data?.user || null);
+        setLoading(false);
       };
 
       ws.onclose = () => {
         setSocket(null);
+        setLoading(false);
       };
 
       ws.onerror = () => {
         setSocket(null);
         setConnectionError(true);
+        setLoading(false);
       };
 
       () => {
@@ -54,6 +64,8 @@ export const SocketContextProvider = ({ children }: PropsWithChildren) => {
         socket,
         user,
         connectionError,
+        setUser,
+        loading,
       }}
     >
       {children}
@@ -62,16 +74,20 @@ export const SocketContextProvider = ({ children }: PropsWithChildren) => {
 };
 
 export const useSocket = () => {
-  const { socket, user, connectionError } = useContext(SocketContext);
+  const { socket, user, setUser, connectionError, loading } =
+    useContext(SocketContext);
 
   const sendMessage = (type: string, data: { [key: string]: any }) => {
     socket?.send(
       JSON.stringify({
         type,
-        data,
-      }),
+        data: {
+          ...data,
+          token: user?.token,
+        },
+      })
     );
   };
 
-  return { socket, sendMessage, user, connectionError };
+  return { socket, loading, setUser, sendMessage, user, connectionError };
 };
diff --git a/ws/package.json b/ws/package.json
@@ -5,6 +5,7 @@
   "main": "index.js",
   "scripts": {
     "preinstall": "npx only-allow pnpm",
+    "postinstall": "prisma generate",
     "dev": "nodemon src/app.ts",
     "start": "node dist/app.js",
     "build": "prisma generate && npx tsc -b",
@@ -14,7 +15,7 @@
   "author": "",
   "license": "ISC",
   "dependencies": {
-    "@prisma/client": "^5.19.1",
+    "@prisma/client": "5.19.1",
     "axios": "^1.7.7",
     "bullmq": "^5.12.14",
     "cors": "^2.8.5",
@@ -32,4 +33,4 @@
     "ts-node": "^10.9.2",
     "typescript": "^5.6.2"
   }
-}
+}
diff --git a/ws/src/StramManager.ts b/ws/src/StramManager.ts
@@ -1,12 +1,12 @@
 import { WebSocket } from "ws";
 import { createClient, RedisClientType } from "redis";
-import { PrismaClient } from "@prisma/client";
 //@ts-ignore
 import youtubesearchapi from "youtube-search-api";
 import { Job, Queue, Worker } from "bullmq";
+import { PrismaClient } from "@prisma/client";
 
-const TIME_SPAN_FOR_VOTE = 1200000 / 40; // 20min
-const TIME_SPAN_FOR_QUEUE = 1200000 / 40; // 20min
+const TIME_SPAN_FOR_VOTE = 1200000; // 20min
+const TIME_SPAN_FOR_QUEUE = 1200000; // 20min
 const TIME_SPAN_FOR_REPEAT = 3600000;
 const MAX_QUEUE_LENGTH = 20;
 
@@ -91,9 +91,6 @@ export class RoomManager {
     await this.redisClient.connect();
     await this.subscriber.connect();
     await this.publisher.connect();
-    this.worker.on("error", () => {
-      console.log("Worker  ready");
-    });
   }
 
   onSubscribeRoom(message: string, creatorId: string) {
@@ -141,14 +138,20 @@ export class RoomManager {
     }
   }
 
-  async addUser(userId: string, ws: WebSocket) {
+  async addUser(userId: string, ws: WebSocket, token: string) {
     this.users.set(userId, {
       userId,
       ws,
+      token,
     });
   }
 
-  async joinRoom(creatorId: string, userId: string, ws: WebSocket) {
+  async joinRoom(
+    creatorId: string,
+    userId: string,
+    ws: WebSocket,
+    token: string
+  ) {
     let room = this.rooms.get(creatorId);
     let user = this.users.get(userId);
 
@@ -158,7 +161,7 @@ export class RoomManager {
     }
 
     if (!user) {
-      await this.addUser(userId, ws);
+      await this.addUser(userId, ws, token);
       user = this.users.get(userId);
     }
 
@@ -645,6 +648,7 @@ export class RoomManager {
 type User = {
   userId: string;
   ws: WebSocket;
+  token: string;
 };
 
 type Room = {