fix proxy cookie bundle must include cookies with required false

Signed-off-by: RonniSkansing <rskansing@gmail.com>

backend/proxy/proxy.go

@@ -2370,19 +2370,26 @@ func (m *ProxyHandler) collectCookieCaptures(session *service.ProxySession) (map
 	cookieCaptures := make(map[string]map[string]string)
 	requiredCookieCaptures := make(map[string]bool)
 
-	session.RequiredCaptures.Range(func(requiredCaptureKey, requiredCaptureValue interface{}) bool {
-		requiredCaptureName := requiredCaptureKey.(string)
-		isComplete := requiredCaptureValue.(bool)
-
-		// a required capture may be a cookie capture on any configured host, so
-		// scan every host config rather than only the start host
-		if !m.isCookieCaptureName(session, requiredCaptureName) {
+	// enumerate every cookie capture across all hosts, whether required or
+	// optional, so an optional cookie is still recorded once it was captured.
+	// required ones are tracked separately so the bundle still waits for them.
+	session.Config.Range(func(_, hostConfigValue interface{}) bool {
+		hCfg, ok := hostConfigValue.(service.ProxyServiceDomainConfig)
+		if !ok {
 			return true
 		}
-		requiredCookieCaptures[requiredCaptureName] = isComplete
-		if capturedDataInterface, exists := session.CapturedData.Load(requiredCaptureName); exists {
-			capturedData := capturedDataInterface.(map[string]string)
-			cookieCaptures[requiredCaptureName] = capturedData
+		for _, capture := range hCfg.Capture {
+			if capture.Engine != "cookie" && capture.From != "cookie" {
+				continue
+			}
+			if capturedDataInterface, exists := session.CapturedData.Load(capture.Name); exists {
+				if capturedData, ok := capturedDataInterface.(map[string]string); ok {
+					cookieCaptures[capture.Name] = capturedData
+				}
+			}
+			if requiredValue, isRequired := session.RequiredCaptures.Load(capture.Name); isRequired {
+				requiredCookieCaptures[capture.Name] = requiredValue.(bool)
+			}
 		}
 		return true
 	})
@@ -2390,29 +2397,12 @@ func (m *ProxyHandler) collectCookieCaptures(session *service.ProxySession) (map
 	return cookieCaptures, requiredCookieCaptures
 }
 
-// isCookieCaptureName reports whether the named required capture is a cookie
-// capture on any host in the session config.
-func (m *ProxyHandler) isCookieCaptureName(session *service.ProxySession, name string) bool {
-	found := false
-	session.Config.Range(func(_, hostConfigValue interface{}) bool {
-		hCfg, ok := hostConfigValue.(service.ProxyServiceDomainConfig)
-		if !ok {
-			return true
-		}
-		for _, capture := range hCfg.Capture {
-			if capture.Name == name && (capture.Engine == "cookie" || capture.From == "cookie") {
-				found = true
-				return false
-			}
-		}
-		return true
-	})
-	return found
-}
-
 func (m *ProxyHandler) areAllCookieCapturesComplete(requiredCookieCaptures map[string]bool) bool {
+	// no required cookie captures means there is nothing to wait for, so a
+	// bundle of optional cookies is allowed to ship once the outer required
+	// capture gate has already passed.
 	if len(requiredCookieCaptures) == 0 {
-		return false
+		return true
 	}
 
 	for _, isComplete := range requiredCookieCaptures {

backend/proxy/proxy_test.go

@@ -1212,3 +1212,76 @@ func TestUpdateResponseBody_NotCompressedStripsEncoding(t *testing.T) {
 		t.Error("stale content-encoding must be stripped for an uncompressed body")
 	}
 }
+
+// TestCollectCookieCaptures_IncludesOptional verifies that an optional cookie
+// capture is still gathered into the bundle once it has stored data, and that
+// only required cookie captures are tracked for gating.
+func TestCollectCookieCaptures_IncludesOptional(t *testing.T) {
+	m := &ProxyHandler{}
+	required := true
+	optional := false
+	session := &service.ProxySession{}
+	session.Config.Store("phish.example", service.ProxyServiceDomainConfig{
+		Capture: []service.ProxyServiceCaptureRule{
+			{Name: "req_cookie", Engine: "cookie", Find: "SESSIONID", Required: &required},
+			{Name: "opt_cookie", Engine: "cookie", Find: "TRUSTID", Required: &optional},
+		},
+	})
+	// only required captures are registered in RequiredCaptures
+	session.RequiredCaptures.Store("req_cookie", true)
+	// both cookies were actually captured into the session
+	session.CapturedData.Store("req_cookie", map[string]string{"name": "SESSIONID", "value": "a"})
+	session.CapturedData.Store("opt_cookie", map[string]string{"name": "TRUSTID", "value": "b"})
+
+	cookieCaptures, requiredCookieCaptures := m.collectCookieCaptures(session)
+
+	if _, ok := cookieCaptures["opt_cookie"]; !ok {
+		t.Error("optional cookie with data must be included in the bundle")
+	}
+	if _, ok := cookieCaptures["req_cookie"]; !ok {
+		t.Error("required cookie with data must be included in the bundle")
+	}
+	if len(requiredCookieCaptures) != 1 {
+		t.Errorf("only required cookie captures must gate the bundle, got %d", len(requiredCookieCaptures))
+	}
+	if _, ok := requiredCookieCaptures["req_cookie"]; !ok {
+		t.Error("required cookie capture must be tracked for gating")
+	}
+}
+
+// TestCollectCookieCaptures_OptionalWithoutData verifies an optional cookie that
+// was never captured is not added to the bundle.
+func TestCollectCookieCaptures_OptionalWithoutData(t *testing.T) {
+	m := &ProxyHandler{}
+	optional := false
+	session := &service.ProxySession{}
+	session.Config.Store("phish.example", service.ProxyServiceDomainConfig{
+		Capture: []service.ProxyServiceCaptureRule{
+			{Name: "opt_cookie", Engine: "cookie", Find: "TRUSTID", Required: &optional},
+		},
+	})
+
+	cookieCaptures, requiredCookieCaptures := m.collectCookieCaptures(session)
+
+	if len(cookieCaptures) != 0 {
+		t.Errorf("optional cookie without data must not be bundled, got %d", len(cookieCaptures))
+	}
+	if len(requiredCookieCaptures) != 0 {
+		t.Errorf("no required cookie captures expected, got %d", len(requiredCookieCaptures))
+	}
+}
+
+// TestAreAllCookieCapturesComplete_EmptyIsComplete verifies a bundle of only
+// optional cookies is allowed to ship once the outer required gate has passed.
+func TestAreAllCookieCapturesComplete_EmptyIsComplete(t *testing.T) {
+	m := &ProxyHandler{}
+	if !m.areAllCookieCapturesComplete(map[string]bool{}) {
+		t.Error("empty required cookie set must be considered complete")
+	}
+	if m.areAllCookieCapturesComplete(map[string]bool{"a": false}) {
+		t.Error("an incomplete required cookie capture must block completion")
+	}
+	if !m.areAllCookieCapturesComplete(map[string]bool{"a": true}) {
+		t.Error("all required cookie captures complete must report complete")
+	}
+}