Add pre-response http interceptor (#52366) (#52404)

* add onPreResponse interceptor

* expose registerPreResponse to plugins

* address comments

* regen docs

docs/development/core/server/kibana-plugin-server.httpservicesetup.md

@@ -23,6 +23,7 @@ export interface HttpServiceSetup
 |  [registerAuth](./kibana-plugin-server.httpservicesetup.registerauth.md) | <code>(handler: AuthenticationHandler) =&gt; void</code> | To define custom authentication and/or authorization mechanism for incoming requests. |
 |  [registerOnPostAuth](./kibana-plugin-server.httpservicesetup.registeronpostauth.md) | <code>(handler: OnPostAuthHandler) =&gt; void</code> | To define custom logic to perform for incoming requests. |
 |  [registerOnPreAuth](./kibana-plugin-server.httpservicesetup.registeronpreauth.md) | <code>(handler: OnPreAuthHandler) =&gt; void</code> | To define custom logic to perform for incoming requests. |
+|  [registerOnPreResponse](./kibana-plugin-server.httpservicesetup.registeronpreresponse.md) | <code>(handler: OnPreResponseHandler) =&gt; void</code> | To define custom logic to perform for the server response. |
 |  [registerRouteHandlerContext](./kibana-plugin-server.httpservicesetup.registerroutehandlercontext.md) | <code>&lt;T extends keyof RequestHandlerContext&gt;(contextName: T, provider: RequestHandlerContextProvider&lt;T&gt;) =&gt; RequestHandlerContextContainer</code> | Register a context provider for a route handler. |
 
 ## Example

docs/development/core/server/kibana-plugin-server.httpservicesetup.registeronpreresponse.md

@@ -0,0 +1,18 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-server](./kibana-plugin-server.md) &gt; [HttpServiceSetup](./kibana-plugin-server.httpservicesetup.md) &gt; [registerOnPreResponse](./kibana-plugin-server.httpservicesetup.registeronpreresponse.md)
+
+## HttpServiceSetup.registerOnPreResponse property
+
+To define custom logic to perform for the server response.
+
+<b>Signature:</b>
+
+```typescript
+registerOnPreResponse: (handler: OnPreResponseHandler) => void;
+```
+
+## Remarks
+
+Doesn't provide the whole response object. Supports extending response with custom headers. See [OnPreResponseHandler](./kibana-plugin-server.onpreresponsehandler.md)<!-- -->.
+

docs/development/core/server/kibana-plugin-server.md

@@ -77,6 +77,9 @@ The plugin integrates with the core system via lifecycle events: `setup`<!-- -->
 |  [LogMeta](./kibana-plugin-server.logmeta.md) | Contextual metadata |
 |  [OnPostAuthToolkit](./kibana-plugin-server.onpostauthtoolkit.md) | A tool set defining an outcome of OnPostAuth interceptor for incoming request. |
 |  [OnPreAuthToolkit](./kibana-plugin-server.onpreauthtoolkit.md) | A tool set defining an outcome of OnPreAuth interceptor for incoming request. |
+|  [OnPreResponseExtensions](./kibana-plugin-server.onpreresponseextensions.md) | Additional data to extend a response. |
+|  [OnPreResponseInfo](./kibana-plugin-server.onpreresponseinfo.md) | Response status code. |
+|  [OnPreResponseToolkit](./kibana-plugin-server.onpreresponsetoolkit.md) | A tool set defining an outcome of OnPreAuth interceptor for incoming request. |
 |  [PackageInfo](./kibana-plugin-server.packageinfo.md) |  |
 |  [Plugin](./kibana-plugin-server.plugin.md) | The interface that should be returned by a <code>PluginInitializer</code>. |
 |  [PluginConfigDescriptor](./kibana-plugin-server.pluginconfigdescriptor.md) | Describes a plugin configuration schema and capabilities. |
@@ -173,6 +176,7 @@ The plugin integrates with the core system via lifecycle events: `setup`<!-- -->
 |  [MutatingOperationRefreshSetting](./kibana-plugin-server.mutatingoperationrefreshsetting.md) | Elasticsearch Refresh setting for mutating operation |
 |  [OnPostAuthHandler](./kibana-plugin-server.onpostauthhandler.md) | See [OnPostAuthToolkit](./kibana-plugin-server.onpostauthtoolkit.md)<!-- -->. |
 |  [OnPreAuthHandler](./kibana-plugin-server.onpreauthhandler.md) | See [OnPreAuthToolkit](./kibana-plugin-server.onpreauthtoolkit.md)<!-- -->. |
