From 38d83b44fff14a793fcf846880b9e6f3f9e0cffd Mon Sep 17 00:00:00 2001
From: ChronosX88 <chronosx88@gmail.com>
Date: Sun, 25 Apr 2021 00:09:42 +0300
Subject: [PATCH] Implement parsing Entity ID with additional attribute, add
 unit tests for it

---
 core/models/entity_id.go      | 66 +++++++++++++++++++++++++++++------
 core/models/entity_id_test.go | 39 +++++++++++++++++++++
 2 files changed, 94 insertions(+), 11 deletions(-)
 create mode 100644 core/models/entity_id_test.go

diff --git a/core/models/entity_id.go b/core/models/entity_id.go
index 49d4288..cec2044 100644
--- a/core/models/entity_id.go
+++ b/core/models/entity_id.go
@@ -11,17 +11,22 @@ const (
 	UsernameType  EntityIDType = "@"
 	ChatAliasType EntityIDType = "#"
 	ChatIDType    EntityIDType = "!"
+	ThirdPIDType  EntityIDType = "%"
+	OtherType     EntityIDType = "&"
 )
 
 type EntityID struct {
-	EntityIDType EntityIDType
-	LocalPart    string
-	ServerPart   string
+	Type       EntityIDType
+	LocalPart  string
+	ServerPart string
+	Attr       string
 }
 
 func NewEntityIDFromString(entityID string) (*EntityID, error) {
 	eid := &EntityID{}
 	typ := string(entityID[0])
+	withAttr := false
+
 	switch EntityIDType(typ) {
 	case UsernameType:
 		fallthrough
@@ -29,18 +34,33 @@ func NewEntityIDFromString(entityID string) (*EntityID, error) {
 		fallthrough
 	case ChatIDType:
 		{
-			eid.EntityIDType = EntityIDType(typ)
+			eid.Type = EntityIDType(typ)
+		}
+	case ThirdPIDType:
+		fallthrough
+	case OtherType:
+		{
+			eid.Type = EntityIDType(typ)
+			withAttr = true
 		}
 	default:
 		return nil, fmt.Errorf("invalid entity id type: %s", typ)
 	}
 
 	localAndServerPart := strings.Split(entityID, "@")
-	if len(localAndServerPart) == 3 {
-		localAndServerPart = localAndServerPart[1:]
+	if len(localAndServerPart) == 3 && localAndServerPart[0] == "" {
+		localAndServerPart = localAndServerPart[0:]
+	}
+	if !withAttr {
+		eid.LocalPart = localAndServerPart[0]
+		eid.ServerPart = localAndServerPart[1]
+	} else {
+		attrAndLocal := strings.Split(localAndServerPart[0], ":")
+		attr := attrAndLocal[0][1:]
+		eid.Attr = attr
+		eid.LocalPart = attrAndLocal[1]
+		eid.ServerPart = localAndServerPart[len(localAndServerPart)-1]
 	}
-	eid.LocalPart = localAndServerPart[0]
-	eid.ServerPart = localAndServerPart[1]
 
 	return eid, nil
 }
@@ -55,7 +75,7 @@ func NewEntityID(typ, localPart, serverPart string) (*EntityID, error) {
 		fallthrough
 	case ChatIDType:
 		{
-			eid.EntityIDType = EntityIDType(typ)
+			eid.Type = EntityIDType(typ)
 		}
 	default:
 		return nil, fmt.Errorf("invalid entity id type: %s", typ)
@@ -67,6 +87,30 @@ func NewEntityID(typ, localPart, serverPart string) (*EntityID, error) {
 	return eid, nil
 }
 
-func (eID *EntityID) String() string {
-	return fmt.Sprintf("%s%s@%s", eID.EntityIDType, eID.LocalPart, eID.ServerPart)
+func NewEntityIDWithAttr(typ, attr, localPart, serverPart string) (*EntityID, error) {
+	eid := &EntityID{}
+
+	switch EntityIDType(typ) {
+	case OtherType:
+		fallthrough
+	case ThirdPIDType:
+		{
+			eid.Type = EntityIDType(typ)
+		}
+	default:
+		return nil, fmt.Errorf("invalid entity id type: %s", typ)
+	}
+
+	eid.Attr = attr
+	eid.LocalPart = localPart
+	eid.ServerPart = serverPart
+
+	return eid, nil
+}
+
+func (eID *EntityID) String() string {
+	if eID.Attr != "" {
+		return fmt.Sprintf("%s%s:%s@%s", eID.Type, eID.Attr, eID.LocalPart, eID.ServerPart)
+	}
+	return fmt.Sprintf("%s%s@%s", eID.Type, eID.LocalPart, eID.ServerPart)
 }
diff --git a/core/models/entity_id_test.go b/core/models/entity_id_test.go
new file mode 100644
index 0000000..d03544c
--- /dev/null
+++ b/core/models/entity_id_test.go
@@ -0,0 +1,39 @@
+package models
+
+import "testing"
+
+func TestNewEntityIDFromString(t *testing.T) {
+	str := "@user@cadmium.org"
+
+	eid, err := NewEntityIDFromString(str)
+	if err != nil {
+		t.Fatal("error must be null")
+	}
+	if eid.Attr != "" || eid.Type != UsernameType || eid.ServerPart != "cadmium.org" || eid.LocalPart != "user" {
+		t.FailNow()
+	}
+}
+
+func TestNewEntityIDFromStringWithAttr(t *testing.T) {
+	str := "%msisdn:18002003040@cadmium.org"
+
+	eid, err := NewEntityIDFromString(str)
+	if err != nil {
+		t.Fatal("error must be null")
+	}
+	if eid.Attr != "msisdn" || eid.Type != ThirdPIDType || eid.ServerPart != "cadmium.org" || eid.LocalPart != "18002003040" {
+		t.FailNow()
+	}
+}
+
+func TestNewEntityIDFromStringWithEmailAttr(t *testing.T) {
+	str := "%email:abslimit_netwhood.online@cadmium.org"
+
+	eid, err := NewEntityIDFromString(str)
+	if err != nil {
+		t.Fatal("error must be null")
+	}
+	if eid.Attr != "email" || eid.Type != ThirdPIDType || eid.ServerPart != "cadmium.org" || eid.LocalPart != "abslimit_netwhood.online" {
+		t.Fatal(eid.String())
+	}
+}