Theo's Typesafe Cult•17mo ago

Authentication in Websockets

In the server side, how do you make sure the client is authenticated, given that Authorisation headers are not allowed in the ws protocol (not even in the first http handshake). Analogue to that, in the client side, how do you authenticate with JWT Tokens ? Do you send the token on every web socket message ?

5 Replies

Unknown User•17mo ago

Message Not Public

oljimenez•17mo ago

I've never done what you're asking, but isn't it easier to use cookies to manage auth? Maybe that fix your not headers thing.

lucca•17mo ago

Following this thread for updates as I don't have a definitive answer for this, but I was able to successfully implement WS auth with JWTs in a TypeScript chat app I built a few months ago: https://github.com/ChromeUniverse/LuccaChat

GitHub

GitHub - ChromeUniverse/LuccaChat: A full-stack live chat app, powe...

A full-stack live chat app, powered by TypeScript, React, Express and ws - GitHub - ChromeUniverse/LuccaChat: A full-stack live chat app, powered by TypeScript, React, Express and ws

lucca•17mo ago

Here's the gist of my approach: keep a Map() or some other sort of memory store of user IDs and websocket connections, this will allow you to look up who sent the incoming message in your DB and perform whatever business logic you need. Unfortunately this won't scale past a single instance of the WS server though, unless you use something like Redis I'd imagine... https://github.com/ChromeUniverse/LuccaChat/blob/main/server/src/wss.ts

// wss.ts (the WebSockets server)

const wss = new WebSocketServer({ port: Number(process.env.WS_PORT) });

const userSocketMap = new Map<string, WebSocket>();

wss.on("listening", () => {
  console.log(`WebSockets server started on port ${process.env.WS_PORT}`);
});

wss.on("connection", async function connection(ws, req) {
  ws.on("message", async function message(rawData) {
    try {
      const jsonData = JSON.parse(rawData.toString());
      const data = baseDataSchema.parse(jsonData);

      console.log("Got parsed data", jsonData);

      // Process chat messages
      if (data.dataType === "auth") {
        handleAuth(ws, userSocketMap, prisma, jsonData);
      }
        
      //.... other message handlers
    } catch (error) {
      console.error(error);
    }
  });

  ws.on("close", function () {
    if (!ws.userId) return;
    console.log(`User ID ${ws.userId} has disconnected`);
    userSocketMap.delete(ws.userId);
  });
});

// wss.ts (the WebSockets server)

const wss = new WebSocketServer({ port: Number(process.env.WS_PORT) });

const userSocketMap = new Map<string, WebSocket>();

wss.on("listening", () => {
  console.log(`WebSockets server started on port ${process.env.WS_PORT}`);
});

wss.on("connection", async function connection(ws, req) {
  ws.on("message", async function message(rawData) {
    try {
      const jsonData = JSON.parse(rawData.toString());
      const data = baseDataSchema.parse(jsonData);

      console.log("Got parsed data", jsonData);

      // Process chat messages
      if (data.dataType === "auth") {
        handleAuth(ws, userSocketMap, prisma, jsonData);
      }
        
      //.... other message handlers
    } catch (error) {
      console.error(error);
    }
  });

  ws.on("close", function () {
    if (!ws.userId) return;
    console.log(`User ID ${ws.userId} has disconnected`);
    userSocketMap.delete(ws.userId);
  });
});

Here is the actual WS authentication handler itself: https://github.com/ChromeUniverse/LuccaChat/blob/main/server/src/websockets-handlers/auth.ts

// src/websocket-handlers/auth.ts

export async function handleAuth(
  ws: WebSocket,
  userSocketMap: Map<string, WebSocket>,
  prisma: PrismaClient,
  jsonData: any
) {
  try {
    const { token } = authSchema.parse(jsonData);

    console.log("Got auth token from client:", token);

    const decoded = (await asyncJWTverify(
      token,
      process.env.WS_JWT_SECRET as string
    )) as WsAuthJwtReceived;

    // check for JWT expiry
    const expiryTime = Number(process.env.JWT_EXPIRY);
    if (Math.round(Date.now() / 1000) - decoded.iat > expiryTime) {
      return ws.close();
    }

    // Passed verification!
    console.log(`User ID ${decoded.id} connected`);

    // Bind user ID to WebSocket, add it to map
    ws.userId = decoded.id;
    userSocketMap.set(decoded.id, ws);

    // alert user
    const ackData: z.infer<typeof authAckSchema> = {
      dataType: "auth-ack",
      error: false,
    };

    ws.send(JSON.stringify(ackData));
  } catch (error) {
    // auth failed for some reason
    console.error(error);

    // alert user
    const ackData: z.infer<typeof authAckSchema> = {
      dataType: "auth-ack",
      error: true,
    };

    ws.send(JSON.stringify(ackData));

    return ws.close();
  }
}

// src/websocket-handlers/auth.ts

export async function handleAuth(
  ws: WebSocket,
  userSocketMap: Map<string, WebSocket>,
  prisma: PrismaClient,
  jsonData: any
) {
  try {
    const { token } = authSchema.parse(jsonData);

    console.log("Got auth token from client:", token);

    const decoded = (await asyncJWTverify(
      token,
      process.env.WS_JWT_SECRET as string
    )) as WsAuthJwtReceived;

    // check for JWT expiry
    const expiryTime = Number(process.env.JWT_EXPIRY);
    if (Math.round(Date.now() / 1000) - decoded.iat > expiryTime) {
      return ws.close();
    }

    // Passed verification!
    console.log(`User ID ${decoded.id} connected`);

    // Bind user ID to WebSocket, add it to map
    ws.userId = decoded.id;
    userSocketMap.set(decoded.id, ws);

    // alert user
    const ackData: z.infer<typeof authAckSchema> = {
      dataType: "auth-ack",
      error: false,
    };

    ws.send(JSON.stringify(ackData));
  } catch (error) {
    // auth failed for some reason
    console.error(error);

    // alert user
    const ackData: z.infer<typeof authAckSchema> = {
      dataType: "auth-ack",
      error: true,
    };

    ws.send(JSON.stringify(ackData));

    return ws.close();
  }
}

I've cropped the snippets above for brevity but I've linked the full files on GitHub as well Here's the basic flow: - User sends an auth request message to the WS server - Auth request message contains a JWT that was originally signed by my REST API - WS server verifies the authenticity of the JWT - If auth succeeds, the user will get pushed to the userSocketMap which binds their WS connection to their user ID for handling later WS message exchange - If auth fails, the WS server just severs the connection Hope this helps, but I'm sure there's a better way to implement this...

nvegater•17mo ago

Thank you for your answer 🙂 The ticket message is indeed the best alternative - I should've elaborate my question a bit more since this is in the context of only one tRPC server that runs both the http and the ws connections. The ticket System is only possible with the ws and the http server being separated. But you're right! Thank you, same as what I told to @LeoRoese this is in the context of only one tRPC server that runs both the http and the ws connections.... this means that I would like to do this in the createContext function of the tRPC handlers. But technically your answer is correct to how I express my question. Thats an efficient way of doing auth with WS and JWT tokens 🙂 thanks! I think I will ask a new question with more context, since this one was a bit incomplete, and your answers are already answering what I asked 🙂 thanks!

Gaming

Programming

Authentication in Websockets