+|  [OnPreResponseHandler](./kibana-plugin-server.onpreresponsehandler.md) | See [OnPreAuthToolkit](./kibana-plugin-server.onpreauthtoolkit.md)<!-- -->. |
 |  [PluginConfigSchema](./kibana-plugin-server.pluginconfigschema.md) | Dedicated type for plugin configuration schema. |
 |  [PluginInitializer](./kibana-plugin-server.plugininitializer.md) | The <code>plugin</code> export at the root of a plugin's <code>server</code> directory should conform to this interface. |
 |  [PluginName](./kibana-plugin-server.pluginname.md) | Dedicated type for plugin name/id that is supposed to make Map/Set/Arrays that use it as a key or value more obvious. |

docs/development/core/server/kibana-plugin-server.onpreresponseextensions.headers.md

@@ -0,0 +1,13 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-server](./kibana-plugin-server.md) &gt; [OnPreResponseExtensions](./kibana-plugin-server.onpreresponseextensions.md) &gt; [headers](./kibana-plugin-server.onpreresponseextensions.headers.md)
+
+## OnPreResponseExtensions.headers property
+
+additional headers to attach to the response
+
+<b>Signature:</b>
+
+```typescript
+headers?: ResponseHeaders;
+```

docs/development/core/server/kibana-plugin-server.onpreresponseextensions.md

@@ -0,0 +1,20 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-server](./kibana-plugin-server.md) &gt; [OnPreResponseExtensions](./kibana-plugin-server.onpreresponseextensions.md)
+
+## OnPreResponseExtensions interface
+
+Additional data to extend a response.
+
+<b>Signature:</b>
+
+```typescript
+export interface OnPreResponseExtensions 
+```
+
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [headers](./kibana-plugin-server.onpreresponseextensions.headers.md) | <code>ResponseHeaders</code> | additional headers to attach to the response |
+

docs/development/core/server/kibana-plugin-server.onpreresponsehandler.md

@@ -0,0 +1,13 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-server](./kibana-plugin-server.md) &gt; [OnPreResponseHandler](./kibana-plugin-server.onpreresponsehandler.md)
+
+## OnPreResponseHandler type
+
+See [OnPreAuthToolkit](./kibana-plugin-server.onpreauthtoolkit.md)<!-- -->.
+
+<b>Signature:</b>
+
+```typescript
+export declare type OnPreResponseHandler = (request: KibanaRequest, preResponse: OnPreResponseInfo, toolkit: OnPreResponseToolkit) => OnPreResponseResult | Promise<OnPreResponseResult>;
+```

docs/development/core/server/kibana-plugin-server.onpreresponseinfo.md

@@ -0,0 +1,20 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-server](./kibana-plugin-server.md) &gt; [OnPreResponseInfo](./kibana-plugin-server.onpreresponseinfo.md)
+
+## OnPreResponseInfo interface
+
+Response status code.
+
+<b>Signature:</b>
+
+```typescript
+export interface OnPreResponseInfo 
+```
+
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [statusCode](./kibana-plugin-server.onpreresponseinfo.statuscode.md) | <code>number</code> |  |
+

docs/development/core/server/kibana-plugin-server.onpreresponseinfo.statuscode.md

@@ -0,0 +1,11 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-server](./kibana-plugin-server.md) &gt; [OnPreResponseInfo](./kibana-plugin-server.onpreresponseinfo.md) &gt; [statusCode](./kibana-plugin-server.onpreresponseinfo.statuscode.md)
+
+## OnPreResponseInfo.statusCode property
+
+<b>Signature:</b>
+
+```typescript
+statusCode: number;
+```

docs/development/core/server/kibana-plugin-server.onpreresponsetoolkit.md

@@ -0,0 +1,20 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-server](./kibana-plugin-server.md) &gt; [OnPreResponseToolkit](./kibana-plugin-server.onpreresponsetoolkit.md)
+
+## OnPreResponseToolkit interface
+
+A tool set defining an outcome of OnPreAuth interceptor for incoming request.
+
+<b>Signature:</b>
+
+```typescript
+export interface OnPreResponseToolkit 
+```
+
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [next](./kibana-plugin-server.onpreresponsetoolkit.next.md) | <code>(responseExtensions?: OnPreResponseExtensions) =&gt; OnPreResponseResult</code> | To pass request to the next handler |
+

