From b2b386b9616ad5913f55c5eea9b76e6d8d0ffe2d Mon Sep 17 00:00:00 2001
From: mykola2312 <49044616+mykola2312@users.noreply.github.com>
Date: Sun, 19 Jan 2025 12:31:52 +0200
Subject: [PATCH] test IP primitives in proto buffer

---
 proto/lux_buffer.go      | 18 ++++++++++++++++++
 tests/lux_buffer_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/proto/lux_buffer.go b/proto/lux_buffer.go
index f9e708e..6f6a56a 100644
--- a/proto/lux_buffer.go
+++ b/proto/lux_buffer.go
@@ -246,3 +246,21 @@ func (buf *LuxBuffer) WriteIPPort(addrPort netip.AddrPort) {
 	buf.WriteIP(addrPort.Addr())
 	buf.WriteUint16(addrPort.Port())
 }
+
+func LuxProtoIPToAddr(ip net.IP) netip.Addr {
+	if octets4 := ip.To4(); octets4 != nil {
+		return netip.AddrFrom4([4]byte(octets4))
+	} else {
+		return netip.AddrFrom16([16]byte(ip.To16()))
+	}
+}
+
+func LuxProtoAddrToIP(addr netip.Addr) net.IP {
+	if addr.Is4() {
+		octets4 := addr.As4()
+		return octets4[:]
+	} else {
+		octets16 := addr.As16()
+		return octets16[:]
+	}
+}
diff --git a/tests/lux_buffer_test.go b/tests/lux_buffer_test.go
index 6a95287..9080668 100644
--- a/tests/lux_buffer_test.go
+++ b/tests/lux_buffer_test.go
@@ -3,6 +3,7 @@ package tests
 import (
 	"bytes"
 	"lux/proto"
+	"net"
 	"testing"
 )
 
@@ -71,3 +72,30 @@ func TestReadOverrun(t *testing.T) {
 		t.Fatalf("no error when reading at offset 1 of 1 byte buffer\n")
 	}
 }
+
+func TestIP(t *testing.T) {
+	ip4 := net.ParseIP("10.1.0.5")
+	ip6 := net.ParseIP("fd10:101::d47d:65ff:fe31:9367")
+
+	wd := proto.NewLuxBuffer()
+	wd.WriteIP(proto.LuxProtoIPToAddr(ip4))
+	wd.WriteIP(proto.LuxProtoIPToAddr(ip6))
+
+	rd := proto.FromSlice(wd.AllBytes())
+
+	newIp4, err := rd.ReadIP()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !proto.LuxProtoAddrToIP(newIp4).Equal(ip4) {
+		t.Fatalf("IPv4 does not equal: %v and %v", ip4.To4(), newIp4.AsSlice())
+	}
+
+	newIp6, err := rd.ReadIP()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !proto.LuxProtoAddrToIP(newIp6).Equal(ip6) {
+		t.Fatalf("IPv6 does not equal: %v and %v", ip4.To4(), newIp4.AsSlice())
+	}
+}