feat(auth): add token and userinfo endpoints

Adopted from jagregory#384

integration-tests/aws-sdk/setup.ts

@@ -86,7 +86,7 @@ export const withCognitoSdk =
           DefaultConfig.TokenConfig
         ),
       });
-      const server = createServer(router, ctx.logger);
+      const server = createServer(router, ctx.logger, cognitoClient);
       httpServer = await server.start({
         hostname: "127.0.0.1",
         port: 0,

package.json

@@ -68,6 +68,7 @@
     "cors": "^2.8.5",
     "express": "^4.17.1",
     "jsonwebtoken": "^8.5.1",
+    "jwk-to-pem": "^2.0.6",
     "lodash.mergewith": "^4.6.2",
     "pino": "^7.11.0",
     "pino-http": "^6.3.3",

src/server/defaults.ts

@@ -81,6 +81,7 @@ export const createDefaultServer = async (
       triggers,
     }),
     logger,
+    cognitoClient,
     {
       development: !!process.env.COGNITO_LOCAL_DEVMODE,
     }

src/server/server.ts

@@ -7,7 +7,12 @@ import * as uuid from "uuid";
 import { CognitoError, UnsupportedError } from "../errors";
 import { Router } from "./Router";
 import PublicKey from "../keys/cognitoLocal.public.json";
+import PrivateKey from "../keys/cognitoLocal.private.json";
+import jwkToPem from "jwk-to-pem";
 import Pino from "pino-http";
+import { CognitoService } from "../services";
+import { AppClient } from "../services/appClient";
+import jwt, { JwtPayload, VerifyOptions } from "jsonwebtoken";
 
 export interface ServerOptions {
   port: number;
@@ -23,6 +28,7 @@ export interface Server {
 export const createServer = (
   router: Router,
   logger: Logger,
+  cognito: CognitoService,
   options: Partial<ServerOptions> = {}
 ): Server => {
   const pino = Pino({
@@ -48,6 +54,7 @@ export const createServer = (
       type: "application/x-amz-json-1.1",
     })
   );
+  app.use(express.urlencoded({ extended: true }));
 
   app.get("/:userPoolId/.well-known/jwks.json", (req, res) => {
     res.status(200).json({
@@ -56,17 +63,237 @@ export const createServer = (
   });
 
   app.get("/:userPoolId/.well-known/openid-configuration", (req, res) => {
+    const proxyHost = req.headers["x-forwarded-host"];
+    const host = proxyHost ? proxyHost : req.headers.host;
+    const url = `http://${host}/${req.params.userPoolId}`;
+
     res.status(200).json({
+      subject_types_supported: ["public", "pairwise"],
+      grant_types_supported: ["authorization_code"],
       id_token_signing_alg_values_supported: ["RS256"],
-      jwks_uri: `http://localhost:9229/${req.params.userPoolId}/.well-known/jwks.json`,
-      issuer: `http://localhost:9229/${req.params.userPoolId}`,
+      jwks_uri: `${url}/.well-known/jwks.json`,
+      issuer: `${url}`,
+      token_endpoint: `${url}/oauth2/token`,
+      userinfo_endpoint: `${url}/oauth2/userinfo`,
     });
   });
 
   app.get("/health", (req, res) => {
     res.status(200).json({ ok: true });
   });
 
+  app.post("/:userPoolId/oauth2/token", (req, res, next) => {
+    const handleRequest = async () => {
+      const grantType = req.body.grant_type;
+
+      if (grantType !== "password") {
+        res.status(400).json({
+          error: "unsupported_grant_type",
+          description: "only 'password' grant type is supported",
+        });
+        return;
+      }
+
+      const clientId = req.body.client_id;
+
+      let userPoolClient: AppClient | null;
+
+      try {
+        userPoolClient = await cognito.getAppClient(
+          { logger: req.log },
+          clientId
+        );
+      } catch (e) {
+        res.status(500).json({
+          error: "server_error",
+          description: "failed to retrieve user pool client" + e,
+        });
+        return;
+      }
+
+      if (!userPoolClient) {
+        res.status(500).json({
+          error: "server_error",
+          description: "failed to retrieve user pool client",
+        });
+        return;
+      }
+
+      const userPool = await cognito.getUserPoolForClientId(
+        { logger: req.log },
+        clientId
+      );
+
+      const user = await userPool.getUserByUsername(
+        { logger: req.log },
+        req.body.username
+      );
+
+      if (!user) {
+        res.status(400).json({
+          error: "server_error",
+          description: "user " + req.body.username + " not found",
+        });
+        return;
+      }
+      const attr = user.Attributes;
+      const appData2 = attr.find(
+        (attribute) => attribute.Name === "esAppData2"
+      );
+      const sub = attr.find((attribute) => attribute.Name === "sub");
+
+      if (!userPoolClient) {
+        res.status(400).json({
+          error: "invalid_client",
+          description: "invalid user pool client",
+        });
+        return;
+      }
+
+      const now = Math.floor(Date.now() / 1000);
+
+      const accessToken = {
+        sub: sub?.Value,
+        client_id: clientId,
+        scope: req.body.scope,
+        jti: uuid.v4(),
+        auth_time: now,
+        iat: now,
+        nbf: now,
+        aud: clientId,
+        token_use: "access",
+        esAppData2: appData2?.Value,
+        login: user.Username,
+        status: user.Enabled,
+      };
+
+      const idToken = {
+        sub: sub,
+        client_id: clientId,
+        jti: uuid.v4(),
+        auth_time: now,
+        iat: now,
+        token_use: "id",
+      };
+
+      res.status(200).json({
+        access_token: jwt.sign(accessToken, PrivateKey.pem, {
+          algorithm: "RS256",
+          issuer: `http://${req.headers.host}/${userPoolClient.UserPoolId}`,
+          expiresIn: 3600,
+          keyid: "CognitoLocal",
+        }),
+        expiresIn: 3600,
+        id_token: jwt.sign(idToken, PrivateKey.pem, {
+          algorithm: "RS256",
+          issuer: `http://${req.headers.host}/${userPoolClient.UserPoolId}`,
+          expiresIn: 3600,
+          keyid: "CognitoLocal",
+        }),
+        token_type: "Bearer",
+      });
+    };
+
+    handleRequest().catch(next);
+  });
+
+  app.get("/:userPoolId/oauth2/userinfo", (req, res, next) => {
+    const handleRequest = async () => {
+      const authHeader = req.headers.authorization;
+      if (!authHeader || !authHeader.startsWith("Bearer ")) {
+        return res.status(401).json({
+          error: "invalid_token",
+          description:
+            "Authorization header must be provided with a Bearer token",
+        });
+      }
+
+      const { userPoolId } = req.params;
+      const token = authHeader.split(" ")[1];
+      const pem = jwkToPem(PublicKey.jwk);
+
+      const verifyOptions: VerifyOptions = {
+        algorithms: ["RS256"],
+        issuer: `http://${req.headers.host}/${userPoolId}`,
+      };
+
+      try {
+        const decodedToken = jwt.verify(
+          token,
+          pem,
+          verifyOptions
+        ) as JwtPayload;
+
+        if (typeof decodedToken === "object" && decodedToken.aud) {
+          const clientId = decodedToken["client_id"];
+
+          const userPool = await cognito.getUserPoolForClientId(
+            { logger: req.log },
+            clientId
+          );
+
+          if (!userPool) {
+            res.status(404).json({
+              error: "user_pool_not_found",
+              description: "User pool not found for the provided client ID",
+            });
+            return;
+          }
+          const users = await userPool.listUsers({ logger: req.log });
+
+          const user = users.find((user) => {
+            const subAttribute = user.Attributes.find(
+              (attr) => attr.Name === "sub"
+            );
+            return subAttribute && subAttribute.Value === decodedToken.sub;
+          });
+
+          if (!user) {
+            res.status(404).json({
+              error: "user_not_found",
+              description: "User not found",
+            });
+            return;
+          }
+
+          const attr = user.Attributes;
+          if (!attr) {
+            res.status(400).json({
+              error: "missing_attributes",
+              description: "User attributes are missing",
+            });
+            return;
+          }
+
+          const userInfo = {
+            given_name: attr.find(
+              (attribute) => attribute.Name === "given_name"
+            )?.Value,
+            family_name: attr.find(
+              (attribute) => attribute.Name === "family_name"
+            )?.Value,
+            sub: decodedToken.sub,
+            email: attr.find((attribute) => attribute.Name === "email")?.Value,
+          };
+
+          return res.status(200).json(userInfo);
+        } else {
+          return res.status(400).json({
+            error: "invalid_token",
+            description: "Token is missing required claims",
+          });
+        }
+      } catch (err) {
+        return res.status(401).json({
+          error: "invalid_token",
+          description: "Token is invalid or expired: " + err,
+        });
+      }
+    };
+
+    handleRequest().catch(next);
+  });
+
   app.post("/", (req, res) => {
     const xAmzTarget = req.headers["x-amz-target"];