Add ability to set custom headers in hook function

neg4n · Jul 26, 2022 · 6168b28 · 6168b28
1 parent 58e72bb
commit 6168b28
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 39 deletions.
diff --git a/README.md b/README.md
@@ -218,6 +218,8 @@ const nextApiOgImageConfig = {
   cacheControl: 'max-age 3600, must-revalidate',
   // Hook function that allows to intercept inner NextApiRequest with `ogImage` prop attached.
   // useful for e.g. saving image in the database after the generation.
+  // The hook function return is Map containing custom headers that will be set BEFORE sending
+  // response to the client.
   hook: null,
   // NOTE: Options within 'dev' object works only when process.env.NODE_ENV === 'development'
   dev: {

diff --git a/src/index.ts b/src/index.ts
@@ -19,7 +19,7 @@ type ImageType = 'png' | 'jpeg' | 'webp'
 type StrategyAwareParams<
   T extends StrategyOption = 'query',
   StrategyDetails extends string | object = string,
-  > = T extends 'body'
+> = T extends 'body'
   ? StrategyDetails
   : Record<StrategyDetails extends string ? StrategyDetails : string, NonNullable<string>>
 
@@ -30,26 +30,26 @@ type NextApiRequestWithOgImage = {
 export type NextApiOgImageConfig<
   Strategy extends StrategyOption,
   StrategyDetails extends string | object = string,
-  > = {
-    template: RequireExactlyOne<
-      Partial<{
-        html: (params: StrategyAwareParams<Strategy, StrategyDetails>) => string | Promise<string>
-        react: (params: StrategyAwareParams<Strategy, StrategyDetails>) => ReactElement | Promise<ReactElement>
-      }>,
-      'html' | 'react'
-    >
-    strategy?: StrategyOption
-    cacheControl?: string
-    width?: number
-    height?: number
-    type?: ImageType
-    quality?: number
-    hook?: (request: NextApiRequestWithOgImage) => void | Promise<void>,
-    dev?: Partial<{
-      inspectHtml: boolean
-      errorsInResponse: boolean
-    }>
-  }
+> = {
+  template: RequireExactlyOne<
+    Partial<{
+      html: (params: StrategyAwareParams<Strategy, StrategyDetails>) => string | Promise<string>
+      react: (params: StrategyAwareParams<Strategy, StrategyDetails>) => ReactElement | Promise<ReactElement>
+    }>,
+    'html' | 'react'
+  >
+  strategy?: StrategyOption
+  cacheControl?: string
+  width?: number
+  height?: number
+  type?: ImageType
+  quality?: number
+  hook?: (request: NextApiRequestWithOgImage) => Map<string, string> | Promise<Map<string, string>>
+  dev?: Partial<{
+    inspectHtml: boolean
+    errorsInResponse: boolean
+  }>
+}
 
 type BrowserEnvironment = {
   envMode: EnvMode
@@ -61,7 +61,7 @@ type BrowserEnvironment = {
 export function withOGImage<
   Strategy extends StrategyOption = 'query',
   StrategyDetails extends string | object = string,
-  >(options: NextApiOgImageConfig<Strategy, StrategyDetails>) {
+>(options: NextApiOgImageConfig<Strategy, StrategyDetails>) {
   const defaultOptions: Except<NextApiOgImageConfig<Strategy, StrategyDetails>, 'template'> = {
     strategy: 'query',
     cacheControl: 'max-age 3600, must-revalidate',
@@ -106,7 +106,7 @@ export function withOGImage<
     createImageFactory({ inspectHtml, type, quality }),
   )
 
-  return async function(request: NextApiRequest, response: NextApiResponse) {
+  return async function (request: NextApiRequest, response: NextApiResponse) {
     checkStrategy(strategy, !isProductionLikeMode(envMode) ? errorsInResponse : false, request, response)
 
     const params = stringifyObjectProps(strategy === 'query' ? request.query : request.body)
@@ -116,26 +116,31 @@ export function withOGImage<
       htmlTemplate && !reactTemplate
         ? await htmlTemplate({ ...params } as StrategyAwareParams<Strategy, StrategyDetails>)
         : renderToStaticMarkup(
-          await reactTemplate({ ...params } as StrategyAwareParams<Strategy, StrategyDetails>),
-        )
+            await reactTemplate({ ...params } as StrategyAwareParams<Strategy, StrategyDetails>),
+          )
 
-    const image = await browserEnvironment.createImage(emojify(html));
+    const image = await browserEnvironment.createImage(emojify(html))
 
     if (!!hook) {
       const extendedRequest: NextApiRequestWithOgImage = {
         ...request,
-        image
+        image,
       }
 
-      await hook(extendedRequest)
+      const headers = await hook(extendedRequest)
+      if (!!headers) {
+        for (const [extendedheaderName, extendedHeaderValue] of headers.entries()) {
+          response.setHeader(extendedheaderName, extendedHeaderValue)
+        }
+      }
     }
 
     response.setHeader(
       'Content-Type',
       !isProductionLikeMode(envMode) && inspectHtml ? 'text/html' : type ? `image/${type}` : 'image/png',
     )
     response.setHeader('Cache-Control', cacheControl)
-    response.write(image);
+    response.write(image)
     response.end()
   }
 }
@@ -212,7 +217,7 @@ function emojify(html: string) {
 }
 
 function pipe(...functions: Array<Function>): () => Promise<BrowserEnvironment> {
-  return async function() {
+  return async function () {
     return await functions.reduce(
       async (acc, fn) => await fn(await acc),
       Promise.resolve({ envMode: process.env.NODE_ENV as EnvMode } as BrowserEnvironment),
@@ -225,14 +230,14 @@ function getChromiumExecutable(browserEnvironment: BrowserEnvironment) {
     process.platform === 'win32'
       ? 'C:\\Program Files (x86)\\Google\\Chrome\\Application\\chrome.exe'
       : process.platform === 'linux'
-        ? '/usr/bin/google-chrome'
-        : '/Applications/Google Chrome.app/Contents/MacOS/Google Chrome'
+      ? '/usr/bin/google-chrome'
+      : '/Applications/Google Chrome.app/Contents/MacOS/Google Chrome'
 
   return { ...browserEnvironment, executable }
 }
 
 function prepareWebPageFactory(viewPort: Viewport) {
-  return async function(browserEnvironment: BrowserEnvironment) {
+  return async function (browserEnvironment: BrowserEnvironment) {
     const { page, envMode, executable } = browserEnvironment
 
     if (page) {
@@ -242,10 +247,10 @@ function prepareWebPageFactory(viewPort: Viewport) {
     const chromiumOptions = !isProductionLikeMode(envMode)
       ? { args: [], executablePath: executable, headless: true }
       : {
-        args: chrome.args,
-        executablePath: await chrome.executablePath,
-        headless: chrome.headless,
-      }
+          args: chrome.args,
+          executablePath: await chrome.executablePath,
+          headless: chrome.headless,
+        }
 
     const browser = await core.launch(chromiumOptions)
     const newPage = await browser.newPage()
@@ -264,12 +269,12 @@ function createImageFactory({
   type: ImageType
   quality: number
 }) {
-  return function(browserEnvironment: BrowserEnvironment) {
+  return function (browserEnvironment: BrowserEnvironment) {
     const { page, envMode } = browserEnvironment
 
     return {
       ...browserEnvironment,
-      createImage: async function(html: string) {
+      createImage: async function (html: string) {
         await page.setContent(html)
         const file =
           !isProductionLikeMode(envMode) && inspectHtml