From 08928fffa45e6356c3dc6bdbeb441c6ff900a93e Mon Sep 17 00:00:00 2001
From: simeng-li <simeng@silverhand.io>
Date: Wed, 4 Dec 2024 14:44:30 +0800
Subject: [PATCH] feat(core): implement SAML IdP response flow

implement SAML IdP respones flow
---
 .../src/saml-applications/routes/anonymous.ts | 171 +++++++++++++++++-
 .../src/saml-applications/routes/utils.ts     |  74 ++++++++
 packages/core/src/sso/types/oidc.ts           |   2 +
 3 files changed, 246 insertions(+), 1 deletion(-)
 create mode 100644 packages/core/src/saml-applications/routes/utils.ts

diff --git a/packages/core/src/saml-applications/routes/anonymous.ts b/packages/core/src/saml-applications/routes/anonymous.ts
index 661eb48bb..131af9dfc 100644
--- a/packages/core/src/saml-applications/routes/anonymous.ts
+++ b/packages/core/src/saml-applications/routes/anonymous.ts
@@ -1,10 +1,33 @@
+import { parseJson } from '@logto/connector-kit';
+import { BindingType } from '@logto/schemas';
+import camelcaseKeys from 'camelcase-keys';
+import { got } from 'got';
+import saml from 'samlify';
 import { z } from 'zod';
 
+import RequestError from '#src/errors/RequestError/index.js';
 import koaGuard from '#src/middleware/koa-guard.js';
 import type { AnonymousRouter, RouterInitArgs } from '#src/routes/types.js';
+import { fetchOidcConfig, getUserInfo } from '#src/sso/OidcConnector/utils.js';
+import { oidcTokenResponseGuard } from '#src/sso/types/oidc.js';
+
+import assertThat from '../../utils/assert-that.js';
+
+import { createSamlTemplateCallback, samlLogInResponseTemplate } from './utils.js';
+
+const samlApplicationSignInCallbackQueryParametersGuard = z.union([
+  z.object({
+    code: z.string(),
+    iss: z.string(),
+  }),
+  z.object({
+    error: z.string(),
+    error_description: z.string().optional(),
+  }),
+]);
 
 export default function samlApplicationAnonymousRoutes<T extends AnonymousRouter>(
-  ...[router, { libraries }]: RouterInitArgs<T>
+  ...[router, { libraries, queries }]: RouterInitArgs<T>
 ) {
   const {
     samlApplications: { getSamlIdPMetadataByApplicationId },
@@ -29,4 +52,150 @@ export default function samlApplicationAnonymousRoutes<T extends AnonymousRouter
       return next();
     }
   );
