From 950dac05e596addbe3b02d5c9d21e71100a5de32 Mon Sep 17 00:00:00 2001
From: Alex Sergeyev <asergeyev@dyn.com>
Date: Mon, 15 Sep 2014 08:16:01 -0400
Subject: [PATCH] Changed label splitting and fixed behavior for final period
 in names.

---
 idn/punycode.go      | 32 ++++++++++++++++++++++++--------
 idn/punycode_test.go | 16 ++++++++++++++++
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/idn/punycode.go b/idn/punycode.go
index 69c689cb..e70b9012 100644
--- a/idn/punycode.go
+++ b/idn/punycode.go
@@ -3,6 +3,8 @@ package idn
 
 import (
 	"bytes"
+	"github.com/miekg/dns"
+	"strings"
 	"unicode"
 )
 
@@ -28,20 +30,34 @@ const (
 // This function would return incorrect result for strings for non-canonical
 // unicode strings.
 func ToPunycode(s string) string {
-	tokens := bytes.Split([]byte(s), []byte{'.'})
-	for i := range tokens {
-		tokens[i] = encodeBytes(tokens[i])
+	tokens := dns.SplitDomainName(s)
+	if s[len(s)-1] == '.' {
+		if tokens == nil {
+			tokens = []string{"", ""}
+		} else {
+			tokens = append(tokens, "")
+		}
 	}
-	return string(bytes.Join(tokens, []byte{'.'}))
+	for i := range tokens {
+		tokens[i] = string(encodeBytes([]byte(tokens[i])))
+	}
+	return strings.Join(tokens, ".")
 }
 
 // FromPunycode returns uncode domain name from provided punycode string.
 func FromPunycode(s string) string {
-	tokens := bytes.Split([]byte(s), []byte{'.'})
-	for i := range tokens {
-		tokens[i] = decodeBytes(tokens[i])
+	tokens := dns.SplitDomainName(s)
+	if s[len(s)-1] == '.' {
+		if tokens == nil {
+			tokens = []string{"", ""}
+		} else {
+			tokens = append(tokens, "")
+		}
 	}
-	return string(bytes.Join(tokens, []byte{'.'}))
+	for i := range tokens {
+		tokens[i] = string(decodeBytes([]byte(tokens[i])))
+	}
+	return strings.Join(tokens, ".")
 }
 
 // digitval converts single byte into meaningful value that's used to calculate decoded unicode character.
diff --git a/idn/punycode_test.go b/idn/punycode_test.go
index 80541074..bd764a9d 100644
--- a/idn/punycode_test.go
+++ b/idn/punycode_test.go
@@ -6,6 +6,7 @@ import (
 )
 
 var testcases = [][2]string{
+	{"", ""},
 	{"a", "a"},
 	{"A-B", "a-b"},
 	{"AbC", "abc"},
@@ -53,6 +54,21 @@ func TestToFromPunycode(t *testing.T) {
 	}
 }
 
+func TestEncodeDecodeFinalPeriod(t *testing.T) {
+	for _, tst := range testcases {
+		// assert unicode.com. == punycode.com.
+		full := ToPunycode(tst[0] + ".")
+		if full != tst[1]+"." {
+			t.Errorf("invalid result from string conversion to punycode when period added at the end, %#v and should be %#v", full, tst[1]+".")
+		}
+		// assert punycode.com. == unicode.com.
+		decoded := FromPunycode(tst[1] + ".")
+		if decoded != strings.ToLower(tst[0]+".") {
+			t.Errorf("invalid result from string conversion to punycode when period added, %#v and should be %#v", decoded, tst[0]+".")
+		}
+	}
+}
+
 var invalid = []string{
 	"xn--*",
 	"xn--",