Merge pull request #190 from lifeomic/EOCI-745-prevent-recursion

[EOCI-745] prevent recursive range error on redirects

Author: evanheisler

src/alpha.ts

@@ -1,16 +1,16 @@
-import pick from 'lodash/pick';
 import merge from 'lodash/merge';
+import pick from 'lodash/pick';
 
+import { Handler } from 'aws-lambda';
 import axios, { Axios, AxiosAdapter, AxiosHeaders, AxiosResponse } from 'axios';
 import cloneDeep from 'lodash/cloneDeep';
 import { AlphaOptions, AlphaResponse, HandlerRequest } from './types';
-import { Handler } from 'aws-lambda';
 
+import { InvocationRequest } from '@aws-sdk/client-lambda';
 import { adapters } from './adapters';
-import { interceptors } from './interceptors';
 import { RequestError } from './adapters/helpers/requestError';
+import { interceptors } from './interceptors';
 import { resolve } from './resolve';
-import { InvocationRequest } from '@aws-sdk/client-lambda';
 
 const ALPHA_CONFIG = ['adapter', 'lambda', 'Lambda', 'retry', '__retryCount'];
 
@@ -72,15 +72,26 @@ export class Alpha extends Axios {
 
     const castResp = response as any as AxiosResponse;
 
-    if (castResp.status === 301 || castResp.status === 302) {
+    const redirectUrl = castResp.headers.location as string;
+
+    if (
+      (castResp.status === 301 || castResp.status === 302) &&
+      redirectUrl
+    ) {
       if (maxRedirects === 0) {
         const request = castResp.request as InvocationRequest | HandlerRequest;
         throw new RequestError('Exceeded maximum number of redirects.', castResp.config, request, castResp);
       }
 
       const redirect = cloneDeep(config);
       redirect.maxRedirects = maxRedirects - 1;
-      redirect.url = resolve(castResp.headers.location as string, castResp.config.url);
+      redirect.url = resolve(redirectUrl, castResp.config.url);
+
+      // Prevent recursive loops
+      if (redirect.url === config.url) {
+        return response as R;
+      }
+
       return this.request(redirect);
     }
 

test/redirects.test.ts

@@ -1,7 +1,7 @@
-import { Alpha } from '../src';
-import nock from 'nock';
 import { InvokeCommand, Lambda } from '@aws-sdk/client-lambda';
 import { mockClient } from 'aws-sdk-client-mock';
+import nock from 'nock';
+import { Alpha } from '../src';
 import { createResponse, prepResponse } from './utils';
 
 const mockLambda = mockClient(Lambda);
@@ -171,3 +171,20 @@ test('Redirects can be explicitly limited', async () => {
   expect(error.response.headers.location).toBe('/two');
   expect(server.isDone()).toBe(true);
 });
+
+test('A Lambda should not follow a redirect back to itself', async () => {
+  createResponse(mockLambda, {
+    StatusCode: 200,
+    Payload: {
+      headers: { location: 'lambda://test' },
+      statusCode: 302,
+    },
+  });
+
+  const response = await ctx.alpha.get('lambda://test');
+
+  // The desired behavior is to receive the 302 response, not to recursively follow it.
+  expect(response.status).toBe(302);
+  // Ensure we only called the lambda once.
+  expect(mockLambda.commandCalls(InvokeCommand)).toHaveLength(1);
+});