feat: add support for custom headers from origins (#261)

* Add support for custom headers from origins

* Chore: Update documented node version

* Fix oopsie with ts version

* Also include minor fix for uri

* feat: Overhaul how origins are injected in the request lifecycle

* Chore: Typo

Co-authored-by: ryan lee <rol@bluetel.co.uk>

Author: ryanolee

README.md

@@ -20,7 +20,7 @@ plugins:
 
 provider:
   name: aws
-  runtime: nodejs8.10
+  runtime: nodejs12.x
 
 functions:
   lambda:
@@ -65,3 +65,13 @@ custom:
   offlineEdgeLambda:
     path: '.build'
 ```
+
+For usage with `serverless-webpack` and `serverless-bundle` the config is similar but the build path changes.
+```yaml
+plugins:
+  - serverless-webpack # or serverless-bundle
+  
+custom:
+  offlineEdgeLambda:
+    path: './.webpack/service/'
+```

examples/serverless.yml

@@ -6,7 +6,7 @@ plugins:
 
 provider:
   name: aws
-  runtime: nodejs8.10
+  runtime: nodejs12.x
 
 functions:
   onViewerRequest:
@@ -16,8 +16,7 @@ functions:
       eventType: 'viewer-request'
       pathPattern: '/lambda'
   onOriginRequest:
-    handler: src/handlers.onViewerRequest
-    lambdaAtEdge:
+    handler: src/handlers.onOriginRequest
       distribution: 'WebsiteDistribution'
       eventType: 'origin-request'
       pathPattern: '/lambda'

package-lock.json

@@ -13,10 +13,10 @@ import { HttpError, InternalServerError } from './errors/http';
 import { FunctionSet } from './function-set';
 import { asyncMiddleware, cloudfrontPost } from './middlewares';
 import { CloudFrontLifecycle, Origin, CacheService } from './services';
-import { ServerlessInstance, ServerlessOptions } from './types';
+import { CFDistribution, ServerlessInstance, ServerlessOptions } from './types';
 import {
 	buildConfig, buildContext, CloudFrontHeadersHelper, ConfigBuilder,
-	convertToCloudFrontEvent, IncomingMessageWithBodyAndCookies
+	convertToCloudFrontEvent, getOriginFromCfDistribution, IncomingMessageWithBodyAndCookies
 } from './utils';
 
 
@@ -30,6 +30,7 @@ export class BehaviorRouter {
 	private builder: ConfigBuilder;
 	private context: Context;
 	private behaviors = new Map<string, FunctionSet>();
+	private cfResources:  Record<string, CFDistribution>;
 
 	private cacheDir: string;
 	private fileDir: string;
@@ -51,6 +52,7 @@ export class BehaviorRouter {
 		this.builder = buildConfig(serverless);
 		this.context = buildContext();
 
+		this.cfResources = serverless.service?.resources?.Resources || {};
 		this.cacheDir = path.resolve(options.cacheDir || path.join(os.tmpdir(), 'edge-lambda'));
 		this.fileDir = path.resolve(options.fileDir || path.join(os.tmpdir(), 'edge-lambda'));
 		this.path = this.serverless.service.custom.offlineEdgeLambda.path || '';
@@ -143,16 +145,22 @@ export class BehaviorRouter {
 				}
 
 				const handler = this.match(req);
-				const cfEvent = convertToCloudFrontEvent(req, this.builder('viewer-request'));
 
 				if (!handler) {
 					res.statusCode = StatusCodes.NOT_FOUND;
 					res.end();
 					return;
 				}
 
+				const customOrigin = handler.distribution in this.cfResources ?
+					getOriginFromCfDistribution(handler.pattern, this.cfResources[handler.distribution]) :
+					null;
+
+				const cfEvent = convertToCloudFrontEvent(req, this.builder('viewer-request'));
+
 				try {
-					const lifecycle = new CloudFrontLifecycle(this.serverless, this.options, cfEvent, this.context, this.cacheService, handler);
+					const lifecycle = new CloudFrontLifecycle(this.serverless, this.options, cfEvent,
+																this.context, this.cacheService, handler, customOrigin);
 					const response = await lifecycle.run(req.url as string);
 
 					if (!response) {
@@ -230,20 +238,30 @@ export class BehaviorRouter {
 		behaviors.clear();
 
 		for await (const [, def] of lambdaDefs) {
+
 			const pattern = def.lambdaAtEdge.pathPattern || '*';
+			const distribution = def.lambdaAtEdge.distribution || '';
 
 			if (!behaviors.has(pattern)) {
 				const origin = this.origins.get(pattern);
-				behaviors.set(pattern, new FunctionSet(pattern, this.log, origin));
+				behaviors.set(pattern, new FunctionSet(pattern, distribution, this.log, origin));
 			}
 
 			const fnSet = behaviors.get(pattern) as FunctionSet;
 
+			// Don't try to register distributions that come from other sources
+			if (fnSet.distribution !== distribution) {
+				this.log(`Warning: pattern ${pattern} has registered handlers for cf distributions ${fnSet.distribution}` +
+						` and ${distribution}. There is no way to tell which distribution should be used so only ${fnSet.distribution}` +
+						` has been registered.` );
+				continue;
+			}
+
 			await fnSet.setHandler(def.lambdaAtEdge.eventType, path.join(this.path, def.handler));
 		}
 
 		if (!behaviors.has('*')) {
-			behaviors.set('*', new FunctionSet('*', this.log, this.origins.get('*')));
+			behaviors.set('*', new FunctionSet('*', '', this.log, this.origins.get('*')));
 		}
 	}
 
@@ -261,7 +279,10 @@ export class BehaviorRouter {
 
 	private logBehaviors() {
 		this.behaviors.forEach((behavior, key) => {
-			this.log(`Lambdas for path pattern ${key}: `);
+
+			this.log(`Lambdas for path pattern ${key}` +
+				(behavior.distribution === '' ? ':' : ` on ${behavior.distribution}:`)
+			);
 
 			behavior.viewerRequest && this.log(`viewer-request => ${behavior.viewerRequest.path || ''}`);
 			behavior.originRequest && this.log(`origin-request => ${behavior.originRequest.path || ''}`);

src/behavior-router.ts

@@ -26,6 +26,7 @@ export class FunctionSet {
 
 	constructor(
 		public readonly pattern: string,
+		public readonly distribution: string,
 		private readonly log: (message: string) => void,
 		public readonly origin: Origin = new Origin()
 	) {

src/function-set.ts

@@ -1,5 +1,5 @@
 import {
-	CloudFrontRequestEvent, CloudFrontResponseResult, CloudFrontResultResponse, Context
+	CloudFrontRequestEvent, CloudFrontResponseResult, CloudFrontOrigin, Context
 } from 'aws-lambda';
 
 import { NoResult } from '../errors';
@@ -19,7 +19,8 @@ export class CloudFrontLifecycle {
 		private event: CloudFrontRequestEvent,
 		private context: Context,
 		private fileService: CacheService,
-		private fnSet: FunctionSet
+		private fnSet: FunctionSet,
+		private origin: CloudFrontOrigin | null,
 	) {
 		this.log = serverless.cli.log.bind(serverless.cli);
 	}
@@ -85,6 +86,7 @@ export class CloudFrontLifecycle {
 			throw new NoResult();
 		} else {
 			this.log('✓ Cache hit');
+			throw new NoResult();
 		}
 
 		const result = toResultResponse(cached);
@@ -99,6 +101,10 @@ export class CloudFrontLifecycle {
 	async onOriginRequest() {
 		this.log('→ origin-request');
 
+		// Inject origin into request once we reach the origin request step
+		// as it is not available in viewer requests
+		this.injectOriginIntoRequest();
+
 		const result = await this.fnSet.originRequest(this.event, this.context);
 
 		if (isResponseResult(result)) {
@@ -116,4 +122,10 @@ export class CloudFrontLifecycle {
 		const event = combineResult(this.event, result);
 		return this.fnSet.originResponse(event, this.context);
 	}
+
+	protected injectOriginIntoRequest() {
+		if (this?.event?.Records[0]?.cf?.request && this.origin !== null) {
+			this.event.Records[0].cf.request.origin = this.origin;
+		}
+	}
 }

src/services/cloudfront.service.ts

@@ -13,11 +13,63 @@ export interface ServerlessInstance {
 		}
 		functions: { [key: string]: ServerlessFunction }
 		package: ServerlessPackage
+		resources?: {
+			Resources?: Record<string, CFDistribution>
+		}
 		getAllFunctions: () => string[]
 	};
 	pluginManager: PluginManager;
 }
 
+/**
+ * A stub for the CF distributions we want details for in the context of this app
+ */
+export interface CFDistribution {
+	Type: string;
+	Properties: {
+		DistributionConfig: {
+			Origins: CFOrigin[]
+			CacheBehaviors: CacheBehaviors[]
+		}
+	};
+}
+
+export interface CacheBehaviors {
+	PathPattern: string;
+	TargetOriginId: string;
+}
+
+export interface CFOrigin {
+	DomainName: string;
+	Id: string;
+	ConnectionTimeout: number;
+	OriginCustomHeaders: CFCustomHeaders[];
+	S3OriginConfig?: {
+		OriginAccessIdentity: string
+	};
+	CustomOriginConfig?: CFCustomOriginConfig;
+}
+
+
+export enum CFOriginProtocolPolicy {
+	HTTP_ONLY = 'http-only',
+	MATCH_VIEWER = 'match-viewer',
+	HTTPS_ONLY = 'https-only'
+}
+
+export interface CFCustomOriginConfig {
+	HTTPPort?: number;
+	HTTPSPort: number;
+	OriginKeepaliveTimeout: string;
+	OriginProtocolPolicy: CFOriginProtocolPolicy;
+	OriginReadTimeout: number;
+	OriginSSLProtocols: ('SSLv3' | 'TLSv1' | 'TLSv1.1' | 'TLSv1.2')[];
+}
+export interface CFCustomHeaders {
+	HeaderName: string;
+	HeaderValue: string;
+}
+
 export interface ServerlessOptions {
 	cacheDir: string;
 	cloudfrontPort: number;

src/types/serverless.types.ts

@@ -17,7 +17,7 @@ export function convertToCloudFrontEvent(req: IncomingMessageWithBodyAndCookies,
 		clientIp: req.socket.remoteAddress as string,
 		method: req.method as string,
 		headers: toCloudFrontHeaders(req.headers),
-		uri: url.href as string,
+		uri: url.pathname as string,
 		querystring:  url.query || '',
 		body: req.body,
 		cookies: req.cookies