protocol/network: Fix deserialization of CookiePartitionKey

Fixes #149

protocol/network/util.go

@@ -10,3 +10,40 @@ func (n Headers) Map() (map[string]string, error) {
 	err := json.Unmarshal(n, &m)
 	return m, err
 }
+
+// UnmarhsalJSON implements json.Unmarshaler for CookiePartitionKey. The
+// protocol incorrectly defines CookiePartitionKey as an object but its
+// return type seems to be string. The ResponseReceivedExtraInfoReply
+// field CookiePartitionKeyOpaque tells us if it is a string, however
+// it's not clear if an object can be returned. For now we will keep the
+// struct just in case and unmarhsla into TopLevelSite.
+//
+// https://chromium.googlesource.com/chromium/src/+/refs/heads/main/net/cookies/cookie_partition_key.cc
+// https://chromium.googlesource.com/chromium/src/+/refs/heads/main/net/cookies/cookie_partition_key.h
+func (k *CookiePartitionKey) UnmarshalJSON(data []byte) error {
+	type alias CookiePartitionKey
+	var v alias
+	if err := json.Unmarshal(data, &v); err != nil {
+		var s string
+		if err2 := json.Unmarshal(data, &s); err2 != nil {
+			return err
+		}
+
+		// TODO(mafredri): Is this correct or can s be parsed further to
+		// determine if it has a cross-site ancestor?
+		*k = CookiePartitionKey{
+			TopLevelSite:         s,
+			HasCrossSiteAncestor: false,
+		}
+
+		return nil
+	}
+
+	*k = CookiePartitionKey(v)
+
+	return nil
+}
+
+func (k CookiePartitionKey) String() string {
+	return k.TopLevelSite
+}

protocol/network/util_test.go

@@ -0,0 +1,47 @@
+package network
+
+import "testing"
+
+func TestCookiePartitionKey_UnmarshalJSON(t *testing.T) {
+	tests := []struct {
+		name    string
+		data    []byte
+		want    CookiePartitionKey
+		wantErr bool
+	}{
+		{
+			name: "string",
+			data: []byte(`"example.com"`),
+			want: CookiePartitionKey{
+				TopLevelSite:         "example.com",
+				HasCrossSiteAncestor: false,
+			},
+			wantErr: false,
+		},
+		{
+			name: "object",
+			data: []byte(`{"topLevelSite":"example.com","hasCrossSiteAncestor":true}`),
+			want: CookiePartitionKey{
+				TopLevelSite:         "example.com",
+				HasCrossSiteAncestor: true,
+			},
+			wantErr: false,
+		},
+		{
+			name:    "invalid",
+			data:    []byte(`invalid`),
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var k CookiePartitionKey
+			if err := k.UnmarshalJSON(tt.data); (err != nil) != tt.wantErr {
+				t.Errorf("CookiePartitionKey.UnmarshalJSON() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if k != tt.want {
+				t.Errorf("CookiePartitionKey.UnmarshalJSON() = %v, want %v", k, tt.want)
+			}
+		})
+	}
+}