From b93a3a2550aa51f238a633ff192f1d9d643831cd Mon Sep 17 00:00:00 2001
From: Mickael Remond <mremond@process-one.net>
Date: Wed, 5 Jun 2019 10:02:24 +0200
Subject: [PATCH] Improve JID parsing

Clean up tests
Fix #1
---
 jid.go      | 51 ++++++++++++++++++++++++++++++++++++++++-----------
 jid_test.go | 48 +++++++++++++++++++++++++++---------------------
 2 files changed, 67 insertions(+), 32 deletions(-)

diff --git a/jid.go b/jid.go
index b795a3d..9594fbd 100644
--- a/jid.go
+++ b/jid.go
@@ -3,6 +3,7 @@ package xmpp // import "gosrc.io/xmpp"
 import (
 	"errors"
 	"strings"
+	"unicode"
 )
 
 type Jid struct {
@@ -11,24 +12,52 @@ type Jid struct {
 	resource string
 }
 
-func NewJid(sjid string) (jid *Jid, err error) {
-	s1 := strings.Split(sjid, "@")
+func NewJid(sjid string) (*Jid, error) {
+	s1 := strings.SplitN(sjid, "@", 2)
 	if len(s1) != 2 {
-		err = errors.New("invalid JID: " + sjid)
-		return
+		return nil, errors.New("invalid JID, missing domain: " + sjid)
 	}
-	jid = new(Jid)
+	jid := new(Jid)
 	jid.username = s1[0]
-
-	s2 := strings.Split(s1[1], "/")
-	if len(s2) > 2 {
-		err = errors.New("invalid JID: " + sjid)
-		return
+	if !isUsernameValid(jid.username) {
+		return jid, errors.New("invalid domain: " + jid.username)
 	}
+
+	s2 := strings.SplitN(s1[1], "/", 2)
+
 	jid.domain = s2[0]
+	if !isDomainValid(jid.domain) {
+		return jid, errors.New("invalid domain: " + jid.domain)
+	}
+
 	if len(s2) == 2 {
 		jid.resource = s2[1]
 	}
 
-	return
+	return jid, nil
+}
+
+func isUsernameValid(username string) bool {
+	invalidRunes := []rune{'@', '/', '\'', '"', ':', '<', '>'}
+	return strings.IndexFunc(username, isInvalid(invalidRunes)) < 0
+}
+
+func isDomainValid(domain string) bool {
+	invalidRunes := []rune{'@', '/'}
+	return strings.IndexFunc(domain, isInvalid(invalidRunes)) < 0
+}
+
+func isInvalid(invalidRunes []rune) func(c rune) bool {
+	isInvalid := func(c rune) bool {
+		if unicode.IsSpace(c) {
+			return true
+		}
+		for _, r := range invalidRunes {
+			if c == r {
+				return true
+			}
+		}
+		return false
+	}
+	return isInvalid
 }
diff --git a/jid_test.go b/jid_test.go
index 705e9f9..d74543f 100644
--- a/jid_test.go
+++ b/jid_test.go
@@ -5,43 +5,49 @@ import (
 )
 
 func TestValidJids(t *testing.T) {
-	var jid *Jid
-	var err error
+	tests := []struct {
+		jidstr   string
+		expected Jid
+	}{
+		{jidstr: "test@domain.com", expected: Jid{"test", "domain.com", ""}},
+		{jidstr: "test@domain.com/resource", expected: Jid{"test", "domain.com", "resource"}},
+		// resource can contain '/' or '@'
+		{jidstr: "test@domain.com/a/b", expected: Jid{"test", "domain.com", "a/b"}},
+		{jidstr: "test@domain.com/a@b", expected: Jid{"test", "domain.com", "a@b"}},
+	}
 
-	goodJids := []string{"test@domain.com", "test@domain.com/resource"}
-
-	for i, sjid := range goodJids {
-		if jid, err = NewJid(sjid); err != nil {
-			t.Error("could not parse correct jid")
-			return
+	for _, tt := range tests {
+		jid, err := NewJid(tt.jidstr)
+		if err != nil {
+			t.Errorf("could not parse correct jid: %s", tt.jidstr)
+			continue
 		}
 
 		if jid == nil {
 			t.Error("jid should not be nil")
 		}
 
-		if jid.username != "test" {
-			t.Error("incorrect jid username")
+		if jid.username != tt.expected.username {
+			t.Errorf("incorrect jid username (%s): %s", tt.expected.username, jid.username)
 		}
 
-		if jid.domain != "domain.com" {
-			t.Error("incorrect jid domain")
+		if jid.username != tt.expected.username {
+			t.Errorf("incorrect jid domain (%s): %s", tt.expected.domain, jid.domain)
 		}
 
-		if i == 0 && jid.resource != "" {
-			t.Error("bare jid resource should be empty")
-		}
-
-		if i == 1 && jid.resource != "resource" {
-			t.Error("incorrect full jid resource")
+		if jid.resource != tt.expected.resource {
+			t.Errorf("incorrect jid resource (%s): %s", tt.expected.resource, jid.resource)
 		}
 	}
 }
 
-// TODO: Check if resource cannot contain a /
 func TestIncorrectJids(t *testing.T) {
-	badJids := []string{"test@domain.com@otherdomain.com",
-		"test@domain.com/test/test"}
+	badJids := []string{
+		"user:name@domain.com",
+		"user<name@domain.com",
+		"test@domain.com@otherdomain.com",
+		"test@domain com/resource",
+	}
 
 	for _, sjid := range badJids {
 		if _, err := NewJid(sjid); err == nil {