Queue integration tasks for import triggered via webhooks (#334)

* WIP

* task handler

* change limit

* changes for gitlab

* refactor

integrations/github/src/index.ts

@@ -21,8 +21,9 @@ import {
 import { configBlock } from './components';
 import { getGitHubAppJWT } from './provider';
 import { triggerExport, updateCommitWithPreviewLinks } from './sync';
-import type { GithubRuntimeContext } from './types';
-import { BRANCH_REF_PREFIX } from './utils';
+import { handleIntegrationTask } from './tasks';
+import type { GithubRuntimeContext, IntegrationTask } from './types';
+import { arrayToHex, BRANCH_REF_PREFIX, safeCompare } from './utils';
 import { handlePullRequestEvents, handlePushEvent, verifyGitHubWebhookSignature } from './webhooks';
 
 const logger = Logger('github');
@@ -38,6 +39,61 @@ const handleFetchEvent: FetchEventCallback<GithubRuntimeContext> = async (reques
         ).pathname,
     });
 
+    async function verifyIntegrationSignature(
+        payload: string,
+        signature: string,
+        secret: string
+    ): Promise<boolean> {
+        if (!signature) {
+            return false;
+        }
+
+        const algorithm = { name: 'HMAC', hash: 'SHA-256' };
+        const enc = new TextEncoder();
+        const key = await crypto.subtle.importKey('raw', enc.encode(secret), algorithm, false, [
+            'sign',
+            'verify',
+        ]);
+        const signed = await crypto.subtle.sign(algorithm.name, key, enc.encode(payload));
+        const expectedSignature = arrayToHex(signed);
+
+        return safeCompare(expectedSignature, signature);
+    }
+
+    /**
+     * Handle integration tasks
+     */
+    router.post('/tasks', async (request) => {
+        const signature = request.headers.get('x-gitbook-integration-signature') ?? '';
+        const payloadString = await request.text();
+
+        const verified = await verifyIntegrationSignature(
+            payloadString,
+            signature,
+            environment.signingSecret!
+        );
+
+        if (!verified) {
+            return new Response('Invalid integration signature', {
+                status: 400,
+            });
+        }
+
+        const { task } = JSON.parse(payloadString) as { task: IntegrationTask };
+        logger.debug('verified & received integration task', task);
+
+        context.waitUntil(
+            (async () => {
+                await handleIntegrationTask(context, task);
+            })()
+        );
+
+        return new Response(JSON.stringify({ acknowledged: true }), {
+            status: 200,
+            headers: { 'content-type': 'application/json' },
+        });
+    });
+
     /**
      * Handle GitHub App webhook events
      */

integrations/github/src/installation.ts

@@ -98,38 +98,35 @@ export async function saveSpaceConfiguration(
 }
 
 /**
- * List space installations that match the given external ID. It takes
- * care of pagination and returns all space installations at once.
+ * List space installations that match the given external ID.
  */
 export async function querySpaceInstallations(
     context: GithubRuntimeContext,
     externalId: string,
-    page?: string
-): Promise<Array<IntegrationSpaceInstallation>> {
+    options: {
+        page?: string;
+        limit?: number;
+    } = {}
+): Promise<{ data: Array<IntegrationSpaceInstallation>; nextPage?: string; total?: number }> {
     const { api, environment } = context;
+    const { page, limit = 100 } = options;
 
-    logger.debug(`Querying space installations for external ID ${externalId} (page: ${page ?? 1})`);
+    logger.debug(
+        `Querying space installations for external ID ${externalId} (${JSON.stringify(options)})`
+    );
 
     const { data } = await api.integrations.listIntegrationSpaceInstallations(
         environment.integration.name,
         {
-            limit: 100,
+            limit,
             externalId,
             page,
         }
     );
 
-    const spaceInstallations = [...data.items];
-
-    // Recursively fetch next pages
-    if (data.next) {
-        const nextSpaceInstallations = await querySpaceInstallations(
-            context,
-            externalId,
-            data.next.page
-        );
-        spaceInstallations.push(...nextSpaceInstallations);
-    }
-
-    return spaceInstallations;
+    return {
+        data: data.items,
+        total: data.count,
+        nextPage: data.next?.page,
+    };
 }

integrations/github/src/tasks.ts

