Continuing where we left off, this part will focus on adding an additional layer to authentication.
Let's start by creating an endpoint to generate tokens for connecting to the socket.
What's needed:
- a route to generate token for logged in users
- token should be a random unique sting, linked to user_id
- token should expire after a certain interval
- May or may not be reused within the interval depending on how you have the retry connection logic on client
I'm using mongodb with mongoose in my app, it supports an expires property which acts as a TTL(time to live) for a document. And _id (ObjectId) servers as a unique token. Going to keep it simple and stupid.
// TokenModel.js
const mongoose = require("mongoose");
const TokenSchema = new mongoose.Schema({
token_type: { type: String, required: true },
user_id: { type: mongoose.Types.ObjectId, ref: "User", required: true },
createdAt: { type: Date, expires: "15m", default: Date.now },
});
module.exports = mongoose.model("Token", TokenSchema);
Now we create an API to generate tokens, something like this
// controller
exports.generateToken = [
auth, // regular jwt middleware
function (req, res) {
try {
// create a new token
const tokenObj = new TokenModel({
token_type: "ws",
user_id: req.user._id,
});
// save the token
tokenObj.save(function (err) {
if (err) {
throw err;
}
return apiResponse.successResponseWithData(
res,
"Token generated successfully",
{ token: tokenObj._id }
);
});
} catch (err) {
return apiResponse.ErrorResponse(res, err);
}
},
];
// route
router.get("/token/ws",YourController.generateToken);
Now let's write an function to validate this token
// authWebSocketToken.js
const mongoose = require("mongoose");
const TokenModel = require("../models/TokenModel");
const toObjectId = (str) => mongoose.Types.ObjectId(str);
// authenticate websocket token
async function authWebSocketToken(token) {
try {
const res = await TokenModel.findById(toObjectId(token));
if (res) {
return res;
}
throw "Token not found";
} catch (err) {
throw "Websocket token authentication failed.";
}
}
module.exports = authWebSocketToken;
All the pieces are ready, time to update the websocket server logic to authenticate using this token, followed by a jwt authentication using first message payload.
// setupWebsocket.js
server.on("upgrade",
/* ... */
// replace the authentication block
if (token) {
const res = await authWebSocketToken(token);
if (res && res.user_id) {
// allow upgrade
wss.handleUpgrade(request, socket, head, function done(ws) {
wss.emit("connection", ws, request);
});
}
} else {
throw "No token found";
}
/* ... */
Next authenticate jwt and make sure the individual/broadcast messages are not sent until authentication is done.
- Move the individual actions out to a function and call it after authenticating the client.
- Keep the broadcast where it is but add a check to make sure a message is send to only authenticated users.
// setupWebsocket.js
wss.on("connection", (ctx) => {
// default value
ctx.is_authenticated = false;
/* ... */
// update the client.on message code
ctx.on("message", (message) => {
const data = JSON.parse(message);
// I expect the client to pass a type
// to distinguish between messages
if (data && data.type == "jwt") {
// the jwt authenticate we did earlier was moved here
authenticateWS({ token: data.token }, {}, (err) => {
if (err) {
ctx.terminate(); // close connection
}
// allow upgrade to web socket
ctx.send("authentication successful");
ctx.is_authenticated = true;
register(ctx); // client specific actions
});
}
});
// somewhere outside
function register(ctx) {
// setup individual pipeline
// ping-pong example
const interval = individualPipeline(ctx);
ctx.on("close", () => {
console.log("connection closed");
clearInterval(interval);
});
ctx.on("message", (message) => {
ctx.send(`echo: ${message}`);
});
}
// pipeline.js
// update broadcast example to check if client is authenticated
/* ... */
for (let c of clients.values()) {
if (c.is_authenticated) {
c.send(`broadcast message ${idx}`);
}
}
/* ... */
Our server is ready for the new authentications scheme, let's move to the client code.
- Client needs to get a token before attempting a connection, where you do it is part of your application structure. I decided to get the token in a higher component and pass it in.
- Modify the hook to accept an initPayload containing the jwt, this payload will be sent as the first message after connection is established
// webSocketHook.js
function useWebSocketLite({
...
// add a new parameter
initPayload = null
}) {
ws.onopen = () => {
/* ... */
// send initialization payload if any
// by the end of the open block
if (initPayload) {
ws.send(JSON.stringify(initPayload));
}
// move the ready state down
setReadyState(true);
/* ... */
}
}
Pass the init payload from the demo component,
// getting the ws token and jwt token is up to application logic
function App() {
const ws = useWebSocketLite({
socketUrl: socketUrl + `/demo?token=${token}`,
initPayload: {
type: "jwt",
token: user.authToken,
},
});
}
Play around with the setup to test different scenarios,
- not passing a ws token or passing an expired ws token fails to establish a connection
- connection is established with ws token but broadcast and individual messages are not sent until jwt step is done
- connection terminates if jwt step fails
Next/Last part of the series will identify the gaps in this approach, list down ideas to make it better and close off with the links to the codebase.
Top comments (0)