From 04fa824bcf6cf817b45b89193fe6bf261b0b7c1e Mon Sep 17 00:00:00 2001
From: Hebilicious <xsh4k3@gmail.com>
Date: Wed, 26 Jul 2023 13:27:25 +0700
Subject: [PATCH] feat: add support for multiple headers

---
 src/utils/response.ts | 20 +++++++++++++++++---
 test/header.test.ts   | 24 +++++++++++++++++++++---
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/src/utils/response.ts b/src/utils/response.ts
index f918823a..c55dc51c 100644
--- a/src/utils/response.ts
+++ b/src/utils/response.ts
@@ -148,9 +148,23 @@ export function appendResponseHeader(
 
 export const appendHeader = appendResponseHeader;
 
-export function removeResponseHeaders(event: H3Event): void {
-  for (const [name] of Object.entries(getResponseHeaders(event))) {
-    removeResponseHeader(event, name);
+/**
+ * Remove all response headers, or only those specified in the headerNames array.
+ * @param event H3 event
+ * @param headerNames Array of header names to remove
+ */
+export function removeResponseHeaders(
+  event: H3Event,
+  headerNames?: string[]
+): void {
+  if (headerNames && headerNames.length > 0) {
+    for (const name of headerNames) {
+      removeResponseHeader(event, name);
+    }
+  } else {
+    for (const [name] of Object.entries(getResponseHeaders(event))) {
+      removeResponseHeader(event, name);
+    }
   }
 }
 
diff --git a/test/header.test.ts b/test/header.test.ts
index bf63cb92..602047aa 100644
--- a/test/header.test.ts
+++ b/test/header.test.ts
@@ -271,7 +271,7 @@ describe("", () => {
   });
 
   describe("removeResponseHeaders", () => {
-    it("can remove multiple response headers", async () => {
+    it("can remove all response headers", async () => {
       app.use(
         "/",
         eventHandler((event) => {
@@ -285,16 +285,34 @@ describe("", () => {
       expect(result.headers["header-2"]).toBeUndefined();
     });
 
-    it("can remove a single response header", async () => {
+    it("can remove multiple response headers", async () => {
       app.use(
         "/",
         eventHandler((event) => {
           appendHeader(event, "header-3", "3");
-          removeResponseHeader(event, "header-3");
+          appendHeader(event, "header-4", "4");
+          appendHeader(event, "header-5", "5");
+          removeResponseHeaders(event, ["header-3", "header-5"]);
         })
       );
       const result = await request.get("/");
       expect(result.headers["header-3"]).toBeUndefined();
+      expect(result.headers["header-4"]).toBe("4");
+      expect(result.headers["header-5"]).toBeUndefined();
+    });
+
+    it("can remove a single response header", async () => {
+      app.use(
+        "/",
+        eventHandler((event) => {
+          appendHeader(event, "header-6", "6");
+          appendHeader(event, "header-7", "7");
+          removeResponseHeader(event, "header-6");
+        })
+      );
+      const result = await request.get("/");
+      expect(result.headers["header-6"]).toBeUndefined();
+      expect(result.headers["header-7"]).toBe("7");
     });
   });
 });