docs/development/core/server/kibana-plugin-server.onpreresponsetoolkit.next.md

@@ -0,0 +1,13 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-server](./kibana-plugin-server.md) &gt; [OnPreResponseToolkit](./kibana-plugin-server.onpreresponsetoolkit.md) &gt; [next](./kibana-plugin-server.onpreresponsetoolkit.next.md)
+
+## OnPreResponseToolkit.next property
+
+To pass request to the next handler
+
+<b>Signature:</b>
+
+```typescript
+next: (responseExtensions?: OnPreResponseExtensions) => OnPreResponseResult;
+```

src/core/server/http/http_server.ts

@@ -16,8 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-
-import { Request, Server } from 'hapi';
+import { Server } from 'hapi';
 import url from 'url';
 
 import { Logger, LoggerFactory } from '../logging';
@@ -26,8 +25,9 @@ import { createServer, getListenerOptions, getServerOptions } from './http_tools
 import { adoptToHapiAuthFormat, AuthenticationHandler } from './lifecycle/auth';
 import { adoptToHapiOnPostAuthFormat, OnPostAuthHandler } from './lifecycle/on_post_auth';
 import { adoptToHapiOnPreAuthFormat, OnPreAuthHandler } from './lifecycle/on_pre_auth';
+import { adoptToHapiOnPreResponseFormat, OnPreResponseHandler } from './lifecycle/on_pre_response';
 
