From 3f593ff8c6a27bcd8e864d48e3635b6e93d6857c Mon Sep 17 00:00:00 2001
From: Steven <steven@ceriously.com>
Date: Wed, 2 Feb 2022 16:27:56 -0500
Subject: [PATCH] Fix image cache race condition (#33883)

- Fixes #33860
---
 packages/next/server/image-optimizer.ts | 55 +++++++++++++++----------
 1 file changed, 34 insertions(+), 21 deletions(-)

diff --git a/packages/next/server/image-optimizer.ts b/packages/next/server/image-optimizer.ts
index 0d603d546ab58b1..8363070892c600f 100644
--- a/packages/next/server/image-optimizer.ts
+++ b/packages/next/server/image-optimizer.ts
@@ -19,6 +19,8 @@ import { getContentType, getExtension } from './serve-static'
 import chalk from 'next/dist/compiled/chalk'
 import { NextUrlWithParsedQuery } from './request-meta'
 
+type XCacheHeader = 'MISS' | 'HIT' | 'STALE'
+
 const AVIF = 'image/avif'
 const WEBP = 'image/webp'
 const PNG = 'image/png'
@@ -29,7 +31,7 @@ const CACHE_VERSION = 3
 const ANIMATABLE_TYPES = [WEBP, PNG, GIF]
 const VECTOR_TYPES = [SVG]
 const BLUR_IMG_SIZE = 8 // should match `next-image-loader`
-const inflightRequests = new Map<string, Promise<undefined>>()
+const inflightRequests = new Map<string, Promise<void>>()
 
 let sharp:
   | ((
@@ -178,18 +180,15 @@ export async function imageOptimizer(
   const imagesDir = join(distDir, 'cache', 'images')
   const hashDir = join(imagesDir, hash)
   const now = Date.now()
-  let staleWhileRevalidate = false
+  let xCache: XCacheHeader = 'MISS'
 
   // If there're concurrent requests hitting the same resource and it's still
   // being optimized, wait before accessing the cache.
   if (inflightRequests.has(hash)) {
     await inflightRequests.get(hash)
   }
-  let dedupeResolver: (val?: PromiseLike<undefined>) => void
-  inflightRequests.set(
-    hash,
-    new Promise((resolve) => (dedupeResolver = resolve))
-  )
+  const dedupe = new Deferred<void>()
+  inflightRequests.set(hash, dedupe.promise)
 
   try {
     if (await fileExists(hashDir, 'directory')) {
@@ -200,8 +199,7 @@ export async function imageOptimizer(
         const expireAt = Number(expireAtSt)
         const contentType = getContentType(extension)
         const fsPath = join(hashDir, file)
-        const isFresh = now < expireAt
-        const xCache = isFresh ? 'HIT' : 'STALE'
+        xCache = now < expireAt ? 'HIT' : 'STALE'
         const result = setResponseHeaders(
           req,
           res,
@@ -214,13 +212,17 @@ export async function imageOptimizer(
           xCache
         )
         if (!result.finished) {
-          createReadStream(fsPath).pipe(res)
+          await new Promise<void>((resolve, reject) => {
+            createReadStream(fsPath)
+              .on('end', resolve)
+              .on('error', reject)
+              .pipe(res)
+          })
         }
-        if (isFresh) {
+        if (xCache === 'HIT') {
           return { finished: true }
         } else {
           await promises.unlink(fsPath)
-          staleWhileRevalidate = true
         }
       }
     }
@@ -338,7 +340,7 @@ export async function imageOptimizer(
           upstreamBuffer,
           isStatic,
           isDev,
-          staleWhileRevalidate
+          xCache
         )
         return { finished: true }
       }
@@ -492,7 +494,7 @@ export async function imageOptimizer(
           optimizedBuffer,
           isStatic,
           isDev,
-          staleWhileRevalidate
+          xCache
         )
       } else {
         throw new Error('Unable to optimize buffer')
@@ -507,14 +509,13 @@ export async function imageOptimizer(
         upstreamBuffer,
         isStatic,
         isDev,
-        staleWhileRevalidate
+        xCache
       )
     }
 
     return { finished: true }
   } finally {
-    // Make sure to remove the hash in the end.
-    dedupeResolver!()
+    dedupe.resolve()
     inflightRequests.delete(hash)
   }
 }
@@ -557,7 +558,7 @@ function setResponseHeaders(
   contentType: string | null,
   isStatic: boolean,
   isDev: boolean,
-  xCache: 'MISS' | 'HIT' | 'STALE'
+  xCache: XCacheHeader
 ) {
   res.setHeader('Vary', 'Accept')
   res.setHeader(
@@ -597,12 +598,11 @@ function sendResponse(
   buffer: Buffer,
   isStatic: boolean,
   isDev: boolean,
-  staleWhileRevalidate: boolean
+  xCache: XCacheHeader
 ) {
-  if (staleWhileRevalidate) {
+  if (xCache === 'STALE') {
     return
   }
-  const xCache = 'MISS'
   const etag = getHash([buffer])
   const result = setResponseHeaders(
     req,
@@ -782,3 +782,16 @@ export async function getImageSize(
   const { width, height } = imageSizeOf(buffer)
   return { width, height }
 }
+
+export class Deferred<T> {
+  promise: Promise<T>
+  resolve!: (value: T) => void
+  reject!: (error?: Error) => void
+
+  constructor() {
+    this.promise = new Promise((resolve, reject) => {
+      this.resolve = resolve
+      this.reject = reject
+    })
+  }
+}