From 102292ec4fadcd6aa53801f154bf434f3e5ac638 Mon Sep 17 00:00:00 2001
From: fenos <fabri.feno@gmail.com>
Date: Wed, 10 Jan 2024 13:40:38 +0000
Subject: [PATCH] feat: stale-while-revalidate

---
 src/config.ts                    |  4 +++
 src/storage/renderer/head.ts     | 27 +++++++++++++++--
 src/storage/renderer/renderer.ts | 50 ++++++++++++++++++++++++++++++--
 3 files changed, 76 insertions(+), 5 deletions(-)

diff --git a/src/config.ts b/src/config.ts
index 6e03015e..3d011bf4 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -85,6 +85,8 @@ type StorageConfigType = {
   tusPath: string
   tusUseFileVersionSeparator: boolean
   enableDefaultMetrics: boolean
+  sMaxAge: number
+  etagHeaders: string[]
 }
 
 function getOptionalConfigFromEnv(key: string): string | undefined {
@@ -246,5 +248,7 @@ export function getConfig(): StorageConfigType {
     tusUseFileVersionSeparator:
       getOptionalConfigFromEnv('TUS_USE_FILE_VERSION_SEPARATOR') === 'true',
     enableDefaultMetrics: getOptionalConfigFromEnv('ENABLE_DEFAULT_METRICS') === 'true',
+    sMaxAge: parseInt(getOptionalConfigFromEnv('S_MAXAGE') || '0', 10),
+    etagHeaders: getOptionalConfigFromEnv('ETAG_HEADERS')?.trim().split(',') || ['if-none-match'],
   }
 }
diff --git a/src/storage/renderer/head.ts b/src/storage/renderer/head.ts
index 2b17f4d0..2f9c78cd 100644
--- a/src/storage/renderer/head.ts
+++ b/src/storage/renderer/head.ts
@@ -1,6 +1,6 @@
 import { AssetResponse, Renderer, RenderOptions } from './renderer'
-import { FastifyRequest } from 'fastify'
-import { StorageBackendAdapter } from '../backend'
+import { FastifyReply, FastifyRequest } from 'fastify'
+import { ObjectMetadata, StorageBackendAdapter } from '../backend'
 import { ImageRenderer, TransformOptions } from './image'
 
 /**
@@ -20,4 +20,27 @@ export class HeadRenderer extends Renderer {
       transformations: ImageRenderer.applyTransformation(request.query as TransformOptions),
     }
   }
+
+  protected handleCacheControl(
+    request: FastifyRequest<any>,
+    response: FastifyReply<any>,
+    metadata: ObjectMetadata
+  ) {
+    const etag = this.findEtagHeader(request)
+
+    const cacheControl = [metadata.cacheControl]
+
+    if (!etag) {
+      response.header('Cache-Control', cacheControl.join(', '))
+      return
+    }
+
+    if (etag !== metadata.eTag) {
+      cacheControl.push('must-revalidate')
+    } else if (this.sMaxAge > 0) {
+      cacheControl.push(`s-maxage=${this.sMaxAge}`)
+    }
+
+    response.header('Cache-Control', cacheControl.join(', '))
+  }
 }
diff --git a/src/storage/renderer/renderer.ts b/src/storage/renderer/renderer.ts
index 03da06a1..fc004819 100644
--- a/src/storage/renderer/renderer.ts
+++ b/src/storage/renderer/renderer.ts
@@ -1,6 +1,7 @@
 import { FastifyReply, FastifyRequest } from 'fastify'
 import { ObjectMetadata } from '../backend'
 import { Readable } from 'stream'
+import { getConfig } from '../../config'
 
 export interface RenderOptions {
   bucket: string
@@ -16,12 +17,16 @@ export interface AssetResponse {
   transformations?: string[]
 }
 
+const { etagHeaders, sMaxAge } = getConfig()
+
 /**
  * Renderer
  * a generic renderer that respond to a request with an asset content
  * and all the important headers
  */
 export abstract class Renderer {
+  protected sMaxAge = sMaxAge
+
   abstract getAsset(request: FastifyRequest, options: RenderOptions): Promise<AssetResponse>
 
   /**
@@ -34,7 +39,7 @@ export abstract class Renderer {
     try {
       const data = await this.getAsset(request, options)
 
-      await this.setHeaders(response, data, options)
+      this.setHeaders(request, response, data, options)
 
       return response.send(data.body)
     } catch (err: any) {
@@ -55,7 +60,12 @@ export abstract class Renderer {
     }
   }
 
-  protected setHeaders(response: FastifyReply<any>, data: AssetResponse, options: RenderOptions) {
+  protected setHeaders(
+    request: FastifyRequest<any>,
+    response: FastifyReply<any>,
+    data: AssetResponse,
+    options: RenderOptions
+  ) {
     response
       .status(data.metadata.httpStatusCode ?? 200)
       .header('Accept-Ranges', 'bytes')
@@ -67,7 +77,7 @@ export abstract class Renderer {
     if (options.expires) {
       response.header('Expires', options.expires)
     } else {
-      response.header('Cache-Control', data.metadata.cacheControl)
+      this.handleCacheControl(request, response, data.metadata)
     }
 
     if (data.metadata.contentRange) {
@@ -95,6 +105,40 @@ export abstract class Renderer {
       }
     }
   }
+
+  protected handleCacheControl(
+    request: FastifyRequest<any>,
+    response: FastifyReply<any>,
+    metadata: ObjectMetadata
+  ) {
+    const etag = this.findEtagHeader(request)
+
+    const cacheControl = [metadata.cacheControl]
+
+    if (!etag) {
+      response.header('Cache-Control', cacheControl.join(', '))
+      return
+    }
+
+    if (this.sMaxAge > 0) {
+      cacheControl.push(`s-maxage=${this.sMaxAge}`)
+    }
+
+    if (etag !== metadata.eTag) {
+      cacheControl.push('stale-while-revalidate=30')
+    }
+
+    response.header('Cache-Control', cacheControl.join(', '))
+  }
+
+  protected findEtagHeader(request: FastifyRequest<any>) {
+    for (const header of etagHeaders) {
+      const etag = request.headers[header]
+      if (etag) {
+        return etag
+      }
+    }
+  }
 }
 
 function normalizeContentType(contentType: string | undefined): string | undefined {