From a4cdc48542fad033ffbdd822930e8e69cf1bedb5 Mon Sep 17 00:00:00 2001
From: vcker <vcker.com@gmail.com>
Date: Mon, 6 May 2024 01:12:33 +0800
Subject: [PATCH] net: add a .port()! method for net.Addr (#21412)

---
 vlib/net/address.c.v      | 22 ++++++++++++++++++++++
 vlib/net/address_test.c.v |  5 +++++
 2 files changed, 27 insertions(+)

diff --git a/vlib/net/address.c.v b/vlib/net/address.c.v
index 3f8c5cf600..59770237c1 100644
--- a/vlib/net/address.c.v
+++ b/vlib/net/address.c.v
@@ -60,6 +60,28 @@ pub fn (a Addr) family() AddrFamily {
 	return unsafe { AddrFamily(a.f) }
 }
 
+// port returns the ip or ip6 port of the given address `a`
+pub fn (a Addr) port() !u16 {
+	match unsafe { AddrFamily(a.f) } {
+		.ip {
+			unsafe {
+				return conv.ntoh16(a.addr.Ip.port)
+			}
+		}
+		.ip6 {
+			unsafe {
+				return conv.ntoh16(a.addr.Ip6.port)
+			}
+		}
+		.unix {
+			return error('unix addr has no port')
+		}
+		.unspec {
+			return error('unspec addr family when obtain port')
+		}
+	}
+}
+
 const max_ip_len = 24
 const max_ip6_len = 46
 
diff --git a/vlib/net/address_test.c.v b/vlib/net/address_test.c.v
index deb5c06d79..e9140d1fec 100644
--- a/vlib/net/address_test.c.v
+++ b/vlib/net/address_test.c.v
@@ -1,5 +1,10 @@
 module net
 
+fn test_ip_port() {
+	assert new_ip(1234, addr_ip_any).port()! == 1234
+	assert new_ip6(1234, addr_ip6_any).port()! == 1234
+}
+
 fn test_diagnostics() {
 	dump(aoffset)
 	eprintln('--------')