+
+  router.get(
+    '/saml-applications/:id/callback',
+    koaGuard({
+      params: z.object({ id: z.string() }),
+      status: [200],
+    }),
+    async (ctx, next) => {
+      const { id } = ctx.guard.params;
+      const queryParameters = ctx.request.query;
+
+      // Find the SAML application secret by application ID
+      const { applications, samlApplicationSecrets, samlApplicationConfigs } = queries;
+      const {
+        secret,
+        oidcClientMetadata: { redirectUris },
+      } = await applications.findApplicationById(id);
+
+      // Sign-in callback handler
+      const query = samlApplicationSignInCallbackQueryParametersGuard.safeParse(queryParameters);
+
+      if (!query.success) {
+        throw new RequestError('guard.invalid_input');
+      }
+
+      if ('error' in query.data) {
+        throw new RequestError({
+          code: 'oidc.invalid_request',
+          status: 400,
+          message: query.data.error_description,
+        });
+      }
+
+      // TODO: need to validate state for SP initiated SAML flow
+      const { code, iss } = query.data;
+
+      const { tokenEndpoint, userinfoEndpoint } = await fetchOidcConfig(iss);
+
+      const headers = {
+        Authorization: `Basic ${Buffer.from(`${id}:${secret}`, 'utf8').toString('base64')}`,
+        'Content-Type': 'application/x-www-form-urlencoded',
+      };
+
+      const tokenRequestParameters = new URLSearchParams({
+        grant_type: 'authorization_code',
+        code,
+        client_id: id,
+        redirect_uri: redirectUris[0] ?? '',
+      });
+
+      // TODO: error handling
+      // Exchange the authorization code for an access token
+      const httpResponse = await got.post(tokenEndpoint, {
+        body: tokenRequestParameters.toString(),
+        headers,
+      });
+
+      const result = oidcTokenResponseGuard.safeParse(parseJson(httpResponse.body));
+
+      if (!result.success) {
+        throw new RequestError({
+          code: 'oidc.invalid_token',
+          status: 400,
+          message: 'Invalid token response',
+        });
+      }
+
+      const { accessToken } = camelcaseKeys(result.data);
+
+      assertThat(accessToken, new RequestError('oidc.access_denied'));
+
+      const userInfo = await getUserInfo(accessToken, userinfoEndpoint);
+
+      const { metadata } = await getSamlIdPMetadataByApplicationId(id);
+      const { privateKey } =
+        await samlApplicationSecrets.findActiveSamlApplicationSecretByApplicationId(id);
+      const { entityId, acsUrl } =
+        await samlApplicationConfigs.findSamlApplicationConfigByApplicationId(id);
+
+      assertThat(entityId && acsUrl, 'application.saml.entity_id_required');
+
+      // eslint-disable-next-line new-cap
+      const idp = saml.IdentityProvider({
+        metadata,
+        privateKey,
+        isAssertionEncrypted: false,
+        loginResponseTemplate: {
+          context: samlLogInResponseTemplate,
+          attributes: [
+            {
+              name: 'email',
+              valueTag: 'email',
+              nameFormat: 'urn:oasis:names:tc:SAML:2.0:attrname-format:basic',
+              valueXsiType: 'xs:string',
+            },
+            {
+              name: 'name',
+              valueTag: 'name',
+              nameFormat: 'urn:oasis:names:tc:SAML:2.0:attrname-format:basic',
+              valueXsiType: 'xs:string',
+            },
+          ],
+        },
+      });
+
+      const binding = acsUrl.binding ?? BindingType.Post;
+
+      // eslint-disable-next-line new-cap
+      const sp = saml.ServiceProvider({
+        entityID: entityId,
+        assertionConsumerService: [
+          {
+            Binding: binding,
+            Location: acsUrl.url,
+          },
+        ],
+      });
+
+      // TODO: fix binding method
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+      const { context, entityEndpoint } = await idp.createLoginResponse(
+        sp,
+        // @ts-expect-error --fix request object later
+        null,
+        'post',
+        userInfo,
+        createSamlTemplateCallback(idp, sp, userInfo)
+      );
+
+      ctx.body = `
+        <html>
+          <body>
+            <form id="redirectForm" action="${entityEndpoint}" method="POST">
+              <input type="hidden" name="SAMLResponse" value="${context}" />
+              <input type="submit" value="Submit" /> 
+            </form>
+            <script>
+              document.getElementById('redirectForm').submit();
+            </script>
+          </body>
+        </html>
+      `;
+
+      return next();
+    }
+  );
 }
