From cfdec3796838e02da2088e151872df275e7e55f6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Przemys=C5=82aw=20=C5=BBydek?= <przemyslawzydek@gmail.com>
Date: Fri, 14 Jun 2024 10:54:40 +0200
Subject: [PATCH] chore: strip port from ipv6 as well

---
 proxy/utils/headers.ts |  9 +--------
 proxy/utils/ip.test.ts | 23 +++++++++++++++++++++++
 proxy/utils/ip.ts      | 28 ++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 8 deletions(-)
 create mode 100644 proxy/utils/ip.test.ts
 create mode 100644 proxy/utils/ip.ts

diff --git a/proxy/utils/headers.ts b/proxy/utils/headers.ts
index 516d3e0..b7169f5 100644
--- a/proxy/utils/headers.ts
+++ b/proxy/utils/headers.ts
@@ -2,6 +2,7 @@ import * as http from 'http'
 import { HttpRequest, HttpRequestHeaders, HttpResponseHeaders, Logger } from '@azure/functions'
 import { updateCacheControlHeader } from './cacheControl'
 import { filterCookie } from './cookies'
+import { stripPort } from './ip'
 
 const CACHE_CONTROL_HEADER_NAME = 'cache-control'
 
@@ -80,14 +81,6 @@ function resolveClientIp(request: HttpRequest, logger?: Logger) {
   return stripPort(clientIp)
 }
 
-function stripPort(ip: string) {
-  if (!ip.includes(':')) {
-    return ip
-  }
-
-  return ip.split(':')[0]
-}
-
 export function getHost(request: Pick<HttpRequest, 'headers' | 'url'>) {
   return request.headers['x-forwarded-host'] || request.headers.host || new URL(request.url).hostname
 }
diff --git a/proxy/utils/ip.test.ts b/proxy/utils/ip.test.ts
new file mode 100644
index 0000000..5922d7a
--- /dev/null
+++ b/proxy/utils/ip.test.ts
@@ -0,0 +1,23 @@
+import { stripPort } from './ip'
+
+describe('Strip port', () => {
+  it('strip port from ipv4 address', () => {
+    expect(stripPort('237.84.2.178:80')).toBe('237.84.2.178')
+  })
+
+  it('ipv6 without port', () => {
+    expect(stripPort('5be8:dde9:7f0b:d5a7:bd01:b3be:9c69:573b')).toBe('5be8:dde9:7f0b:d5a7:bd01:b3be:9c69:573b')
+  })
+
+  it('ipv4 without port', () => {
+    expect(stripPort('237.84.2.178')).toBe('237.84.2.178')
+  })
+
+  it('strip port from ipv6 address', () => {
+    expect(stripPort('[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:443')).toBe('2001:0db8:85a3:0000:0000:8a2e:0370:7334')
+  })
+
+  it.each(['127.0', 'invalid', 'localhost', '2001:0db8:85a3:0000:0000'])('invalid ip: %s', (data) => {
+    expect(stripPort(data)).toBe(data)
+  })
+})
diff --git a/proxy/utils/ip.ts b/proxy/utils/ip.ts
new file mode 100644
index 0000000..cf8015e
--- /dev/null
+++ b/proxy/utils/ip.ts
@@ -0,0 +1,28 @@
+import { isIPv4, isIPv6 } from 'net'
+
+export function stripPort(ip: string) {
+  // Check if it's an IPv6 address with a port
+  if (ip.startsWith('[')) {
+    // IPv6 address with port
+    const closingBracketIndex = ip.indexOf(']')
+    if (closingBracketIndex !== -1) {
+      return ip.substring(1, closingBracketIndex)
+    }
+  } else {
+    // IPv4 address with port or IPv6 without brackets
+    const colonIndex = ip.lastIndexOf(':')
+    if (colonIndex !== -1) {
+      const ipWithoutPort = ip.substring(0, colonIndex)
+      // Validate if the part before the colon is a valid IPv4 or IPv6 address
+      if (isValidIp(ipWithoutPort)) {
+        return ipWithoutPort
+      }
+    }
+  }
+  // If no port is found, return the original IP
+  return ip
+}
+
+function isValidIp(ip: string) {
+  return isIPv4(ip) || isIPv6(ip)
+}