@@ -0,0 +1,121 @@
+import { GitBookAPI } from '@gitbook/api';
+import { Logger } from '@gitbook/runtime';
+
+import { querySpaceInstallations } from './installation';
+import { triggerImport } from './sync';
+import type { GithubRuntimeContext, IntegrationTask, IntegrationTaskImportSpaces } from './types';
+
+const logger = Logger('github:tasks');
+
+/**
+ * Queue a task for the integration to import spaces.
+ */
+export async function queueTaskForImportSpaces(
+    context: GithubRuntimeContext,
+    task: IntegrationTaskImportSpaces
+): Promise<void> {
+    const { api, environment } = context;
+    await api.integrations.queueIntegrationTask(environment.integration.name, {
+        task: {
+            type: task.type,
+            payload: task.payload,
+        },
+    });
+}
+
+/**
+ * Handle an integration task.
+ */
+export async function handleIntegrationTask(
+    context: GithubRuntimeContext,
+    task: IntegrationTask
+): Promise<void> {
+    switch (task.type) {
+        case 'import:spaces':
+            await handleImportDispatchForSpaces(context, task.payload);
+            break;
+        default:
+            throw new Error(`Unknown integration task type: ${task}`);
+    }
+}
+
+/**
+ * This function is used to trigger an import for all the spaces that match the given config query.
+ * It will handle pagination by queueing itself if there are more spaces to import.
+ *
+ * `NOTE`: It is important that the total number of external network calls in this function is less
+ * than 50 as that is the limit imposed by Cloudflare workers.
+ */
+export async function handleImportDispatchForSpaces(
+    context: GithubRuntimeContext,
+    payload: IntegrationTaskImportSpaces['payload']
+): Promise<number | undefined> {
+    const { configQuery, page, standaloneRef } = payload;
+
+    logger.debug(`handling import dispatch for spaces with payload: ${JSON.stringify(payload)}`);
+
+    const {
+        data: spaceInstallations,
+        nextPage,
+        total,
+    } = await querySpaceInstallations(context, configQuery, {
+        limit: 10,
+        page,
+    });
+
+    await Promise.allSettled(
+        spaceInstallations.map(async (spaceInstallation) => {
+            try {
+                // Obtain the installation API token needed to trigger the import
+                const { data: installationAPIToken } =
+                    await context.api.integrations.createIntegrationInstallationToken(
+                        spaceInstallation.integration,
+                        spaceInstallation.installation
+                    );
+
+                // Set the token in the duplicated context to be used by the API client
+                const installationContext: GithubRuntimeContext = {
+                    ...context,
+                    api: new GitBookAPI({
+                        endpoint: context.environment.apiEndpoint,
+                        authToken: installationAPIToken.token,
+                    }),
+                    environment: {
+                        ...context.environment,
+                        authToken: installationAPIToken.token,
+                    },
+                };
+
+                await triggerImport(installationContext, spaceInstallation, {
+                    standalone: standaloneRef
+                        ? {
+                              ref: standaloneRef,
+                          }
+                        : undefined,
+                });
+            } catch (error) {
+                logger.error(
+                    `error while triggering ${
+                        standaloneRef ? `standalone (${standaloneRef})` : ''
+                    } import for space ${spaceInstallation.space}`,
+                    error
+                );
+            }
+        })
+    );
+
+    // Queue the next page if there is one
+    if (nextPage) {
+        logger.debug(`queueing next page ${nextPage} of import dispatch for spaces`);
+        await queueTaskForImportSpaces(context, {
+            type: 'import:spaces',
+            payload: {
+                page: nextPage,
+                configQuery,
+                standaloneRef,
+            },
+        });
+    }
+
+    return total;
+}

integrations/github/src/types.ts

@@ -79,3 +79,21 @@ export type GithubConfigureState = Omit<
     withCustomTemplate?: boolean;
     commitMessagePreview?: string;
 };
+
+export type IntegrationTaskType = 'import:spaces';
+
+export type BaseIntegrationTask<Type extends IntegrationTaskType, Payload extends object> = {
+    type: Type;
+    payload: Payload;
+};
+
+export type IntegrationTaskImportSpaces = BaseIntegrationTask<
+    'import:spaces',
+    {
+        configQuery: string;
+        page?: string;
+        standaloneRef?: string;
+    }
+>;
+
+export type IntegrationTask = IntegrationTaskImportSpaces;

integrations/github/src/utils.ts

@@ -83,3 +83,29 @@ export function assertIsDefined<T>(
         throw new Error(`Expected value (${options.label}) to be defined, but received ${value}`);
     }
 }
+
+/**
+ * Convert an array buffer to a hex string
+ */
+export function arrayToHex(arr: ArrayBuffer) {
+    return [...new Uint8Array(arr)].map((x) => x.toString(16).padStart(2, '0')).join('');
+}
+
+/**
+ * Constant-time string comparison. Equivalent of `crypto.timingSafeEqual`.
+ **/
+export function safeCompare(expected: string, actual: string) {
+    const lenExpected = expected.length;
+    let result = 0;
+
+    if (lenExpected !== actual.length) {
+        actual = expected;
+        result = 1;
+    }
+
+    for (let i = 0; i < lenExpected; i++) {
+        result |= expected.charCodeAt(i) ^ actual.charCodeAt(i);
+    }
+
+    return result === 0;
+}