From 30f5931d84d887589589b09b2b50d98fadf5fc8e Mon Sep 17 00:00:00 2001
From: "Michael J. Roberts" <michael@observerly.com>
Date: Mon, 29 May 2023 14:45:40 +0100
Subject: [PATCH] feat: Added useOriginAgentCluster() composable.

feat: Added useOriginAgentCluster() composable.
---
 src/internals/useOriginAgentCluster.ts | 33 ++++++++++++++++
 tests/mocks/helmet.ts                  | 41 +++++++++++++++++++
 tests/useOriginAgentCluster.spec.ts    | 55 ++++++++++++++++++++++++++
 3 files changed, 129 insertions(+)
 create mode 100644 src/internals/useOriginAgentCluster.ts
 create mode 100644 tests/useOriginAgentCluster.spec.ts

diff --git a/src/internals/useOriginAgentCluster.ts b/src/internals/useOriginAgentCluster.ts
new file mode 100644
index 0000000..a0653ed
--- /dev/null
+++ b/src/internals/useOriginAgentCluster.ts
@@ -0,0 +1,33 @@
+/*****************************************************************************************************************/
+
+// @author         Michael Roberts <michael@observerly.com>
+// @package        @observerly/nitro-helmet
+// @license        Copyright © 2021-2023 observerly
+
+/*****************************************************************************************************************/
+
+import { type H3Event, setHeader } from 'h3'
+
+import { type OriginAgentCluster } from './types'
+
+/*****************************************************************************************************************/
+
+export const defaultOriginAgentCluster = '?1'
+
+/*****************************************************************************************************************/
+
+export const useOriginAgentCluster = (
+  event: H3Event,
+  policy: OriginAgentCluster | boolean = defaultOriginAgentCluster
+) => {
+  if (typeof policy === 'boolean' && !policy) {
+    return
+  }
+
+  // If the policy is a boolean, then use the default policy:
+  if (typeof policy === 'boolean') {
+    policy = defaultOriginAgentCluster
+  }
+
+  setHeader(event, 'Origin-Agent-Cluster', policy)
+}
diff --git a/tests/mocks/helmet.ts b/tests/mocks/helmet.ts
index 364dc96..1360061 100644
--- a/tests/mocks/helmet.ts
+++ b/tests/mocks/helmet.ts
@@ -30,6 +30,11 @@ import {
   defaultCrossOriginResourcePolicy
 } from '../../src/internals/useCrossOriginResourcePolicy'
 
+import {
+  useOriginAgentCluster,
+  defaultOriginAgentCluster
+} from '../../src/internals/useOriginAgentCluster'
+
 /*****************************************************************************************************************/
 
 export const helmetHandlers: Handler[] = [
@@ -160,6 +165,42 @@ export const helmetHandlers: Handler[] = [
     handler: eventHandler(async event => {
       useCrossOriginResourcePolicy(event, false)
 
+      return {
+        policy: false
+      }
+    })
+  },
+
+  //
+
+  {
+    method: 'GET',
+    url: '/helmet/origin-agent-cluster',
+    handler: eventHandler(async event => {
+      useOriginAgentCluster(event, '?1')
+
+      return {
+        policy: '?1'
+      }
+    })
+  },
+  {
+    method: 'GET',
+    url: '/helmet/default-origin-agent-cluster',
+    handler: eventHandler(async event => {
+      useOriginAgentCluster(event)
+
+      return {
+        policy: defaultOriginAgentCluster
+      }
+    })
+  },
+  {
+    method: 'GET',
+    url: '/helmet/no-origin-agent-cluster',
+    handler: eventHandler(async event => {
+      useOriginAgentCluster(event, false)
+
       return {
         policy: false
       }
diff --git a/tests/useOriginAgentCluster.spec.ts b/tests/useOriginAgentCluster.spec.ts
new file mode 100644
index 0000000..da12426
--- /dev/null
+++ b/tests/useOriginAgentCluster.spec.ts
@@ -0,0 +1,55 @@
+/*****************************************************************************************************************/
+
+// @author         Michael Roberts <michael@observerly.com>
+// @package        @observerly/nitro-helmet
+// @license        Copyright © 2021-2023 observerly
+
+/*****************************************************************************************************************/
+
+import { beforeEach, describe, expect, it } from 'vitest'
+
+import supertest, { type SuperTest, type Test } from 'supertest'
+
+import { toNodeListener } from 'h3'
+
+import { server } from './utilities/server'
+
+import {
+  defaultOriginAgentCluster,
+  useOriginAgentCluster
+} from '../src/internals/useOriginAgentCluster'
+
+/*****************************************************************************************************************/
+
+describe('useOriginAgentCluster', () => {
+  let request: SuperTest<Test>
+
+  beforeEach(() => {
+    request = supertest(toNodeListener(server))
+  })
+
+  it('should be defined', () => {
+    expect(useOriginAgentCluster).toBeDefined()
+  })
+
+  it('should set Origin-Agent-Cluster on the request', async () => {
+    const res = await request.get('/helmet/origin-agent-cluster', {
+      method: 'GET'
+    })
+    expect(res.header['origin-agent-cluster']).toEqual('?1')
+  })
+
+  it('should set Origin-Agent-Cluster on the request', async () => {
+    const res = await request.get('/helmet/default-origin-agent-cluster', {
+      method: 'GET'
+    })
+    expect(res.header['origin-agent-cluster']).toEqual(defaultOriginAgentCluster)
+  })
+
+  it('should not set Origin-Agent-Cluster on the request', async () => {
+    const res = await request.get('/helmet/no-origin-agent-cluster', {
+      method: 'GET'
+    })
+    expect(res.header['origin-agent-cluster']).toEqual(undefined)
+  })
+})