Using redis cache to support multiple server (#5)

`redis or ur or useRedis: (default false)`- If you want to run multiple servers behind a load balancer (or multiple pods in Kubernetes) use this option to keep token information in one shared Redis server. Redis information can be provided with environment variables `REDIS_HOST` and `REDIS_PORT`.
pankajupadhyay29 · Mar 5, 2024 · 392d4e1 · 392d4e1
1 parent 112be01
commit 392d4e1
Show file tree

Hide file tree

Showing 7 changed files with 134 additions and 34 deletions.
diff --git a/README.md b/README.md
@@ -82,6 +82,7 @@ This is a mock server implementation for Open ID Connect base authentication. Th
     }
   }
   ```
-
+- `redis or ur or useRedis: (default false)`-
+If you want to run multiple server behind load balancer (or multiple pods in Kubernetes) use this option to keep token information in one shared Redis server. Redis information can be provided with environment variables `REDIS_HOST` and `REDIS_PORT`.
 # If you want to contribute 
 Coming soon
diff --git a/api/oauth.js b/api/oauth.js
@@ -2,14 +2,14 @@ const utils = require('../utils');
 const { getToken, addToken, removeToken } = require('./tokens');
 const { getUser, activateUser, deactivateUser } = require('./users');
 
-const authorize = (req, res) => {
+const authorize = async (req, res) => {
     const options = utils.getOptions();
     const { redirect_uri, response_type, scope, client_id } = req.query;
     if (options.skipLogin || req.cookies.mock_auth_session) {
         const sessionID = !req.cookies.mock_auth_session && options.skipLogin
             ? activateUser('', req.query[options.connectionKey])
             : req.cookies.mock_auth_session;
-        const user = getUser(sessionID)
+        const user = await getUser(sessionID)
         if (user) {
             redirectAfterLogin(req.query, req, res, user, sessionID);
             return;
@@ -18,8 +18,9 @@ const authorize = (req, res) => {
     res.redirect(`/login?protocol=oauth2&redirect_uri=${redirect_uri}&response_type=${response_type}&scope=${scope}&client_id=${client_id}`);
 };
 
-const token = (req, res) => {
-    res.send(getToken(req.body.code))
+const token = async (req, res) => {
+    const result = await getToken(req.body.code);
+    res.send(result);
 }
 
 const jwks = (req, res) => {
@@ -33,32 +34,31 @@ const jwks = (req, res) => {
     );
 }
 
-const login = (req, res) => {
-    const userName =  req.body.username;
-    const password =  req.body.password;
+const login = async (req, res) => {
+    const userName = req.body.username;
+    const password = req.body.password;
     if (userName === password) {
         const options = utils.getOptions();
-        const sessionID = activateUser(userName, req.query[options.connectionKey]);
-        const user = getUser(sessionID);
+        const sessionID = await activateUser(userName, req.query[options.connectionKey]);
+        const user = await getUser(sessionID);
         redirectAfterLogin(req.body, req, res, user, sessionID);
     } else {
         res.status(401).send('Incorrect credentials');
     }
 }
 
-const logout = (req, res) => {
+const logout = async (req, res) => {
     const sessionID = req.cookies.mock_auth_session;
-    deactivateUser(sessionID);
-    removeToken(sessionID);
+    await deactivateUser(sessionID);
+    await removeToken(sessionID);
     setAuthCookie(req, res, '');
     res.send('You are logged out successfully.');
 }
 
-function redirectAfterLogin(data, req, res, user, sessionID) {
+async function redirectAfterLogin(data, req, res, user, sessionID) {
     const { client_id, redirect_uri, response_type, scope, connection } = data;
     console.log(client_id, redirect_uri, response_type, scope);
-    console.log(data.client_id, data.redirect_uri, data.response_type, data.scope);
-    addToken(sessionID, user, scope, client_id, connection);
+    await addToken(sessionID, user, scope, client_id, connection);
     setAuthCookie(req, res, sessionID);
     res.redirect(`${redirect_uri}&${response_type}=${sessionID}`);
 }

diff --git a/api/tokens.js b/api/tokens.js
@@ -1,13 +1,23 @@
 const _ = require('lodash');
 const utils = require('../utils');
+const { setData, getData } = require('../utils/redis');
 
 const tokens = {};
 
-const getToken = (code) => {
-    return tokens[code];
+const getToken = async (code) => {
+    let token = tokens[code];
+    if (!token && utils.getOptions().useRedis) {
+        console.log('getting token from redis', code);
+        const tokenData = await getData('token', code);
+        if (tokenData) {
+            token = JSON.parse(tokenData);
+            _.set(tokens, code, token);
+        }
+    }
+    return token;
 };
 
-const addToken = (key, user, scope, audience) => {
+const addToken = async (key, user, scope, audience) => {
     const options = utils.getOptions();
     const access_token = utils.generateJWT({}, options.keys.privateKey, audience);
     const id_token = utils.generateJWT(user, options.keys.privateKey, audience);
@@ -18,12 +28,22 @@ const addToken = (key, user, scope, audience) => {
         expires_in: 86400,
         token_type: 'Bearer',
     };
+    await setToken(key, token);
+};
+
+const setToken = async (key, token) => {
     _.set(tokens, key, token);
-    return token;
-}
+    if (utils.getOptions().useRedis) {
+        console.log('setting token in redis', key, token);
+        await setData('token', key, JSON.stringify(token));
+    }
+};
 
 const removeToken = (key) => {
-    return tokens[key];
+    delete _.unset(tokens, key);
+    if (utils.getOptions().useRedis) {
+        deleteKey('token', key);
+    }
 };
 
 module.exports = { getToken, addToken, removeToken };
diff --git a/api/users.js b/api/users.js
@@ -1,6 +1,7 @@
 const _ = require('lodash');
 const { faker } = require('@faker-js/faker');
 const { getOptions } = require('../utils');
+const { getData, setData, deleteKey } = require('../utils/redis');
 const activeUsers = {};
 
 var userStore = null;
@@ -22,13 +23,12 @@ const getRandomUser = (users, options) => {
 const getUserFromStore = (id, connection, options) => {
     const users = userStore.users || userStore[connection]
     const user = users[id];
-    return !user && options.skipLogin 
+    return !user && options.skipLogin
         ? getRandomUser(users, options)
         : { ...user, [options.idField]: id };
 }
 
-
-const activateUser = (id, connection) => {
+const activateUser = async (id, connection = '') => {
     ensureUserDB();
     const options = getOptions();
     const user = options.users
@@ -39,18 +39,37 @@ const activateUser = (id, connection) => {
             sub: faker.string.uuid(10),
             [options.idField]: id
         };
-    const key = userStore && userStore.users ? user[options.idField] : `${connection}_${user[options.idField]}`;
-    console.log(options.users, user, key);
-    _.set(activeUsers, key, user);
-    return key;
+    const sessionID = faker.string.uuid(10);
+    console.log(options.users, user, sessionID);
+    await addActiveUser(sessionID, user, options);
+    return sessionID;
 }
 
-const getUser = (key) => {
-    return _.get(activeUsers, key);
+const addActiveUser = async (key, user, options) => {
+    _.set(activeUsers, key, user);
+    if (options.useRedis) {
+        console.log('setting user in redis', key, user);
+        await setData('user', key, user);
+    }
 }
 
+const getUser = async (key) => {
+    let user = _.get(activeUsers, key);
+    if (!user && getOptions().useRedis) {
+        console.log('getting user from redis', key);
+        user = await getData('user', key);
+        if (user) {
+            _.set(activeUsers, key, user);
+        }
+    }
+    return user;
+};
+
 const deactivateUser = (key) => {
     delete _.unset(activeUsers, key);
+    if (getOptions().useRedis) {
+        deleteKey('user', key);
+    }
 }
 
 module.exports = { activateUser, getUser, deactivateUser };
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "mock-auth-server",
-  "version": "1.0.0",
+  "version": "2.0.0",
   "description": "Mock for Open ID Connect base authentication server for testing",
   "repository": {
     "type": "git",
@@ -34,6 +34,7 @@
     "express": "4.18.2",
     "jsonwebtoken": "9.0.2",
     "lodash": "4.17.21",
-    "pem-jwk": "2.0.0"
+    "pem-jwk": "2.0.0",
+    "redis": "4.6.13"
   }
 }
diff --git a/utils/common.js b/utils/common.js
@@ -67,7 +67,7 @@ const populateOptions = async () => {
     }
     options.keys = await getJWTKeys(args.pvtk || args.privateKey, args.pubk || args.publicKey);
     options.ssl = options.sslKey && options.sslCert;
-
+    options.useRedis = args.ur || args.redis || args.useRedis;
     return Promise.resolve(options);
 }
 

diff --git a/utils/redis.js b/utils/redis.js
@@ -0,0 +1,59 @@
+const redis = require('redis');
+
+let client = null;
+
+
+
+async function ensureConnection() {
+    if (!client) {
+        if (process.env.REDIS_HOST && process.env.REDIS_PORT) {
+            client = redis.createClient({
+                socket: {
+                    host: process.env.REDIS_HOST,
+                    port: process.env.REDIS_PORT
+                },
+            });
+            await openConnection();
+        }
+    } else if (!client.isReady) {
+        await openConnection();
+    }
+}
+
+async function openConnection() {
+    if (!client.isOpen) await client.connect().catch(console.error);
+    await client.ping().then(console.log).catch(console.error);
+}
+
+async function setData(store, key, value) {
+    await ensureConnection();
+    const data = typeof value === 'object' ? JSON.stringify(value) : value;
+
+    await client.set(`${store}_${key}`, data);
+    return true;
+}
+
+async function getData(store, key) {
+    await ensureConnection();
+    return client.get(`${store}_${key}`)
+        .then((reply) => {
+            console.log(`Data retrieved successfully: ${reply}`);
+            try {
+                return JSON.parse(reply);
+            } catch(e) {
+                console.error('Error parsing data: ', e.message);
+                throw e;
+            }
+        });
+}
+
+async function deleteKey(store, key) {
+    await ensureConnection();
+    return client.del(`${store}_${key}`);
+}
+
+module.exports = {
+    setData,
+    getData,
+    deleteKey,
+};