feat(api): add support for oauth2 token endpoint

package.json

@@ -31,7 +31,7 @@
     "@semantic-release/exec": "^6.0.3",
     "@semantic-release/git": "^10.0.1",
     "@semantic-release/npm": "^9.0.0",
-    "@types/aws-lambda": "^8.10.85",
+    "@types/aws-lambda": "8.10.85",
     "@types/body-parser": "^1.19.0",
     "@types/cors": "^2.8.6",
     "@types/express": "^4.17.13",
@@ -49,8 +49,8 @@
     "eslint-config-prettier": "^6.15.0",
     "husky": "^7.0.4",
     "jest": "^27.4.3",
-    "jest-extended": "^3.0.1",
     "jest-date-mock": "^1.0.8",
+    "jest-extended": "^3.0.1",
     "lint-staged": "^10.1.3",
     "markdown-toc": "^1.2.0",
     "nodemon": "^2.0.3",

src/__tests__/mockTokenGenerator.ts

@@ -2,4 +2,5 @@ import { TokenGenerator } from "../services/tokenGenerator";
 
 export const newMockTokenGenerator = (): jest.Mocked<TokenGenerator> => ({
   generate: jest.fn(),
+  generateWithClientCreds: jest.fn(),
 });

src/errors.ts

@@ -101,3 +101,9 @@ export class InvalidParameterError extends CognitoError {
     super("InvalidParameterException", message);
   }
 }
+
+export class NotImplementedError extends CognitoError {
+  public constructor() {
+    super("NotImplementedException", "Function not implemented");
+  }
+}

src/server/server.ts

@@ -4,7 +4,7 @@ import express from "express";
 import * as http from "http";
 import type { Logger } from "pino";
 import * as uuid from "uuid";
-import { CognitoError, UnsupportedError } from "../errors";
+import { CognitoError, UnsupportedError, NotAuthorizedError } from "../errors";
 import { Router } from "./Router";
 import PublicKey from "../keys/cognitoLocal.public.json";
 import Pino from "pino-http";
@@ -59,6 +59,60 @@ export const createServer = (
     res.status(200).json({ ok: true });
   });
 
