Unmarshal Capabilities correctly in all structs that use them

Instead of targeting NamespaceAdV2 like I did previously, this change moves the unmarshal logic to the level of the Capabilities struct. Why I didn't do that in the first place? I couldn't tell ya In addition, I now test that JSON for ServerAd and OriginAdvertiseV2 (which may also arrive to the Director with old/new capabilities) unmarshal correctly.
jhiemstrawisc · Nov 5, 2024 · 38bcb71 · 38bcb71
1 parent 8e17764
commit 38bcb71
Show file tree

Hide file tree

Showing 2 changed files with 102 additions and 44 deletions.
diff --git a/server_structs/director.go b/server_structs/director.go
@@ -187,48 +187,36 @@ type (
 	}
 )
 
-// A helper function to handle JSON->NSAdV2 unmarshalling across multiple deprecated JSON keys.
+// A helper function to handle JSON->Caps unmarshalling across multiple deprecated JSON keys.
 //
-// When the Director sends a list of NamespaceAdV2 structs as JSON, it may contain
-// the old capabilities struct. This function checks if the raw JSON contains the
-// "FallbackRead" field, indicating the old capabilities struct, and unmarshals
-// the JSON into the new struct accordingly. This is to ensure backwards compatibility
-// We can probably think about removing this function when we don't see origins/directors
+// Old Caches/Origins will send various bits of JSON to the Director containing the JSON representation
+// of the OldCapabilities struct. To make sure these old JSON representations are unmarshalled correctly
+// into the new Capabilities struct, this function determines which JSON format is being sent and handles
+// conversion if necessary. We can probably think about removing this function when we don't see origins/directors
 // running Pelican <= v7.11.0.
