From 417d6e2d6132f587dc67ac45b6c7aca15f6139b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=8D=E8=B4=A5=E7=9A=84=E7=B1=B3?= <675491258@qq.com>
Date: Sat, 26 Dec 2020 22:01:24 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E7=BB=93=E6=9E=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md        |  2 +-
 attributetype.go | 27 ++++++++++++++++++-
 message.go       |  2 --
 nattype.go       | 69 ++++++++++++++++++++----------------------------
 4 files changed, 55 insertions(+), 45 deletions(-)

diff --git a/README.md b/README.md
index db92700..f3e7335 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@
 
 [nat-type]() 是 [RFC 3489][1]
 的go实现，从 [NatTypeDetector](https://github.com/cdnbye/NatTypeDetector) 移植而来。
-代码写得烂，思想跟不上，就这样了，补上一句一只企鹅狗司马
+代码写得烂，思想跟不上，就这样了
 
 ## 使用
 
diff --git a/attributetype.go b/attributetype.go
index acff60c..3d364f5 100644
--- a/attributetype.go
+++ b/attributetype.go
@@ -1,6 +1,6 @@
 package stun
 
-type AttributeType int
+type AttributeType uint
 
 const (
 	Undefined AttributeType = iota
@@ -19,3 +19,28 @@ const (
 	XorOnly
 	ServerName
 )
+
+var attributeTypeNames = []string{
+	"Undefined",
+	"MappedAddress",
+	"ResponseAddress",
+	"ChangeRequest",
+	"SourceAddress",
+	"ChangedAddress",
+	"Username",
+	"Password",
+	"MessageIntegrity",
+	"ErrorCode",
+	"UnknownAttribute",
+	"ReflectedFrom",
+	"XorMappedAddress",
+	"XorOnly",
+	"ServerName",
+}
+
+func (t AttributeType) String() string {
+	if t <= ServerName {
+		return attributeTypeNames[t]
+	}
+	return ""
+}
diff --git a/message.go b/message.go
index bd2f26f..9beaa4d 100644
--- a/message.go
+++ b/message.go
@@ -158,9 +158,7 @@ func (message *Message) Parse(data []byte) error {
 		offset += 2
 		switch attributeType {
 		case MappedAddress:
-			//                System.out.println("type == AttributeType.MappedAddress");
 			message.mappedAddress = parseIPAddr(data, offset)
-			//                System.out.println("mappedAddress " + mappedAddress.getAddress() + " " + mappedAddress.getPort());
 			offset += 8
 		case ResponseAddress:
 			// RESPONSE-ADDRESS
diff --git a/nattype.go b/nattype.go
index 5caa107..e1f3170 100644
--- a/nattype.go
+++ b/nattype.go
@@ -1,48 +1,29 @@
 package stun
 
-type NatType struct {
-	int
-	string
-}
+type NatType uint
 
-func (natType NatType) String() string {
-	return natType.string
-}
-
-var (
+const (
 	/// <summary>
 	/// UDP is always blocked.
 	/// </summary>
-	UdpBlocked = NatType{
-		int:    0,
-		string: "UdpBlocked",
-	}
+	UdpBlocked NatType = iota
 
 	/// <summary>
 	/// No NAT, public IP, no firewall.
 	/// </summary>
-	OpenInternet = NatType{
-		int:    1,
-		string: "OpenInternet",
-	}
+	OpenInternet
 
 	/// <summary>
 	/// No NAT, public IP, but symmetric UDP firewall.
 	/// </summary>
-	SymmetricUdpFirewall = NatType{
-		int:    2,
-		string: "SymmetricUdpFirewall",
-	}
+	SymmetricUdpFirewall
 
 	/// <summary>
 	/// A full cone NAT is one where all requests from the same internal IP address and port are
 	/// mapped to the same external IP address and port. Furthermore, any external host can send
 	/// a packet to the internal host, by sending a packet to the mapped external address.
 	/// </summary>
-	FullCone = NatType{
-		int:    3,
-		string: "FullCone",
-	}
+	FullCone
 
 	/// <summary>
 	/// A restricted cone NAT is one where all requests from the same internal IP address and
@@ -50,10 +31,7 @@ var (
 	/// host (with IP address X) can send a packet to the internal host only if the internal host
 	/// had previously sent a packet to IP address X.
 	/// </summary>
-	RestrictedCone = NatType{
-		int:    4,
-		string: "RestrictedCone",
-	}
+	RestrictedCone
 
 	/// <summary>
 	/// A port restricted cone NAT is like a restricted cone NAT, but the restriction
@@ -61,10 +39,7 @@ var (
 	/// address X and source port P, to the internal host only if the internal host had previously
 	/// sent a packet to IP address X and port P.
 	/// </summary>
-	PortRestrictedCone = NatType{
-		int:    5,
-		string: "PortRestrictedCone",
-	}
+	PortRestrictedCone
 
 	/// <summary>
 	/// A symmetric NAT is one where all requests from the same internal IP address and port,
@@ -73,13 +48,25 @@ var (
 	/// a different destination, a different mapping is used. Furthermore, only the external host that
 	/// receives a packet can send a UDP packet back to the internal host.
 	/// </summary>
-	Symmetric = NatType{
-		int:    6,
-		string: "Symmetric",
-	}
+	Symmetric
 
-	Unknown = NatType{
-		int:    7,
-		string: "Unknown",
-	}
+	Unknown
 )
+
+var natTypeNames = []string{
+	"UdpBlocked",
+	"OpenInternet",
+	"SymmetricUdpFirewall",
+	"FullCone",
+	"RestrictedCone",
+	"PortRestrictedCone",
+	"Symmetric",
+	"Unknown",
+}
+
+func (natType NatType) String() string {
+	if natType <= Unknown {
+		return natTypeNames[natType]
+	}
+	return ""
+}