-import { ResponseHeaders, IRouter } from './router';
+import { IRouter } from './router';
 import {
   SessionStorageCookieOptions,
   createCookieSessionStorageFactory,
@@ -50,6 +50,7 @@ export interface HttpServerSetup {
   registerAuth: HttpServiceSetup['registerAuth'];
   registerOnPreAuth: HttpServiceSetup['registerOnPreAuth'];
   registerOnPostAuth: HttpServiceSetup['registerOnPostAuth'];
+  registerOnPreResponse: HttpServiceSetup['registerOnPreResponse'];
   isTlsEnabled: HttpServiceSetup['isTlsEnabled'];
   auth: {
     get: GetAuthState;
@@ -103,6 +104,7 @@ export class HttpServer {
       registerRouter: this.registerRouter.bind(this),
       registerOnPreAuth: this.registerOnPreAuth.bind(this),
       registerOnPostAuth: this.registerOnPostAuth.bind(this),
+      registerOnPreResponse: this.registerOnPreResponse.bind(this),
       createCookieSessionStorageFactory: <T>(cookieOptions: SessionStorageCookieOptions<T>) =>
         this.createCookieSessionStorageFactory(cookieOptions, config.basePath),
       registerAuth: this.registerAuth.bind(this),
@@ -232,6 +234,14 @@ export class HttpServer {
     this.server.ext('onRequest', adoptToHapiOnPreAuthFormat(fn, this.log));
   }
 
+  private registerOnPreResponse(fn: OnPreResponseHandler) {
+    if (this.server === undefined) {
+      throw new Error('Server is not created yet');
+    }
+
+    this.server.ext('onPreResponse', adoptToHapiOnPreResponseFormat(fn, this.log));
+  }
+
   private async createCookieSessionStorageFactory<T>(
     cookieOptions: SessionStorageCookieOptions<T>,
     basePath?: string
@@ -289,39 +299,9 @@ export class HttpServer {
     // https://github.com/hapijs/hapi/blob/master/API.md#-serverauthdefaultoptions
     this.server.auth.default('session');
 
-    this.server.ext('onPreResponse', (request, t) => {
+    this.registerOnPreResponse((request, preResponseInfo, t) => {
       const authResponseHeaders = this.authResponseHeaders.get(request);
-      this.extendResponseWithHeaders(request, authResponseHeaders);
-      return t.continue;
-    });
-  }
-
-  private extendResponseWithHeaders(request: Request, headers?: ResponseHeaders) {
-    const response = request.response;
-    if (!headers || !response) return;
-
-    if (response instanceof Error) {
-      this.findHeadersIntersection(response.output.headers, headers);
-      // hapi wraps all error response in Boom object internally
-      response.output.headers = {
-        ...response.output.headers,
-        ...(headers as any), // hapi types don't specify string[] as valid value
-      };
-    } else {
-      for (const [headerName, headerValue] of Object.entries(headers)) {
-        this.findHeadersIntersection(response.headers, headers);
-        response.header(headerName, headerValue as any); // hapi types don't specify string[] as valid value
-      }
-    }
-  }
-
-  // NOTE: responseHeaders contains not a full list of response headers, but only explicitly set on a response object.
-  // any headers added by hapi internally, like `content-type`, `content-length`, etc. do not present here.
-  private findHeadersIntersection(responseHeaders: ResponseHeaders, headers: ResponseHeaders) {
-    Object.keys(headers).forEach(headerName => {
-      if (responseHeaders[headerName] !== undefined) {
-        this.log.warn(`Server rewrites a response header [${headerName}].`);
-      }
+      return t.next({ headers: authResponseHeaders });
     });
   }
 }

src/core/server/http/http_service.mock.ts

@@ -51,6 +51,7 @@ const createSetupContractMock = () => {
     registerAuth: jest.fn(),
     registerOnPostAuth: jest.fn(),
     registerRouteHandlerContext: jest.fn(),
+    registerOnPreResponse: jest.fn(),
     createRouter: jest.fn().mockImplementation(() => mockRouter.create({})),
     basePath: createBasePathMock(),
     auth: {

src/core/server/http/index.ts

@@ -64,6 +64,12 @@ export {
   AuthResultType,
 } from './lifecycle/auth';
 export { OnPostAuthHandler, OnPostAuthToolkit } from './lifecycle/on_post_auth';
+export {
+  OnPreResponseHandler,
+  OnPreResponseToolkit,
+  OnPreResponseExtensions,
+  OnPreResponseInfo,
+} from './lifecycle/on_pre_response';
 export { SessionStorageFactory, SessionStorage } from './session_storage';
 export {
   SessionStorageCookieOptions,

src/core/server/http/integration_tests/lifecycle.test.ts

@@ -161,7 +161,7 @@ describe('OnPreAuth', () => {
     expect(result.header['www-authenticate']).toBe('challenge');
   });
 
-  it("doesn't expose error details if interceptor throws", async () => {
+  it('does not expose error details if interceptor throws', async () => {
     const { registerOnPreAuth, server: innerServer, createRouter } = await server.setup(setupDeps);
     const router = createRouter('/');
 
@@ -734,7 +734,7 @@ describe('Auth', () => {
     expect(loggingServiceMock.collect(logger).warn).toMatchInlineSnapshot(`
       Array [
         Array [
-          "Server rewrites a response header [www-authenticate].",
+          "onPreResponseHandler rewrote a response header [www-authenticate].",
         ],
       ]
     `);
@@ -769,7 +769,7 @@ describe('Auth', () => {
     expect(loggingServiceMock.collect(logger).warn).toMatchInlineSnapshot(`
       Array [
         Array [
-          "Server rewrites a response header [www-authenticate].",
+          "onPreResponseHandler rewrote a response header [www-authenticate].",
         ],
       ]
     `);
@@ -893,3 +893,165 @@ describe('Auth', () => {
       .expect(200, { customField: 'undefined' });
   });
 });
+
+describe('OnPreResponse', () => {
+  it('supports registering response inceptors', async () => {
+    const { registerOnPreResponse, server: innerServer, createRouter } = await server.setup(
+      setupDeps
+    );
+    const router = createRouter('/');
+
+    router.get({ path: '/', validate: false }, (context, req, res) => res.ok({ body: 'ok' }));
+
+    const callingOrder: string[] = [];
+    registerOnPreResponse((req, res, t) => {
+      callingOrder.push('first');
+      return t.next();
+    });
+
+    registerOnPreResponse((req, res, t) => {
+      callingOrder.push('second');
+      return t.next();
+    });
+    await server.start();
+
+    await supertest(innerServer.listener)
+      .get('/')
+      .expect(200, 'ok');
+
+    expect(callingOrder).toEqual(['first', 'second']);
+  });
+
+  it('supports additional headers attachments', async () => {
+    const { registerOnPreResponse, server: innerServer, createRouter } = await server.setup(
+      setupDeps
+    );
+    const router = createRouter('/');
+
+    router.get({ path: '/', validate: false }, (context, req, res) =>
+      res.ok({
+        headers: {
+          'x-my-header': 'foo',
+        },
+      })
+    );
+
+    registerOnPreResponse((req, res, t) =>
+      t.next({
+        headers: {
+          'x-kibana-header': 'value',
+        },
+      })
+    );
+    await server.start();
+
+    const result = await supertest(innerServer.listener)
+      .get('/')
+      .expect(200);
+
+    expect(result.header['x-kibana-header']).toBe('value');
+    expect(result.header['x-my-header']).toBe('foo');
+  });
+
+  it('logs a warning if interceptor rewrites response header', async () => {
+    const { registerOnPreResponse, server: innerServer, createRouter } = await server.setup(
+      setupDeps
+    );
+    const router = createRouter('/');
+
+    router.get({ path: '/', validate: false }, (context, req, res) =>
+      res.ok({
+        headers: { 'x-kibana-header': 'value' },
+      })
+    );
+
+    registerOnPreResponse((req, res, t) =>
+      t.next({
+        headers: { 'x-kibana-header': 'value' },
+      })
+    );
+    await server.start();
+
+    await supertest(innerServer.listener)
+      .get('/')
+      .expect(200);
+
+    expect(loggingServiceMock.collect(logger).warn).toMatchInlineSnapshot(`
+      Array [
+        Array [
+          "onPreResponseHandler rewrote a response header [x-kibana-header].",
+        ],
+      ]
+    `);
+  });
+
+  it("doesn't expose error details if interceptor throws", async () => {
+    const { registerOnPreResponse, server: innerServer, createRouter } = await server.setup(
+      setupDeps
+    );
+    const router = createRouter('/');
+
+    router.get({ path: '/', validate: false }, (context, req, res) => res.ok(undefined));
+    registerOnPreResponse((req, res, t) => {
+      throw new Error('reason');
+    });
+    await server.start();
+
+    const result = await supertest(innerServer.listener)
+      .get('/')
+      .expect(500);
+
+    expect(result.body.message).toBe('An internal server error occurred.');
+    expect(loggingServiceMock.collect(logger).error).toMatchInlineSnapshot(`
+      Array [
+        Array [
+          [Error: reason],
+        ],
+      ]
+    `);
+  });
+
+  it('returns internal error if interceptor returns unexpected result', async () => {
+    const { registerOnPreResponse, server: innerServer, createRouter } = await server.setup(
+      setupDeps
+    );
+    const router = createRouter('/');
+
+    router.get({ path: '/', validate: false }, (context, req, res) => res.ok());
+    registerOnPreResponse((req, res, t) => ({} as any));
+    await server.start();
+
+    const result = await supertest(innerServer.listener)
+      .get('/')
+      .expect(500);
+
+    expect(result.body.message).toBe('An internal server error occurred.');
+    expect(loggingServiceMock.collect(logger).error).toMatchInlineSnapshot(`
+      Array [
+        Array [
+          [Error: Unexpected result from OnPreResponse. Expected OnPreResponseResult, but given: [object Object].],
+        ],
+      ]
+    `);
+  });
+
+  it('cannot change response statusCode', async () => {
+    const { registerOnPreResponse, server: innerServer, createRouter } = await server.setup(
+      setupDeps
+    );
+    const router = createRouter('/');
+
+    registerOnPreResponse((req, res, t) => {
+      res.statusCode = 500;
+      return t.next();
+    });
+
+    router.get({ path: '/', validate: false }, (context, req, res) => res.ok({ body: 'ok' }));
+
+    await server.start();
+
+    await supertest(innerServer.listener)
+      .get('/')
+      .expect(200);
+  });
+});

src/core/server/http/lifecycle/on_pre_response.ts

@@ -0,0 +1,155 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { Lifecycle, Request, ResponseToolkit as HapiResponseToolkit } from 'hapi';
+import Boom from 'boom';
+import { Logger } from '../../logging';
+
+import { HapiResponseAdapter, KibanaRequest, ResponseHeaders } from '../router';
+
+enum ResultType {
+  next = 'next',
+}
+
+interface Next {
+  type: ResultType.next;
+  headers?: ResponseHeaders;
+}
+
+/**
+ * @internal
+ */
+type OnPreResponseResult = Next;
+
+/**
+ * Additional data to extend a response.
+ * @public
+ */
+export interface OnPreResponseExtensions {
+  /** additional headers to attach to the response */
+  headers?: ResponseHeaders;
+}
+
+/**
+ * Response status code.
+ * @public
+ */
+export interface OnPreResponseInfo {
+  statusCode: number;
+}
+
+const preResponseResult = {
+  next(responseExtensions?: OnPreResponseExtensions): OnPreResponseResult {
+    return { type: ResultType.next, headers: responseExtensions?.headers };
+  },
+  isNext(result: OnPreResponseResult): result is Next {
+    return result && result.type === ResultType.next;
+  },
+};
+
+/**
+ * A tool set defining an outcome of OnPreAuth interceptor for incoming request.
+ * @public
+ */
+export interface OnPreResponseToolkit {
+  /** To pass request to the next handler */
+  next: (responseExtensions?: OnPreResponseExtensions) => OnPreResponseResult;
+}
+
+const toolkit: OnPreResponseToolkit = {
+  next: preResponseResult.next,
+};
+
+/**
+ * See {@link OnPreAuthToolkit}.
+ * @public
+ */
+export type OnPreResponseHandler = (
+  request: KibanaRequest,
+  preResponse: OnPreResponseInfo,
+  toolkit: OnPreResponseToolkit
+) => OnPreResponseResult | Promise<OnPreResponseResult>;
+
+/**
+ * @public
+ * Adopt custom request interceptor to Hapi lifecycle system.
+ * @param fn - an extension point allowing to perform custom logic for
+ * incoming HTTP requests.
+ */
+export function adoptToHapiOnPreResponseFormat(fn: OnPreResponseHandler, log: Logger) {
+  return async function interceptPreResponse(
+    request: Request,
+    responseToolkit: HapiResponseToolkit
+  ): Promise<Lifecycle.ReturnValue> {
+    const response = request.response;
+
+    try {
+      if (response) {
+        const statusCode: number = isBoom(response)
+          ? response.output.statusCode
+          : response.statusCode;
+
+        const result = await fn(KibanaRequest.from(request), { statusCode }, toolkit);
+        if (!preResponseResult.isNext(result)) {
+          throw new Error(
+            `Unexpected result from OnPreResponse. Expected OnPreResponseResult, but given: ${result}.`
+          );
+        }
+        if (result.headers) {
+          if (isBoom(response)) {
+            findHeadersIntersection(response.output.headers, result.headers, log);
+            // hapi wraps all error response in Boom object internally
+            response.output.headers = {
+              ...response.output.headers,
+              ...(result.headers as any), // hapi types don't specify string[] as valid value
+            };
+          } else {
+            for (const [headerName, headerValue] of Object.entries(result.headers)) {
+              findHeadersIntersection(response.headers, result.headers, log);
+              response.header(headerName, headerValue as any); // hapi types don't specify string[] as valid value
+            }
+          }
+        }
+      }
+    } catch (error) {
+      log.error(error);
+      const hapiResponseAdapter = new HapiResponseAdapter(responseToolkit);
+      return hapiResponseAdapter.toInternalError();
+    }
+    return responseToolkit.continue;
+  };
+}
+
+function isBoom(response: any): response is Boom {
+  return response instanceof Boom;
+}
+
+// NOTE: responseHeaders contains not a full list of response headers, but only explicitly set on a response object.
+// any headers added by hapi internally, like `content-type`, `content-length`, etc. are not present here.
+function findHeadersIntersection(
+  responseHeaders: ResponseHeaders,
+  headers: ResponseHeaders,
+  log: Logger
+) {
+  Object.keys(headers).forEach(headerName => {
+    if (responseHeaders[headerName] !== undefined) {
+      log.warn(`onPreResponseHandler rewrote a response header [${headerName}].`);
+    }
+  });
+}

src/core/server/http/types.ts

@@ -24,6 +24,7 @@ import { SessionStorageFactory } from './session_storage';
 import { AuthenticationHandler } from './lifecycle/auth';
 import { OnPreAuthHandler } from './lifecycle/on_pre_auth';
 import { OnPostAuthHandler } from './lifecycle/on_post_auth';
+import { OnPreResponseHandler } from './lifecycle/on_pre_response';
 import { IBasePath } from './base_path_service';
 import { PluginOpaqueId, RequestHandlerContext } from '..';
 
@@ -163,6 +164,18 @@ export interface HttpServiceSetup {
    */
   registerOnPostAuth: (handler: OnPostAuthHandler) => void;
 
+  /**
+   * To define custom logic to perform for the server response.
+   *
+   * @remarks
+   * Doesn't provide the whole response object.
+   * Supports extending response with custom headers.
+   * See {@link OnPreResponseHandler}.
+   *
+   * @param handler {@link OnPreResponseHandler} - function to call.
+   */
+  registerOnPreResponse: (handler: OnPreResponseHandler) => void;
+
   /**
    * Access or manipulate the Kibana base path
    * See {@link IBasePath}.

src/core/server/index.ts

@@ -105,6 +105,10 @@ export {
   OnPreAuthToolkit,
   OnPostAuthHandler,
   OnPostAuthToolkit,
+  OnPreResponseHandler,
+  OnPreResponseToolkit,
+  OnPreResponseExtensions,
+  OnPreResponseInfo,
   RedirectResponseOptions,
   RequestHandler,
   RequestHandlerContextContainer,

src/core/server/legacy/legacy_service.ts

@@ -270,6 +270,7 @@ export class LegacyService implements CoreService {
         registerOnPreAuth: setupDeps.core.http.registerOnPreAuth,
         registerAuth: setupDeps.core.http.registerAuth,
         registerOnPostAuth: setupDeps.core.http.registerOnPostAuth,
+        registerOnPreResponse: setupDeps.core.http.registerOnPreResponse,
         basePath: setupDeps.core.http.basePath,
         isTlsEnabled: setupDeps.core.http.isTlsEnabled,
       },

src/core/server/mocks.ts

@@ -90,6 +90,7 @@ function createCoreSetupMock() {
     registerOnPreAuth: httpService.registerOnPreAuth,
     registerAuth: httpService.registerAuth,
     registerOnPostAuth: httpService.registerOnPostAuth,
+    registerOnPreResponse: httpService.registerOnPreResponse,
     basePath: httpService.basePath,
     isTlsEnabled: httpService.isTlsEnabled,
     createRouter: jest.fn(),

src/core/server/plugins/plugin_context.ts

@@ -159,6 +159,7 @@ export function createPluginSetupContext<TPlugin, TPluginDependencies>(
       registerOnPreAuth: deps.http.registerOnPreAuth,
       registerAuth: deps.http.registerAuth,
       registerOnPostAuth: deps.http.registerOnPostAuth,
+      registerOnPreResponse: deps.http.registerOnPreResponse,
       basePath: deps.http.basePath,
       isTlsEnabled: deps.http.isTlsEnabled,
     },

src/core/server/server.api.md

@@ -697,6 +697,7 @@ export interface HttpServiceSetup {
     registerAuth: (handler: AuthenticationHandler) => void;
     registerOnPostAuth: (handler: OnPostAuthHandler) => void;
     registerOnPreAuth: (handler: OnPreAuthHandler) => void;
+    registerOnPreResponse: (handler: OnPreResponseHandler) => void;
     registerRouteHandlerContext: <T extends keyof RequestHandlerContext>(contextName: T, provider: RequestHandlerContextProvider<T>) => RequestHandlerContextContainer;
 }
 
@@ -976,6 +977,27 @@ export interface OnPreAuthToolkit {
     rewriteUrl: (url: string) => OnPreAuthResult;
 }
 
+// @public
+export interface OnPreResponseExtensions {
+    headers?: ResponseHeaders;
+}
+
+// Warning: (ae-forgotten-export) The symbol "OnPreResponseResult" needs to be exported by the entry point index.d.ts
+// 
+// @public
+export type OnPreResponseHandler = (request: KibanaRequest, preResponse: OnPreResponseInfo, toolkit: OnPreResponseToolkit) => OnPreResponseResult | Promise<OnPreResponseResult>;
+
+// @public
+export interface OnPreResponseInfo {
+    // (undocumented)
+    statusCode: number;
+}
+
+// @public
+export interface OnPreResponseToolkit {
+    next: (responseExtensions?: OnPreResponseExtensions) => OnPreResponseResult;
+}
+
 // @public (undocumented)
 export interface PackageInfo {
     // (undocumented)