-func (n *NamespaceAdV2) UnmarshalJSON(data []byte) error {
-	// Use alias struct of NamespaceAdV2 to prevent infinite recursion of UnmarshalJSON
-	type Alias NamespaceAdV2
-	aux := &struct {
-		Caps json.RawMessage `json:"Caps"`
-		*Alias
-	}{
-		Alias: (*Alias)(n),
-	}
-
-	if err := json.Unmarshal(data, &aux); err != nil {
-		return err
-	}
-
-	// Check if the raw Caps JSON contains "FallbackRead", indicating
-	// the old caps struct. Since the struct-to-JSON code has never included
-	// 'omitempty', we can safely assume that the field is ALWAYS present in
-	// old versions of the JSON.
-	if strings.Contains(string(aux.Caps), "FallBackRead") {
+func (c *Capabilities) UnmarshalJSON(data []byte) error {
+	// Check if it's the old format by looking for a unique field, e.g., "FallBackRead"
+	if strings.Contains(string(data), "FallBackRead") {
+		// Detected old JSON format, so unmarshal into OldCapabilities
 		var oldCaps OldCapabilities
-		if err := json.Unmarshal(aux.Caps, &oldCaps); err != nil {
+		if err := json.Unmarshal(data, &oldCaps); err != nil {
 			return err
 		}
 
-		// Map old capabilities to new capabilities
-		n.Caps.PublicReads = oldCaps.PublicRead
-		n.Caps.Reads = oldCaps.Read
-		n.Caps.Writes = oldCaps.Write
-		n.Caps.Listings = oldCaps.Listing
-		n.Caps.DirectReads = oldCaps.FallBackRead
+		// Map the old fields to the new struct
+		c.PublicReads = oldCaps.PublicRead
+		c.Reads = oldCaps.Read
+		c.Writes = oldCaps.Write
+		c.Listings = oldCaps.Listing
+		c.DirectReads = oldCaps.FallBackRead
 	} else {
-		if err := json.Unmarshal(aux.Caps, &n.Caps); err != nil {
+		// Assume it's the new JSON format and unmarshal directly into Capabilities
+		type Alias Capabilities // Create an alias to avoid recursive calls to UnmarshalJSON
+		var newCaps Alias
+		if err := json.Unmarshal(data, &newCaps); err != nil {
 			return err
 		}
+		*c = Capabilities(newCaps) // Copy unmarshalled values back to the original struct
 	}
 
 	return nil

diff --git a/server_structs/director_test.go b/server_structs/director_test.go
@@ -283,47 +283,117 @@ func TestXPelTokGenParsing(t *testing.T) {
 	})
 }
 
-func TestNSAdV2UnmarshalJSON(t *testing.T) {
-	oldJSON := `{"PublicRead":true,"Caps":{"PublicRead":true,"Read":true,"Write":false,"Listing":false,"FallBackRead":true},"path":"/ncar","token-generation":[{"strategy":"","vault-server":"","max-scope-depth":0,"issuer":{"Scheme":"","Opaque":"","User":null,"Host":"","Path":"","RawPath":"","OmitHost":false,"ForceQuery":false,"RawQuery":"","Fragment":"","RawFragment":""}}],"token-issuer":[],"from-topology":true}`
-	newJSON := `{"Caps":{"PublicReads":false,"Reads":true,"Writes":false,"Listings":false,"DirectReads":true},"path":"/ncar","token-generation":[{"strategy":"","vault-server":"","max-scope-depth":0,"issuer":{"Scheme":"","Opaque":"","User":null,"Host":"","Path":"","RawPath":"","OmitHost":false,"ForceQuery":false,"RawQuery":"","Fragment":"","RawFragment":""}}],"token-issuer":[],"from-topology":true}`
+func TestCapsUnmarshalJSON(t *testing.T) {
+	oldCaps := `{"PublicRead":true,"Read":true,"Write":false,"Listing":false,"FallBackRead":true}}`
+	newCaps := `{"PublicReads":false,"Reads":true,"Writes":false,"Listings":false,"DirectReads":true}}`
+
+	nsAdV2NoCap := `{"path":"/ncar","token-generation":[{"strategy":"","vault-server":"","max-scope-depth":0,"issuer":{}}],"token-issuer":[],"from-topology":true,"Caps":`
+	oAdV2NoCap := `{"name": "example-server","registry-prefix": "/origins/example-server","broker-url": "http://example-broker.com","data-url": "http://example-data.com","web-url": "http://example-web.com","namespaces": [{"path": "/example-namespace","token-generation": [],"token-issuer": [],"from-topology": false}],"token-issuer": [],"storageType": "POSIX","directorTest": false,"capabilities":`
+	sAdV2NoCap := `{"name": "example-server","storageType": "POSIX","directorTest": false,"auth_url": {"Scheme": "http", "Host": "example-auth.com"},"broker_url": {"Scheme": "http", "Host": "example-auth.com"},"url": {"Scheme": "http", "Host": "example-auth.com"},"web_url": {"Scheme": "http", "Host": "example-auth.com"},"type": "cache","latitude": 40.7128,"longitude": -74.0060,"from_topology": true,"io_load": 0.75,"capabilities":`
 
 	tests := []struct {
 		name     string
 		jsonData string
 		expected Capabilities
+		target   interface{}
 	}{
 		{
-			name:     "Old JSON format",
-			jsonData: oldJSON,
+			name:     "NamespaceAdV2 with old caps",
+			jsonData: nsAdV2NoCap + oldCaps,
+			expected: Capabilities{
+				PublicReads: true,
+				Reads:       true,
+				Writes:      false,
+				Listings:    false,
+				DirectReads: true,
+			},
+			target: &NamespaceAdV2{},
+		},
+		{
+			name:     "NamespaceAdV2 with new caps",
+			jsonData: nsAdV2NoCap + newCaps,
+			expected: Capabilities{
+				PublicReads: false,
+				Reads:       true,
+				Writes:      false,
+				Listings:    false,
+				DirectReads: true,
+			},
+			target: &NamespaceAdV2{},
+		},
+		{
+			name:     "OriginAdvertiseV2 with old caps",
+			jsonData: oAdV2NoCap + oldCaps,
 			expected: Capabilities{
 				PublicReads: true,
 				Reads:       true,
 				Writes:      false,
 				Listings:    false,
 				DirectReads: true,
 			},
+			target: &OriginAdvertiseV2{},
 		},
 		{
-			name:     "New JSON format",
-			jsonData: newJSON,
+			name:     "OriginAdvertiseV2 with new caps",
+			jsonData: oAdV2NoCap + newCaps,
 			expected: Capabilities{
 				PublicReads: false,
 				Reads:       true,
 				Writes:      false,
 				Listings:    false,
 				DirectReads: true,
 			},
+			target: &OriginAdvertiseV2{},
+		},
+		{
+			name:     "ServerAd with old caps",
+			jsonData: sAdV2NoCap + oldCaps,
+			expected: Capabilities{
+				PublicReads: true,
+				Reads:       true,
+				Writes:      false,
+				Listings:    false,
+				DirectReads: true,
+			},
+			target: &ServerAd{},
+		},
+		{
+			name:     "ServerAd with new caps",
+			jsonData: sAdV2NoCap + newCaps,
+			expected: Capabilities{
+				PublicReads: false,
+				Reads:       true,
+				Writes:      false,
+				Listings:    false,
+				DirectReads: true,
+			},
+			target: &ServerAd{},
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			var ns NamespaceAdV2
-			if err := json.Unmarshal([]byte(tt.jsonData), &ns); err != nil {
+			// Unmarshal JSON into the appropriate struct
+			if err := json.Unmarshal([]byte(tt.jsonData), tt.target); err != nil {
 				t.Fatalf("UnmarshalJSON() error = %v", err)
 			}
-			if ns.Caps != tt.expected {
-				t.Errorf("UnmarshalJSON() = %v, want %v", ns.Caps, tt.expected)
+
+			// Check Caps field in the struct
+			switch v := tt.target.(type) {
+			case *NamespaceAdV2:
+				if v.Caps != tt.expected {
+					t.Errorf("NamespaceAdV2 Caps = %v, want %v", v.Caps, tt.expected)
+				}
+			case *OriginAdvertiseV2:
+				if v.Caps != tt.expected {
+					t.Errorf("OriginAdvertiseV2 Caps = %v, want %v", v.Caps, tt.expected)
+				}
+			case *ServerAd:
+				if v.Caps != tt.expected {
+					t.Errorf("ServerAd Caps = %v, want %v", v.Caps, tt.expected)
+				}
+			default:
+				t.Errorf("Unknown struct type: %T", tt.target)
 			}
 		})
 	}