feat(server): add OIDC for AFFiNE (#6991)

Co-authored-by: LongYinan <lynweklm@gmail.com>
Co-authored-by: DarkSky <darksky2048@gmail.com>

packages/backend/server/src/plugins/oauth/providers/index.ts

@@ -1,4 +1,9 @@
 import { GithubOAuthProvider } from './github';
 import { GoogleOAuthProvider } from './google';
+import { OIDCProvider } from './oidc';
 
-export const OAuthProviders = [GoogleOAuthProvider, GithubOAuthProvider];
+export const OAuthProviders = [
+  GoogleOAuthProvider,
+  GithubOAuthProvider,
+  OIDCProvider,
+];

packages/backend/server/src/plugins/oauth/providers/oidc.ts

@@ -0,0 +1,213 @@
+import {
+  BadRequestException,
+  Injectable,
+  InternalServerErrorException,
+  OnModuleInit,
+} from '@nestjs/common';
+import { z } from 'zod';
+
+import { Config, URLHelper } from '../../../fundamentals';
+import { AutoRegisteredOAuthProvider } from '../register';
+import { OAuthOIDCProviderConfig, OAuthProviderName, OIDCArgs } from '../types';
+import { OAuthAccount, Tokens } from './def';
+
+const OIDCTokenSchema = z.object({
+  access_token: z.string(),
+  expires_in: z.number(),
+  refresh_token: z.string(),
+  scope: z.string(),
+  token_type: z.string(),
+});
+
+const OIDCUserInfoSchema = z.object({
+  id: z.string(),
+  email: z.string().email(),
+  name: z.string(),
+  groups: z.array(z.string()).optional(),
+});
+
+type OIDCUserInfo = z.infer<typeof OIDCUserInfoSchema>;
+
+const OIDCConfigurationSchema = z.object({
+  authorization_endpoint: z.string().url(),
+  token_endpoint: z.string().url(),
+  userinfo_endpoint: z.string().url(),
+  end_session_endpoint: z.string().url(),
+});
+
+type OIDCConfiguration = z.infer<typeof OIDCConfigurationSchema>;
+
+class OIDCClient {
+  private static async fetch<T = any>(
+    url: string,
+    options: RequestInit,
+    verifier: z.Schema<T>
+  ): Promise<T> {
+    const response = await fetch(url, options);
+
+    if (!response.ok) {
+      if (response.status >= 400 && response.status < 500) {
+        throw new BadRequestException(`Invalid OIDC configuration`, {
+          cause: await response.json(),
+          description: response.statusText,
+        });
+      } else {
+        throw new InternalServerErrorException(`Failed to configure client`, {
+          cause: await response.json(),
+          description: response.statusText,
+        });
+      }
+    }
+    return verifier.parse(response.json());
+  }
+
+  static async create(config: OAuthOIDCProviderConfig, url: URLHelper) {
+    const { args, clientId, clientSecret, issuer } = config;
+    if (!url.verify(issuer)) {
+      throw new Error('OIDC Issuer is invalid.');
+    }
+    const oidcConfig = await OIDCClient.fetch(
+      `${issuer}/.well-known/openid-configuration`,
+      {
+        method: 'GET',
+        headers: { Accept: 'application/json' },
+      },
+      OIDCConfigurationSchema
+    );
+
+    return new OIDCClient(clientId, clientSecret, args, oidcConfig, url);
+  }
+
+  private constructor(
+    private readonly clientId: string,
+    private readonly clientSecret: string,
+    private readonly args: OIDCArgs | undefined,
+    private readonly config: OIDCConfiguration,
+    private readonly url: URLHelper
+  ) {}
+
+  authorize(state: string): string {
+    const args = Object.assign({}, this.args);
+    if ('claim_id' in args) delete args.claim_id;
+    if ('claim_email' in args) delete args.claim_email;
+    if ('claim_name' in args) delete args.claim_name;
+
+    return `${this.config.authorization_endpoint}?${this.url.stringify({
+      client_id: this.clientId,
+      redirect_uri: this.url.link('/oauth/callback'),
+      response_type: 'code',
+      ...args,
+      scope: this.args?.scope || 'openid profile email',
+      state,
+    })}`;
+  }
+
+  async token(code: string): Promise<Tokens> {
+    const token = await OIDCClient.fetch(
+      this.config.token_endpoint,
+      {
+        method: 'POST',
+        body: this.url.stringify({
+          code,
+          client_id: this.clientId,
+          client_secret: this.clientSecret,
+          redirect_uri: this.url.link('/oauth/callback'),
+          grant_type: 'authorization_code',
+        }),
+        headers: {
+          Accept: 'application/json',
+          'Content-Type': 'application/x-www-form-urlencoded',
+        },
+      },
+      OIDCTokenSchema
+    );
+
+    return {
+      accessToken: token.access_token,
+      refreshToken: token.refresh_token,
+      expiresAt: new Date(Date.now() + token.expires_in * 1000),
+      scope: token.scope,
+    };
+  }
+
+  private mapUserInfo(
+    user: Record<string, any>,
+    claimsMap: Record<string, string>
+  ): OIDCUserInfo {
+    const mappedUser: Partial<OIDCUserInfo> = {};
+    for (const [key, value] of Object.entries(claimsMap)) {
+      if (user[value] !== undefined) {
+        mappedUser[key as keyof OIDCUserInfo] = user[value];
+      }
+    }
+    return mappedUser as OIDCUserInfo;
+  }
+
+  async userinfo(token: string) {
+    const user = await OIDCClient.fetch(
+      this.config.userinfo_endpoint,
+      {
+        method: 'GET',
+        headers: {
+          Accept: 'application/json',
+          Authorization: `Bearer ${token}`,
+        },
+      },
+      OIDCUserInfoSchema
+    );
+
+    const claimsMap = {
+      id: this.args?.claim_id || 'preferred_username',
+      email: this.args?.claim_email || 'email',
+      name: this.args?.claim_name || 'name',
+    };
+    const userinfo = this.mapUserInfo(user, claimsMap);
+    return { id: userinfo.id, email: userinfo.email };
+  }
+}
+
+@Injectable()
+export class OIDCProvider
+  extends AutoRegisteredOAuthProvider
+  implements OnModuleInit
+{
+  override provider = OAuthProviderName.OIDC;
+  private client: OIDCClient | null = null;
+
+  constructor(
+    protected readonly AFFiNEConfig: Config,
+    private readonly url: URLHelper
+  ) {
+    super();
+  }
+
+  override async onModuleInit() {
+    const config = this.optionalConfig as OAuthOIDCProviderConfig;
+    if (config && config.issuer && config.clientId && config.clientSecret) {
+      this.client = await OIDCClient.create(config, this.url);
+      super.onModuleInit();
+    }
+  }
+
+  private checkOIDCClient(
+    client: OIDCClient | null
+  ): asserts client is OIDCClient {
+    if (!client) {
+      throw new Error('OIDC client has not been loaded yet.');
+    }
+  }
+
+  getAuthUrl(state: string): string {
+    this.checkOIDCClient(this.client);
+    return this.client.authorize(state);
+  }
+
+  async getToken(code: string): Promise<Tokens> {
+    this.checkOIDCClient(this.client);
+    return await this.client.token(code);
+  }
+  async getUser(token: string): Promise<OAuthAccount> {
+    this.checkOIDCClient(this.client);
+    return await this.client.userinfo(token);
+  }
+}

packages/backend/server/src/plugins/oauth/types.ts

@@ -4,12 +4,31 @@ export interface OAuthProviderConfig {
   args?: Record<string, string>;
 }
 
+export type OIDCArgs = {
+  scope?: string;
+  claim_id?: string;
+  claim_email?: string;
+  claim_name?: string;
+};
+
+export interface OAuthOIDCProviderConfig extends OAuthProviderConfig {
+  issuer: string;
+  args?: OIDCArgs;
+}
+
 export enum OAuthProviderName {
   Google = 'google',
   GitHub = 'github',
+  OIDC = 'oidc',
 }
 
+type OAuthProviderConfigMapping = {
+  [OAuthProviderName.Google]: OAuthProviderConfig;
+  [OAuthProviderName.GitHub]: OAuthProviderConfig;
+  [OAuthProviderName.OIDC]: OAuthOIDCProviderConfig;
+};
+
 export interface OAuthConfig {
   enabled: boolean;
-  providers: Partial<{ [key in OAuthProviderName]: OAuthProviderConfig }>;
+  providers: Partial<OAuthProviderConfigMapping>;
 }