From 7361dac8110bfb49f732a19fa52dfeb3c2f04ccc Mon Sep 17 00:00:00 2001
From: atavism <paul@getlantern.org>
Date: Tue, 18 Jun 2024 07:56:16 -0700
Subject: [PATCH] clean-ups, add AddInternalHeaders

---
 internalsdk/common/headers.go | 15 ++++++++++-----
 internalsdk/pro/pro.go        | 12 ++++--------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/internalsdk/common/headers.go b/internalsdk/common/headers.go
index e135b32d6..7933c7409 100644
--- a/internalsdk/common/headers.go
+++ b/internalsdk/common/headers.go
@@ -49,11 +49,7 @@ var (
 func AddCommonNonUserHeaders(uc UserConfig, req *http.Request) {
 	req.Header.Set(AppVersionHeader, ApplicationVersion)
 	req.Header.Set(LibraryVersionHeader, LibraryVersion)
-	for k, v := range uc.GetInternalHeaders() {
-		if v != "" {
-			req.Header.Set(k, v)
-		}
-	}
+	AddInternalHeaders(uc, req)
 	if len(uc.GetEnabledExperiments()) > 0 {
 		req.Header.Set("x-lantern-dev-experiments", strings.Join(uc.GetEnabledExperiments(), ","))
 	}
@@ -74,6 +70,15 @@ func AddCommonNonUserHeaders(uc UserConfig, req *http.Request) {
 	req.Header.Add(RandomNoiseHeader, randomizedString())
 }
 
+// AddInternalHeaders adds the internal headers from the user config to the headers of the given http.Request
+func AddInternalHeaders(uc UserConfig, req *http.Request) {
+	for k, v := range uc.GetInternalHeaders() {
+		if v != "" {
+			req.Header.Set(k, v)
+		}
+	}
+}
+
 // AddCommonUserHeaders adds common headers that are user or device specific.
 func AddCommonUserHeaders(uc UserConfig, req *http.Request) {
 	if req.Header.Get(DeviceIdHeader) == "" {
diff --git a/internalsdk/pro/pro.go b/internalsdk/pro/pro.go
index 2bbe93ad5..ea3a0b04d 100644
--- a/internalsdk/pro/pro.go
+++ b/internalsdk/pro/pro.go
@@ -59,16 +59,12 @@ func NewClient(baseURL string, opts *webclient.Opts) ProClient {
 	client := &proClient{
 		userConfig: opts.UserConfig,
 	}
-	client.webclient = webclient.NewRESTClient(defaultwebclient.SendToURL(httpClient, baseURL, func(rc *resty.Client, req *resty.Request) error {
+	client.webclient = webclient.NewRESTClient(defaultwebclient.SendToURL(httpClient, baseURL, func(rc *resty.Client, restyReq *resty.Request) error {
 		uc := client.userConfig()
+		req := restyReq.RawRequest
 		req.Header.Set(common.Referer, "http://localhost:37457/")
-		common.AddCommonUserHeaders(uc, req.RawRequest)
-		headers := map[string]string{}
-		// Import all the internal headers
-		for k, v := range uc.GetInternalHeaders() {
-			headers[k] = v
-		}
-		req.SetHeaders(headers)
+		common.AddCommonUserHeaders(uc, req)
+		common.AddInternalHeaders(uc, req)
 		return nil
 	}, nil))
 	return client