+  app.post("/:userPoolId/oauth2/token", (req, res) => {
+    let rawBody = "";
+    req.setEncoding("utf8");
+    req.on("data", function (chunk) {
+      rawBody += chunk;
+    });
+    req.on("end", function () {
+      const target = "GetToken";
+      const route = router(target);
+
+      const parsed = new URLSearchParams(rawBody);
+      const params = {
+        grant_type: parsed.get("grant_type"),
+        client_id: parsed.get("client_id"),
+        client_secret: parsed.get("client_secret"),
+        refresh_token: parsed.get("refresh_token"),
+      };
+
+      const auth = req.get("Authorization");
+      if (auth && auth.startsWith("Basic ")) {
+        const sliced = auth.slice("Basic ".length);
+        const buff = new Buffer(sliced, "base64");
+        const decoded = buff.toString("ascii");
+        const creds = decoded.split(":");
+        if (creds.length == 2) {
+          params.client_id = creds[0];
+          params.client_secret = creds[1];
+        }
+      }
+
+      route({ logger: req.log }, params).then(
+        (output) => {
+          res.status(200).type("json").send(JSON.stringify(output));
+        },
+        (ex) => {
+          req.log.warn(ex, `Error handling target: ${target}`);
+          if (ex instanceof NotAuthorizedError) {
+            res.status(401).json(ex);
+            return;
+          } else if (ex instanceof CognitoError) {
+            res.status(400).json({
+              code: ex.code,
+              message: ex.message,
+            });
+            return;
+          } else {
+            res.status(500).json(ex);
+            return;
+          }
+        }
+      );
+    });
+  });
+
   app.post("/", (req, res) => {
     const xAmzTarget = req.headers["x-amz-target"];
 

src/services/tokenGenerator.ts

@@ -86,8 +86,8 @@ const applyTokenOverrides = (
 
 export interface Tokens {
   readonly AccessToken: string;
-  readonly IdToken: string;
-  readonly RefreshToken: string;
+  readonly IdToken?: string;
+  readonly RefreshToken?: string;
 }
 
 export interface TokenGenerator {
@@ -104,6 +104,10 @@ export interface TokenGenerator {
       | "NewPasswordChallenge"
       | "RefreshTokens"
   ): Promise<Tokens>;
+  generateWithClientCreds(
+    ctx: Context,
+    userPoolClient: AppClient
+  ): Promise<Tokens>;
 }
 
 const formatExpiration = (
@@ -240,4 +244,38 @@ export class JwtTokenGenerator implements TokenGenerator {
       ),
     };
   }
+
+  public async generateWithClientCreds(
+    ctx: Context,
+    userPoolClient: AppClient
+  ): Promise<Tokens> {
+    const eventId = uuid.v4();
+    const authTime = Math.floor(this.clock.get().getTime() / 1000);
+
+    const accessToken: RawToken = {
+      auth_time: authTime,
+      client_id: userPoolClient.ClientId,
+      event_id: eventId,
+      iat: authTime,
+      jti: uuid.v4(),
+      scope: "aws.cognito.signin.user.admin", // TODO: scopes
+      sub: userPoolClient.ClientId,
+      token_use: "access",
+    };
+
+    const issuer = `${this.tokenConfig.IssuerDomain}/${userPoolClient.UserPoolId}`;
+
+    return await Promise.resolve({
+      AccessToken: jwt.sign(accessToken, PrivateKey.pem, {
+        algorithm: "RS256",
+        issuer,
+        expiresIn: formatExpiration(
+          userPoolClient.AccessTokenValidity,
+          userPoolClient.TokenValidityUnits?.AccessToken ?? "hours",
+          "24h"
+        ),
+        keyid: "CognitoLocal",
+      }),
+    });
+  }
 }

src/targets/getToken.test.ts

@@ -0,0 +1,100 @@
+import { newMockCognitoService } from "../__tests__/mockCognitoService";
+import { newMockTokenGenerator } from "../__tests__/mockTokenGenerator";
+import { newMockTriggers } from "../__tests__/mockTriggers";
+import { newMockUserPoolService } from "../__tests__/mockUserPoolService";
+import { TestContext } from "../__tests__/testContext";
+import * as TDB from "../__tests__/testDataBuilder";
+import { CognitoService, Triggers, UserPoolService } from "../services";
+import { TokenGenerator } from "../services/tokenGenerator";
+
+import { GetToken, GetTokenTarget } from "./getToken";
+
+describe("GetToken target", () => {
+  let getToken: GetTokenTarget;
+  let mockCognitoService: jest.Mocked<CognitoService>;
+  let mockTokenGenerator: jest.Mocked<TokenGenerator>;
+  let mockTriggers: jest.Mocked<Triggers>;
+  let mockUserPoolService: jest.Mocked<UserPoolService>;
+  const userPoolClient = TDB.appClient();
+
+  beforeEach(() => {
+    mockUserPoolService = newMockUserPoolService({
+      Id: userPoolClient.UserPoolId,
+    });
+    mockCognitoService = newMockCognitoService(mockUserPoolService);
+    mockCognitoService.getAppClient.mockResolvedValue(userPoolClient);
+    mockTriggers = newMockTriggers();
+    mockTokenGenerator = newMockTokenGenerator();
+    getToken = GetToken({
+      cognito: mockCognitoService,
+      tokenGenerator: mockTokenGenerator,
+    });
+  });
+
+  it("issues access tokens via refresh tokens", async () => {
+    mockTokenGenerator.generate.mockResolvedValue({
+      AccessToken: "access",
+      IdToken: "id",
+      RefreshToken: "refresh",
+    });
+
+    const existingUser = TDB.user({
+      RefreshTokens: ["refresh-orig"],
+    });
+    mockUserPoolService.getUserByRefreshToken.mockResolvedValue(existingUser);
+    mockUserPoolService.listUserGroupMembership.mockResolvedValue([]);
+
+    const response = await getToken(TestContext, {
+      client_id: userPoolClient.ClientId,
+      grant_type: "refresh_token",
+      refresh_token: "refresh-orig",
+    });
+    expect(mockUserPoolService.getUserByRefreshToken).toHaveBeenCalledWith(
+      TestContext,
+      "refresh-orig"
+    );
+    expect(mockUserPoolService.storeRefreshToken).not.toHaveBeenCalled();
+
+    expect(response.access_token).toEqual("access");
+    expect(response.refresh_token).toEqual("refresh");
+  });
+});
+
+describe("GetToken target - Client Creds", () => {
+  let getToken: GetTokenTarget;
+  let mockCognitoService: jest.Mocked<CognitoService>;
+  let mockTokenGenerator: jest.Mocked<TokenGenerator>;
+  let mockUserPoolService: jest.Mocked<UserPoolService>;
+  const userPoolClient = TDB.appClient({
+    ClientSecret: "secret",
+    ClientId: "id",
+  });
+
+  beforeEach(() => {
+    mockUserPoolService = newMockUserPoolService({
+      Id: userPoolClient.UserPoolId,
+    });
+    mockCognitoService = newMockCognitoService(mockUserPoolService);
+    mockCognitoService.getAppClient.mockResolvedValue(userPoolClient);
+    mockTokenGenerator = newMockTokenGenerator();
+    getToken = GetToken({
+      cognito: mockCognitoService,
+      tokenGenerator: mockTokenGenerator,
+    });
+  });
+
+  it("issues access tokens via client credentials", async () => {
+    mockTokenGenerator.generateWithClientCreds.mockResolvedValue({
+      AccessToken: "access",
+      RefreshToken: null,
+      IdToken: null,
+    });
+
+    const response = await getToken(TestContext, {
+      client_id: userPoolClient.ClientId,
+      client_secret: userPoolClient.ClientSecret,
+      grant_type: "client_credentials",
+    });
+    expect(response.access_token).toEqual("access");
+  });
+});

src/targets/getToken.ts

@@ -0,0 +1,126 @@
+import {
+  InvalidParameterError,
+  NotAuthorizedError,
+  NotImplementedError,
+} from "../errors";
+import { Services } from "../services";
+import { Context } from "../services/context";
+import { Target } from "../targets/Target";
+
+type HandleTokenServices = Pick<Services, "cognito" | "tokenGenerator">;
+
+export type GetTokenRequest =
+  | GetTokenRequestClientCreds
+  | GetTokenRequestRefreshToken
+  | GetTokenRequestAuthCode;
+
+interface GetTokenRequestGrantType {
+  grant_type: "authorization_code" | "client_credentials" | "refresh_token";
+  client_id: string;
+}
+
+interface GetTokenRequestClientCreds extends GetTokenRequestGrantType {
+  client_secret: string;
+}
+
+type GetTokenRequestAuthCode = GetTokenRequestGrantType;
+
+interface GetTokenRequestRefreshToken extends GetTokenRequestGrantType {
+  refresh_token: string;
+}
+
+interface GetTokenResponse {
+  access_token: string;
+  refresh_token?: string;
+}
+
+export type GetTokenTarget = Target<GetTokenRequest, GetTokenResponse>;
+
+async function getWithRefreshToken(
+  ctx: Context,
+  services: HandleTokenServices,
+  params: GetTokenRequestRefreshToken
+) {
+  const clientId = params.client_id;
+  const userPool = await services.cognito.getUserPoolForClientId(ctx, clientId);
+  const userPoolClient = await services.cognito.getAppClient(ctx, clientId);
+  const user = await userPool.getUserByRefreshToken(ctx, params.refresh_token);
+  if (!user || !userPoolClient) {
+    throw new NotAuthorizedError();
+  }
+
+  const userGroups = await userPool.listUserGroupMembership(ctx, user);
+
+  const tokens = await services.tokenGenerator.generate(
+    ctx,
+    user,
+    userGroups,
+    userPoolClient,
+    undefined,
+    "RefreshTokens"
+  );
+
+  return {
+    access_token: tokens.AccessToken,
+    refresh_token: tokens.RefreshToken,
+  };
+}
+
+async function getWithClientCredentials(
+  ctx: Context,
+  services: HandleTokenServices,
+  params: GetTokenRequestClientCreds
+) {
+  const clientId = params.client_id;
+  const clientSecret = params.client_secret;
+  const userPoolClient = await services.cognito.getAppClient(ctx, clientId);
+  if (!userPoolClient) {
+    throw new NotAuthorizedError();
+  }
+  if (
+    userPoolClient.ClientSecret &&
+    userPoolClient.ClientSecret != clientSecret
+  ) {
+    throw new NotAuthorizedError();
+  }
+
+  const tokens = await services.tokenGenerator.generateWithClientCreds(
+    ctx,
+    userPoolClient
+  );
+  if (!tokens) {
+    throw new NotAuthorizedError();
+  }
+
+  return {
+    access_token: tokens.AccessToken,
+  };
+}
+
+export const GetToken =
+  (services: HandleTokenServices): GetTokenTarget =>
+  async (ctx, req) => {
+    switch (req.grant_type) {
+      case "authorization_code": {
+        throw new NotImplementedError();
+      }
+      case "client_credentials": {
+        return getWithClientCredentials(
+          ctx,
+          services,
+          req as GetTokenRequestClientCreds
+        );
+      }
+      case "refresh_token": {
+        return getWithRefreshToken(
+          ctx,
+          services,
+          req as GetTokenRequestRefreshToken
+        );
+      }
+      default: {
+        console.log("Invalid grant type passed:", req.grant_type);
+        throw new InvalidParameterError();
+      }
+    }
+  };