diff --git a/packages/core/src/saml-applications/routes/utils.ts b/packages/core/src/saml-applications/routes/utils.ts
new file mode 100644
index 000000000..3d0e5f342
--- /dev/null
+++ b/packages/core/src/saml-applications/routes/utils.ts
@@ -0,0 +1,74 @@
+import { generateStandardId } from '@logto/shared';
+import saml from 'samlify';
+
+import { type idTokenProfileStandardClaims } from '#src/sso/types/oidc.js';
+
+export const samlLogInResponseTemplate = `
+<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" ID="{ID}" Version="2.0" IssueInstant="{IssueInstant}" Destination="{Destination}" InResponseTo="{InResponseTo}">
+  <saml:Issuer>{Issuer}</saml:Issuer>
+  <samlp:Status>
+    <samlp:StatusCode Value="{StatusCode}"/>
+  </samlp:Status>
+  <saml:Assertion ID="{AssertionID}" Version="2.0" IssueInstant="{IssueInstant}" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion">
+    <saml:Issuer>{Issuer}</saml:Issuer>
+    <saml:Subject>
+      <saml:NameID Format="{NameIDFormat}">{NameID}</saml:NameID>
+      <saml:SubjectConfirmation Method="urn:oasis:names:tc:SAML:2.0:cm:bearer">
+        <saml:SubjectConfirmationData NotOnOrAfter="{SubjectConfirmationDataNotOnOrAfter}" Recipient="{SubjectRecipient}" InResponseTo="{InResponseTo}"/>
+      </saml:SubjectConfirmation>
+    </saml:Subject>
+    <saml:Conditions NotBefore="{ConditionsNotBefore}" NotOnOrAfter="{ConditionsNotOnOrAfter}">
+      <saml:AudienceRestriction>
+        <saml:Audience>{Audience}</saml:Audience>
+      </saml:AudienceRestriction>
+    </saml:Conditions>
+    {AttributeStatement}
+  </saml:Assertion>
+</samlp:Response>`;
+
+export const createSamlTemplateCallback =
+  (
+    idp: saml.IdentityProviderInstance,
+    sp: saml.ServiceProviderInstance,
+    user: idTokenProfileStandardClaims
+  ) =>
+  (template: string) => {
+    const assertionConsumerServiceUrl = sp.entityMeta.getAssertionConsumerService(
+      saml.Constants.wording.binding.post
+    );
+
+    const { nameIDFormat } = idp.entitySetting;
+    const selectedNameIDFormat = Array.isArray(nameIDFormat) ? nameIDFormat[0] : nameIDFormat;
+
+    const id = `ID_${generateStandardId()}`;
+    const now = new Date();
+    const expireAt = new Date(now.getTime() + 10 * 60 * 1000); // 10 minutes later
+
+    const tagValues = {
+      ID: id,
+      AssertionID: `ID_${generateStandardId()}`,
+      Destination: assertionConsumerServiceUrl,
+      Audience: sp.entityMeta.getEntityID(),
+      EntityID: sp.entityMeta.getEntityID(),
+      SubjectRecipient: assertionConsumerServiceUrl,
+      Issuer: idp.entityMeta.getEntityID(),
+      IssueInstant: now.toISOString(),
+      AssertionConsumerServiceURL: assertionConsumerServiceUrl,
+      StatusCode: 'urn:oasis:names:tc:SAML:2.0:status:Success',
+      ConditionsNotBefore: now.toISOString(),
+      ConditionsNotOnOrAfter: expireAt.toISOString(),
+      SubjectConfirmationDataNotOnOrAfter: expireAt.toISOString(),
+      NameIDFormat: selectedNameIDFormat,
+      NameID: user.sub,
+      InResponseTo: 'null',
+      attrEmail: user.email,
+      attrName: user.name,
+    };
+
+    const context = saml.SamlLib.replaceTagsByValue(template, tagValues);
+
+    return {
+      id,
+      context,
+    };
+  };
diff --git a/packages/core/src/sso/types/oidc.ts b/packages/core/src/sso/types/oidc.ts
index 3125ea020..935109f9c 100644
--- a/packages/core/src/sso/types/oidc.ts
+++ b/packages/core/src/sso/types/oidc.ts
@@ -77,3 +77,5 @@ export const idTokenProfileStandardClaimsGuard = z.object({
   profile: z.string().nullish(),
   nonce: z.string().nullish(),
 });
+
+export type idTokenProfileStandardClaims = z.infer<typeof idTokenProfileStandardClaimsGuard>;