Change Swagger auth (#264)

datalens-tech · Feb 13, 2025 · cb0d464 · cb0d464
1 parent 019c80a
commit cb0d464
Show file tree

Hide file tree

Showing 6 changed files with 134 additions and 68 deletions.
diff --git a/api/components.ts b/api/components.ts
@@ -28,6 +28,11 @@ export {
     initSwagger,
     ApiTag,
     getAdditionalHeaders,
+    getAdditionalSecuritySchemes,
+    SecurityType,
+    SecuritySchemeObject,
+    GetAdditionalHeadersResult,
+    GetAdditionalSecuritySchemesResult,
 } from '../src/components/api-docs';
 export {
     makeParser,

diff --git a/src/components/api-docs/tags.ts → src/components/api-docs/constants.ts b/src/components/api-docs/tags.ts → src/components/api-docs/constants.ts
@@ -1,3 +1,9 @@
+export enum SecurityType {
+    MasterToken = 'Master token',
+    BearerAuth = 'Access token',
+    ServiceUserToken = 'Service user access token',
+}
+
 export enum ApiTag {
     Collections = 'Collections',
     Workbooks = 'Workbooks',

diff --git a/src/components/api-docs/index.ts b/src/components/api-docs/index.ts
@@ -1,24 +1,62 @@
 import {OpenAPIRegistry, OpenApiGeneratorV31} from '@asteasolutions/zod-to-openapi';
-import {AuthPolicy, ExpressKit} from '@gravity-ui/expresskit';
+import {ExpressKit} from '@gravity-ui/expresskit';
 import swaggerUi from 'swagger-ui-express';
-import {ZodType, z} from 'zod';
+import {ZodType} from 'zod';
 
-import {US_MASTER_TOKEN_HEADER} from '../../const';
 import type {ExtendedAppRouteDescription} from '../../routes';
 
-import type {Method} from './types';
+import type {GetAdditionalHeadersResult, GetAdditionalSecuritySchemesResult, Method} from './types';
 
-export {ApiTag} from './tags';
-export * from './utils';
+export {getAdditionalHeaders, getAdditionalSecuritySchemes} from './utils';
+export {ApiTag, SecurityType} from './constants';
+export {
+    SecuritySchemeObject,
+    GetAdditionalHeadersResult,
+    GetAdditionalSecuritySchemesResult,
+} from './types';
 
 const openApiRegistry = new OpenAPIRegistry();
 
+export const initSwagger = (
+    app: ExpressKit,
+    securitySchemes?: GetAdditionalSecuritySchemesResult,
+) => {
+    const {config} = app;
+
+    const installationText = `Installation – <b>${config.appInstallation}</b>`;
+    const envText = `Env – <b>${config.appEnv}</b>`;
+    const descriptionText = `<br />Storage for DataLens entities.`;
+
+    setImmediate(() => {
+        if (securitySchemes) {
+            Object.keys(securitySchemes).forEach((securityType) => {
+                openApiRegistry.registerComponent('securitySchemes', securityType, {
+                    ...securitySchemes[securityType],
+                });
+            });
+        }
+
+        app.express.use(
+            '/api-docs',
+            swaggerUi.serve,
+            swaggerUi.setup(
+                new OpenApiGeneratorV31(openApiRegistry.definitions).generateDocument({
+                    openapi: '3.1.0',
+                    info: {
+                        version: `${config.appVersion}`,
+                        title: `United Storage `,
+                        description: [installationText, envText, descriptionText].join('<br />'),
+                    },
+                    servers: [{url: '/'}],
+                }),
+            ),
+        );
+    });
+};
+
 export const registerApiRoute = (
     routeDescription: ExtendedAppRouteDescription<unknown>,
-    {
-        headers: additionalHeaders,
-        security: additionalSecurity,
-    }: {headers: ZodType<unknown>[]; security: {[key: string]: any}[]},
+    {headers: additionalHeaders, security: additionalSecurity}: GetAdditionalHeadersResult,
 ) => {
     const {route, handler} = routeDescription;
     const {api} = handler;
@@ -46,14 +84,6 @@ export const registerApiRoute = (
 
         const headers: ZodType<unknown>[] = [];
 
-        if (routeDescription.private) {
-            headers.push(
-                z.strictObject({
-                    [US_MASTER_TOKEN_HEADER]: z.string(),
-                }),
-            );
-        }
-
         if (additionalHeaders) {
             headers.push(...additionalHeaders);
         }
@@ -79,39 +109,3 @@ export const registerApiRoute = (
         });
     }
 };
-
-export const initSwagger = (app: ExpressKit) => {
-    const {config} = app;
-
-    const installationText = `Installation – <b>${config.appInstallation}</b>`;
-    const envText = `Env – <b>${config.appEnv}</b>`;
-    const descriptionText = `<br />Storage for DataLens entities.`;
-
-    setImmediate(() => {
-        const authDisabled = config.appAuthPolicy === AuthPolicy.disabled;
-
-        if (!authDisabled && (config.zitadelEnabled || config.isAuthEnabled)) {
-            openApiRegistry.registerComponent('securitySchemes', 'bearerAuth', {
-                type: 'http',
-                scheme: 'bearer',
-                bearerFormat: 'JWT',
-            });
-        }
-
-        app.express.use(
-            '/api-docs',
-            swaggerUi.serve,
-            swaggerUi.setup(
-                new OpenApiGeneratorV31(openApiRegistry.definitions).generateDocument({
-                    openapi: '3.1.0',
-                    info: {
-                        version: `${config.appVersion}`,
-                        title: `United Storage `,
-                        description: [installationText, envText, descriptionText].join('<br />'),
-                    },
-                    servers: [{url: '/'}],
-                }),
-            ),
-        );
-    });
-};
diff --git a/src/components/api-docs/types.ts b/src/components/api-docs/types.ts
@@ -1,2 +1,23 @@
+import type {ZodType} from 'zod';
+
 // Copied from @asteasolutions/zod-to-openapi
 export type Method = 'get' | 'post' | 'put' | 'delete' | 'patch' | 'head' | 'options' | 'trace';
+
+export type SecuritySchemeType = 'apiKey' | 'http' | 'oauth2' | 'openIdConnect';
+
+export type SecuritySchemeObject = {
+    type: SecuritySchemeType;
+    description?: string;
+    name?: string;
+    in?: string;
+    scheme?: string;
+    bearerFormat?: string;
+    openIdConnectUrl?: string;
+};
+
+export type GetAdditionalHeadersResult = {
+    headers: ZodType<unknown>[];
+    security: {[key: string]: []}[];
+};
+
+export type GetAdditionalSecuritySchemesResult = Record<string, SecuritySchemeObject>;
diff --git a/src/components/api-docs/utils.ts b/src/components/api-docs/utils.ts
@@ -1,17 +1,56 @@
 import {AuthPolicy} from '@gravity-ui/expresskit';
 import {NodeKit} from '@gravity-ui/nodekit';
-import {ZodType} from 'zod';
 
-import {DL_SERVICE_USER_ACCESS_TOKEN} from '../../const';
+import {DL_SERVICE_USER_ACCESS_TOKEN, US_MASTER_TOKEN_HEADER} from '../../const';
 import type {ExtendedAppRouteDescription} from '../../routes';
-import {z} from '../zod';
+
+import {SecurityType} from './constants';
+import type {GetAdditionalHeadersResult, GetAdditionalSecuritySchemesResult} from './types';
+
+export const getAdditionalSecuritySchemes = (
+    nodekit: NodeKit,
+): GetAdditionalSecuritySchemesResult => {
+    const result: GetAdditionalSecuritySchemesResult = {};
+
+    result[SecurityType.MasterToken] = {
+        type: 'apiKey',
+        in: 'header',
+        name: US_MASTER_TOKEN_HEADER,
+    };
+
+    const {config} = nodekit;
+
+    const authDisabled = config.appAuthPolicy === AuthPolicy.disabled;
+
+    if (!authDisabled && (config.zitadelEnabled || config.isAuthEnabled)) {
+        result[SecurityType.BearerAuth] = {
+            type: 'http',
+            scheme: 'bearer',
+            bearerFormat: 'JWT',
+        };
+
+        if (config.zitadelEnabled) {
+            result[SecurityType.ServiceUserToken] = {
+                type: 'apiKey',
+                in: 'header',
+                name: DL_SERVICE_USER_ACCESS_TOKEN,
+            };
+        }
+    }
+
+    return result;
+};
 
 export const getAdditionalHeaders = (
     routeDescription: ExtendedAppRouteDescription,
     nodekit: NodeKit,
-) => {
-    const headers: ZodType<unknown>[] = [];
-    const security: {[key: string]: unknown}[] = [];
+): GetAdditionalHeadersResult => {
+    const headers: GetAdditionalHeadersResult['headers'] = [];
+    const security: GetAdditionalHeadersResult['security'] = [];
+
+    if (routeDescription.private) {
+        security.push({[SecurityType.MasterToken]: []});
+    }
 
     const {config} = nodekit;
 
@@ -20,14 +59,10 @@ export const getAdditionalHeaders = (
         routeDescription.authPolicy === AuthPolicy.disabled;
 
     if (!authDisabled && (config.zitadelEnabled || config.isAuthEnabled)) {
-        security.push({bearerAuth: []});
+        security.push({[SecurityType.BearerAuth]: []});
 
         if (config.zitadelEnabled) {
-            headers.push(
-                z.strictObject({
-                    ...(config.zitadelEnabled ? {[DL_SERVICE_USER_ACCESS_TOKEN]: z.string()} : {}),
-                }),
-            );
+            security.push({[SecurityType.ServiceUserToken]: []});
         }
     }
 

diff --git a/src/index.ts b/src/index.ts
@@ -21,7 +21,12 @@ import {registry} from './registry';
 import {getRoutes} from './routes';
 import {setRegistryToContext} from './components/app-context';
 import {isEnabledFeature} from './components/features';
-import {getAdditionalHeaders, initSwagger, registerApiRoute} from './components/api-docs';
+import {
+    getAdditionalHeaders,
+    getAdditionalSecuritySchemes,
+    initSwagger,
+    registerApiRoute,
+} from './components/api-docs';
 import {objectKeys} from './utils/utility-types';
 
 setRegistryToContext(nodekit, registry);
@@ -90,7 +95,7 @@ const app = new ExpressKit(nodekit, routes);
 registry.setupApp(app);
 
 if (nodekit.config.swaggerEnabled) {
-    initSwagger(app);
+    initSwagger(app, getAdditionalSecuritySchemes(nodekit));
 }
 
 if (require.main === module) {