From 72d2f4e7e3552b6e37af56b852c0a50e7eeb5509 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Wed, 26 Jul 2023 18:23:15 +0700
Subject: [PATCH] internal/controld: add support for parsing client id from raw
 UID

---
 internal/controld/config.go        | 23 +++++++++++++++++---
 internal/controld/config_test.go   | 29 ++++++++++++-------------
 internal/controld/controld_test.go | 34 ++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 19 deletions(-)
 create mode 100644 internal/controld/controld_test.go

diff --git a/internal/controld/config.go b/internal/controld/config.go
index 852aa8a..320fd4c 100644
--- a/internal/controld/config.go
+++ b/internal/controld/config.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/Control-D-Inc/ctrld"
@@ -53,12 +54,18 @@ func (u UtilityErrorResponse) Error() string {
 }
 
 type utilityRequest struct {
-	UID string `json:"uid"`
+	UID      string `json:"uid"`
+	ClientID string `json:"client_id,omitempty"`
 }
 
 // FetchResolverConfig fetch Control D config for given uid.
-func FetchResolverConfig(uid, version string, cdDev bool) (*ResolverConfig, error) {
-	body, _ := json.Marshal(utilityRequest{UID: uid})
+func FetchResolverConfig(rawUID, version string, cdDev bool) (*ResolverConfig, error) {
+	uid, clientID := ParseRawUID(rawUID)
+	uReq := utilityRequest{UID: uid}
+	if clientID != "" {
+		uReq.ClientID = clientID
+	}
+	body, _ := json.Marshal(uReq)
 	apiUrl := resolverDataURLCom
 	if cdDev {
 		apiUrl = resolverDataURLDev
@@ -120,3 +127,13 @@ func FetchResolverConfig(uid, version string, cdDev bool) (*ResolverConfig, erro
 	}
 	return &ur.Body.Resolver, nil
 }
+
+// ParseRawUID parse the input raw UID, returning real UID and ClientID.
+// The raw UID can have 2 forms:
+//
+// - <uid>
+// - <uid>/<client_id>
+func ParseRawUID(rawUID string) (string, string) {
+	uid, clientID, _ := strings.Cut(rawUID, "/")
+	return uid, clientID
+}
diff --git a/internal/controld/config_test.go b/internal/controld/config_test.go
index 2c00247..b266142 100644
--- a/internal/controld/config_test.go
+++ b/internal/controld/config_test.go
@@ -1,34 +1,31 @@
-//go:build controld
-
 package controld
 
 import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
-func TestFetchResolverConfig(t *testing.T) {
+func Test_parseUID(t *testing.T) {
 	tests := []struct {
-		name    string
-		uid     string
-		dev     bool
-		wantErr bool
+		name         string
+		uid          string
+		wantUID      string
+		wantClientID string
 	}{
-		{"valid com", "p2", false, false},
-		{"valid dev", "p2", true, false},
-		{"invalid uid", "abcd1234", false, true},
+		{"empty", "", "", ""},
+		{"only uid", "abcd1234", "abcd1234", ""},
+		{"with client id", "abcd1234/clientID", "abcd1234", "clientID"},
+		{"with empty clientID", "abcd1234/", "abcd1234", ""},
 	}
+
 	for _, tc := range tests {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			got, err := FetchResolverConfig(tc.uid, "dev-test", tc.dev)
-			require.False(t, (err != nil) != tc.wantErr, err)
-			if !tc.wantErr {
-				assert.NotEmpty(t, got.DOH)
-			}
+			gotUID, gotClientID := ParseRawUID(tc.uid)
+			assert.Equal(t, tc.wantUID, gotUID)
+			assert.Equal(t, tc.wantClientID, gotClientID)
 		})
 	}
 }
diff --git a/internal/controld/controld_test.go b/internal/controld/controld_test.go
new file mode 100644
index 0000000..2c00247
--- /dev/null
+++ b/internal/controld/controld_test.go
@@ -0,0 +1,34 @@
+//go:build controld
+
+package controld
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestFetchResolverConfig(t *testing.T) {
+	tests := []struct {
+		name    string
+		uid     string
+		dev     bool
+		wantErr bool
+	}{
+		{"valid com", "p2", false, false},
+		{"valid dev", "p2", true, false},
+		{"invalid uid", "abcd1234", false, true},
+	}
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			got, err := FetchResolverConfig(tc.uid, "dev-test", tc.dev)
+			require.False(t, (err != nil) != tc.wantErr, err)
+			if !tc.wantErr {
+				assert.NotEmpty(t, got.DOH)
+			}
+		})
+	}
+}