From 24d2735b4c86d852efb769596dd829be66bccf76 Mon Sep 17 00:00:00 2001
From: Nick Taylor <nick@iamdeveloper.com>
Date: Wed, 2 Aug 2023 16:30:52 -0400
Subject: [PATCH] test: added tests for WAF fix

---
 src/index.ts       |   6 +--
 test/index.test.ts | 124 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 126 insertions(+), 4 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 6310c59..8ef4422 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -4,7 +4,7 @@ import { createIPX, handleRequest, IPXOptions } from 'ipx'
 import { builder, Handler } from '@netlify/functions'
 import { parseURL } from 'ufo'
 import etag from 'etag'
-import { loadSourceImage } from './http'
+import { loadSourceImage as defaultLoadSourceImage } from './http'
 import { decodeBase64Params, doPatternsMatchUrl, RemotePattern } from './utils'
 
 // WAF is Web Application Firewall
@@ -46,7 +46,7 @@ const plainText = {
   'Content-Type': 'text/plain'
 }
 
-export function createIPXHandler({
+export function createIPXHandler ({
   cacheDir = join(tmpdir(), 'ipx-cache'),
   basePath = '/_ipx/',
   propsEncoding,
@@ -55,7 +55,7 @@ export function createIPXHandler({
   responseHeaders,
   localPrefix,
   ...opts
-}: IPXHandlerOptions = {}) {
+}: IPXHandlerOptions = {}, loadSourceImage = defaultLoadSourceImage) {
   const ipx = createIPX({ ...opts, dir: join(cacheDir, 'cache') })
   if (!basePath.endsWith('/')) {
     basePath = `${basePath}/`
diff --git a/test/index.test.ts b/test/index.test.ts
index 148bb01..8523206 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -5,7 +5,7 @@ import test from 'ava'
 import { readFile, statSync, emptyDir, readdirSync } from 'fs-extra'
 
 import { createIPXHandler } from '../src/index'
-import { CACHE_PRUNING_THRESHOLD } from '../src/http'
+import { CACHE_PRUNING_THRESHOLD, SourceImageResult } from '../src/http'
 
 test('source image cache pruning', async (t) => {
   const filePath = join(__dirname, '..', 'example', 'public', 'img', 'test.jpg')
@@ -96,3 +96,125 @@ test('source image cache pruning', async (t) => {
     'cache size should not be equal to number of images * image size if we exceed threshold'
   )
 })
+
+test('should add WAF headers to local images being transformed', async (t) => {
+  const handler = createIPXHandler({
+    basePath: '/_ipx/',
+    cacheDir: '/tmp/ipx-cache',
+    bypassDomainCheck: true
+  }, (sourceImageOptions) => {
+    t.assert(sourceImageOptions.requestHeaders && sourceImageOptions.requestHeaders['X-Nf-Waf-Bypass-Token'] === 'some token')
+
+    return Promise.resolve({ finalize: () => { } } as SourceImageResult)
+  })
+
+  await handler(
+    {
+      rawUrl: 'http://localhost:3000/some-path',
+      path: '/_ipx/w_500/no-file.jpg',
+      headers: { 'X-Nf-Waf-Bypass-Token': 'some token' },
+      rawQuery: '',
+      httpMethod: 'GET',
+      queryStringParameters: {},
+      multiValueQueryStringParameters: {},
+      multiValueHeaders: {},
+      isBase64Encoded: false,
+      body: null
+    },
+    {
+      functionName: 'ipx',
+      callbackWaitsForEmptyEventLoop: false,
+      functionVersion: '1',
+      invokedFunctionArn: '',
+      awsRequestId: '',
+      logGroupName: '',
+      logStreamName: '',
+      memoryLimitInMB: '',
+      getRemainingTimeInMillis: () => 1000,
+      done: () => { },
+      fail: () => { },
+      succeed: () => { }
+    }
+  )
+})
+
+test('should not add WAF headers to local images being transformed', async (t) => {
+  const handler = createIPXHandler({
+    basePath: '/_ipx/',
+    cacheDir: '/tmp/ipx-cache',
+    bypassDomainCheck: true
+  }, (sourceImageOptions) => {
+    t.assert(sourceImageOptions.requestHeaders && sourceImageOptions.requestHeaders['X-Nf-Waf-Bypass-Token'] === undefined)
+
+    return Promise.resolve({ finalize: () => { } } as SourceImageResult)
+  })
+
+  await handler(
+    {
+      rawUrl: 'http://localhost:3000/some-path',
+      path: '/_ipx/w_500/https%3A%2F%2Fsome-site.com%2Fno-file.jpg',
+      headers: { 'X-Nf-Waf-Bypass-Token': 'some token' },
+      rawQuery: '',
+      httpMethod: 'GET',
+      queryStringParameters: {},
+      multiValueQueryStringParameters: {},
+      multiValueHeaders: {},
+      isBase64Encoded: false,
+      body: null
+    },
+    {
+      functionName: 'ipx',
+      callbackWaitsForEmptyEventLoop: false,
+      functionVersion: '1',
+      invokedFunctionArn: '',
+      awsRequestId: '',
+      logGroupName: '',
+      logStreamName: '',
+      memoryLimitInMB: '',
+      getRemainingTimeInMillis: () => 1000,
+      done: () => { },
+      fail: () => { },
+      succeed: () => { }
+    }
+  )
+})
+test('should not add WAF headers to local images if WAF is disabled', async (t) => {
+  const handler = createIPXHandler({
+    basePath: '/_ipx/',
+    cacheDir: '/tmp/ipx-cache',
+    bypassDomainCheck: true
+  }, (sourceImageOptions) => {
+    t.assert(sourceImageOptions.requestHeaders && sourceImageOptions.requestHeaders['X-Nf-Waf-Bypass-Token'] === undefined)
+
+    return Promise.resolve({ finalize: () => { } } as SourceImageResult)
+  })
+
+  await handler(
+    {
+      rawUrl: 'http://localhost:3000/some-path',
+      path: '/_ipx/w_500/https%3A%2F%2Fsome-site.com%2Fno-file.jpg',
+      headers: {},
+      rawQuery: '',
+      httpMethod: 'GET',
+      queryStringParameters: {},
+      multiValueQueryStringParameters: {},
+      multiValueHeaders: {},
+      isBase64Encoded: false,
+      body: null
+    },
+    {
+      functionName: 'ipx',
+      callbackWaitsForEmptyEventLoop: false,
+      functionVersion: '1',
+      invokedFunctionArn: '',
+      awsRequestId: '',
+      logGroupName: '',
+      logStreamName: '',
+      memoryLimitInMB: '',
+      getRemainingTimeInMillis: () => 1000,
+      done: () => { },
+      fail: () => { },
+      succeed: () => { }
+    }
+  )
+})