From ba2156133af98cd49921077fde33bf0499ccab7a Mon Sep 17 00:00:00 2001
From: Bel LaPointe <breel@qualtrics.com>
Date: Thu, 27 Jan 2022 11:00:05 -0700
Subject: [PATCH] cookie jar go

---
 broker/broker.go      | 22 ++++++++------------
 broker/broker_test.go | 48 ++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 52 insertions(+), 18 deletions(-)

diff --git a/broker/broker.go b/broker/broker.go
index 0d7a37c..1c0ab52 100644
--- a/broker/broker.go
+++ b/broker/broker.go
@@ -26,6 +26,10 @@ type Broker interface {
 }
 
 func do(r *http.Request) (*http.Response, error) {
+	return _do(config.Get().DB(), r)
+}
+
+func _do(db storage.DB, r *http.Request) (*http.Response, error) {
 	limiter.Wait(context.Background())
 	if strings.Contains(strings.ToLower(r.URL.Path), "login") {
 		authlimiter.Wait(context.Background())
@@ -39,8 +43,8 @@ func do(r *http.Request) (*http.Response, error) {
 	}
 	client.Jar = newjar
 
-	cookieJarKey := r.URL.Host
-	cookies, err := getCookies(cookieJarKey)
+	cookieJarKey := "cookies_" + r.URL.Host
+	cookies, err := getCookies(db, cookieJarKey)
 	if err != nil {
 		logtr.Errorf("failed to load cookies: %v", err)
 	} else {
@@ -50,17 +54,13 @@ func do(r *http.Request) (*http.Response, error) {
 	if err != nil {
 		return nil, err
 	}
-	if err := setCookies(cookieJarKey, client.Jar.Cookies(r.URL)); err != nil {
+	if err := setCookies(db, cookieJarKey, client.Jar.Cookies(r.URL)); err != nil {
 		logtr.Errorf("failed to set cookies: %v", err)
 	}
 	return resp, err
 }
 
-func getCookies(host string) ([]*http.Cookie, error) {
-	return _getCookies(config.Get().DB(), host)
-}
-
-func _getCookies(db storage.DB, host string) ([]*http.Cookie, error) {
+func getCookies(db storage.DB, host string) ([]*http.Cookie, error) {
 	b, err := db.Get(host)
 	if err != nil {
 		return nil, err
@@ -70,11 +70,7 @@ func _getCookies(db storage.DB, host string) ([]*http.Cookie, error) {
 	return result, err
 }
 
-func setCookies(host string, cookies []*http.Cookie) error {
-	return _setCookies(config.Get().DB(), host, cookies)
-}
-
-func _setCookies(db storage.DB, host string, cookies []*http.Cookie) error {
+func setCookies(db storage.DB, host string, cookies []*http.Cookie) error {
 	b, err := json.Marshal(cookies)
 	if err != nil {
 		return err
diff --git a/broker/broker_test.go b/broker/broker_test.go
index 703c20a..9ef8085 100644
--- a/broker/broker_test.go
+++ b/broker/broker_test.go
@@ -1,21 +1,59 @@
 package broker
 
 import (
+	"fmt"
 	"local/storage"
+	"net/http"
+	"net/http/httptest"
+	"strings"
 	"testing"
+	"time"
+
+	"golang.org/x/time/rate"
 )
 
 func TestDoCookies(t *testing.T) {
-	// func _setCookies(db storage.DB, host string, cookies []*http.Cookie) error {
-	// func _getCookies(db storage.DB, host string) cookies, error {
+	limiter = rate.NewLimiter(rate.Limit(20.0), 1)
+	calls := 0
 	db := storage.NewMap()
-	host := "host"
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		calls += 1
+		if calls == 1 {
+			http.SetCookie(w, &http.Cookie{
+				Name:    "name",
+				Value:   "value",
+				Expires: time.Now().Add(time.Hour),
+			})
+		} else {
+			if !strings.Contains(fmt.Sprint(r.Header["Cookie"]), "name=value") {
+				w.WriteHeader(http.StatusBadRequest)
+				t.Error("cookie not set on calls after first")
+			}
+		}
+	}))
+	defer s.Close()
+	req, _ := http.NewRequest(http.MethodGet, s.URL, nil)
 
-	if cookies, err := _getCookies(db, host); err == nil {
+	if cookies, err := getCookies(db, "cookies_"+req.URL.Host); err == nil {
 		t.Fatal(err)
 	} else if len(cookies) != 0 {
 		t.Fatal(cookies)
 	}
 
-	t.Fatal("not impl")
+	for i := 0; i < 3; i++ {
+		resp, err := _do(db, req)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if resp.StatusCode != http.StatusOK {
+			t.Fatal(resp.StatusCode)
+		}
+		resp.Body.Close()
+
+		if cookies, err := getCookies(db, "cookies_"+req.URL.Host); err != nil {
+			t.Fatal(err)
+		} else if len(cookies) == 0 {
+			t.Fatal(cookies)
+		}
+	}
 }