From a755de3bbd49cf23f65b969f3bc188b4f117fb23 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Wed, 2 Aug 2023 19:39:18 +0800
Subject: [PATCH] Improve serializer

---
 common/metadata/family.go     |  7 +++--
 common/metadata/serializer.go | 55 +++++++++++++++++++++--------------
 2 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/common/metadata/family.go b/common/metadata/family.go
index c6551c2..c2a5bed 100644
--- a/common/metadata/family.go
+++ b/common/metadata/family.go
@@ -3,7 +3,8 @@ package metadata
 type Family = byte
 
 const (
-	AddressFamilyIPv4 Family = 0x01
-	AddressFamilyIPv6 Family = 0x04
-	AddressFamilyFqdn Family = 0x03
+	AddressFamilyIPv4  Family = 0x01
+	AddressFamilyIPv6  Family = 0x04
+	AddressFamilyFqdn  Family = 0x03
+	AddressFamilyEmpty Family = 0xff
 )
diff --git a/common/metadata/serializer.go b/common/metadata/serializer.go
index 62d31c1..9dec47d 100644
--- a/common/metadata/serializer.go
+++ b/common/metadata/serializer.go
@@ -50,14 +50,20 @@ func NewSerializer(options ...SerializerOption) *Serializer {
 
 func (s *Serializer) WriteAddress(buffer *buf.Buffer, addr Socksaddr) error {
 	var af Family
-	if addr.IsIPv4() {
+	if !addr.IsValid() {
+		af = AddressFamilyFqdn
+	} else if addr.IsIPv4() {
 		af = AddressFamilyIPv4
 	} else if addr.IsIPv6() {
 		af = AddressFamilyIPv6
 	} else {
 		af = AddressFamilyFqdn
 	}
-	err := buffer.WriteByte(s.familyByteMap[af])
+	afByte, loaded := s.familyByteMap[af]
+	if !loaded {
+		return E.New("unsupported address")
+	}
+	err := buffer.WriteByte(afByte)
 	if err != nil {
 		return err
 	}
@@ -70,7 +76,14 @@ func (s *Serializer) WriteAddress(buffer *buf.Buffer, addr Socksaddr) error {
 }
 
 func (s *Serializer) AddressLen(addr Socksaddr) int {
-	if addr.IsIPv4() {
+	if !addr.IsValid() {
+		_, supportEmpty := s.familyByteMap[AddressFamilyEmpty]
+		if supportEmpty {
+			return 1
+		} else {
+			return 2
+		}
+	} else if addr.IsIPv4() {
 		return 5
 	} else if addr.IsIPv6() {
 		return 17
@@ -129,26 +142,24 @@ func (s *Serializer) ReadAddress(reader io.Reader) (Socksaddr, error) {
 			return Socksaddr{}, E.Cause(err, "read fqdn")
 		}
 		return ParseSocksaddrHostPort(fqdn, 0), nil
-	default:
-		switch family {
-		case AddressFamilyIPv4:
-			var addr [4]byte
-			err = common.Error(reader.Read(addr[:]))
-			if err != nil {
-				return Socksaddr{}, E.Cause(err, "read ipv4 address")
-			}
-			return Socksaddr{Addr: netip.AddrFrom4(addr)}, nil
-		case AddressFamilyIPv6:
-			var addr [16]byte
-			err = common.Error(reader.Read(addr[:]))
-			if err != nil {
-				return Socksaddr{}, E.Cause(err, "read ipv6 address")
-			}
-
-			return Socksaddr{Addr: netip.AddrFrom16(addr)}.Unwrap(), nil
-		default:
-			return Socksaddr{}, E.New("unknown address family: ", af)
+	case AddressFamilyIPv4:
+		var addr [4]byte
+		_, err = io.ReadFull(reader, addr[:])
+		if err != nil {
+			return Socksaddr{}, E.Cause(err, "read ipv4 address")
 		}
+		return Socksaddr{Addr: netip.AddrFrom4(addr)}, nil
+	case AddressFamilyIPv6:
+		var addr [16]byte
+		_, err = io.ReadFull(reader, addr[:])
+		if err != nil {
+			return Socksaddr{}, E.Cause(err, "read ipv6 address")
+		}
+		return Socksaddr{Addr: netip.AddrFrom16(addr)}.Unwrap(), nil
+	case AddressFamilyEmpty:
+		return Socksaddr{}, nil
+	default:
+		return Socksaddr{}, E.New("unknown address family: ", af)
 	}
 }