Merge pull request cedar-policy#64 from strongdm/json-encoding-cleanup

types: clean up JSON unmarshaling story for extension types
strongdm · Nov 12, 2024 · c7cc1e3 · c7cc1e3
2 parents 96c854e + 3b282ac
commit c7cc1e3
Show file tree

Hide file tree

Showing 10 changed files with 235 additions and 543 deletions.
diff --git a/types/datetime.go b/types/datetime.go
@@ -1,9 +1,7 @@
 package types
 
 import (
-	"bytes"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"strconv"
 	"time"
@@ -247,46 +245,16 @@ func (a Datetime) String() string {
 
 // UnmarshalJSON implements encoding/json.Unmarshaler for Datetime
 //
-// It is capable of unmarshaling 4 different representations supported by Cedar
-// - { "__extn": { "fn": "datetime", "arg": "1970-01-01" }}
-// - { "fn": "datetime", "arg": "1970-01-01" }
-// - "datetime(\"1970-01-01\")"
-// - "1970-01-01"
+// It is capable of unmarshaling 3 different representations supported by Cedar
+//   - { "__extn": { "fn": "datetime", "arg": "1970-01-01" }}
+//   - { "fn": "datetime", "arg": "1970-01-01" }
+//   - "1970-01-01"
 func (a *Datetime) UnmarshalJSON(b []byte) error {
-	var arg string
-	if bytes.HasPrefix(b, []byte(`"datetime(\"`)) && bytes.HasSuffix(b, []byte(`\")"`)) {
-		arg = string(b[12 : len(b)-4])
-	} else if len(b) > 0 && b[0] == '"' {
-		if err := json.Unmarshal(b, &arg); err != nil {
-			return errors.Join(errJSONDecode, err)
-		}
-	} else {
-		var res extValueJSON
-		if err := json.Unmarshal(b, &res); err != nil {
-			return errors.Join(errJSONDecode, err)
-		}
-		if res.Extn == nil {
-			// If we didn't find an Extn, maybe it's just an extn.
-			var res2 extn
-			json.Unmarshal(b, &res2)
-			// We've tried Ext.Fn and Fn, so no good.
-			if res2.Fn == "" {
-				return errJSONExtNotFound
-			}
-			if res2.Fn != "datetime" {
-				return errJSONExtFnMatch
-			}
-			arg = res2.Arg
-		} else if res.Extn.Fn != "datetime" {
-			return errJSONExtFnMatch
-		} else {
-			arg = res.Extn.Arg
-		}
-	}
-	aa, err := ParseDatetime(arg)
+	aa, err := unmarshalExtensionValue(b, "datetime", ParseDatetime)
 	if err != nil {
 		return err
 	}
+
 	*a = aa
 	return nil
 }

diff --git a/types/datetime_test.go b/types/datetime_test.go
@@ -1,6 +1,7 @@
 package types_test
 
 import (
+	"encoding/json"
 	"fmt"
 	"testing"
 	"time"
@@ -197,8 +198,25 @@ func TestDatetime(t *testing.T) {
 
 	t.Run("MarshalJSON", func(t *testing.T) {
 		t.Parallel()
-		bs, err := types.NewDatetime(time.UnixMilli(42)).MarshalJSON()
+		expected := `{
+			"__extn": {
+				"fn": "datetime",
+				"arg": "1970-01-01T00:00:00.042Z"
+			}
+		}`
+		dt1 := types.NewDatetime(time.UnixMilli(42))
+		testutil.JSONMarshalsTo(t, dt1, expected)
+
+		var dt2 types.Datetime
+		err := json.Unmarshal([]byte(expected), &dt2)
 		testutil.OK(t, err)
-		testutil.Equals(t, string(bs), `{"__extn":{"fn":"datetime","arg":"1970-01-01T00:00:00.042Z"}}`)
+		testutil.Equals(t, dt1, dt2)
+	})
+
+	t.Run("UnmarshalJSON/error", func(t *testing.T) {
+		t.Parallel()
+		var dt2 types.Datetime
+		err := json.Unmarshal([]byte("{}"), &dt2)
+		testutil.Error(t, err)
 	})
 }
diff --git a/types/decimal.go b/types/decimal.go
@@ -160,34 +160,19 @@ func (d Decimal) String() string {
 	return res[:right]
 }
 
+// UnmarshalJSON implements encoding/json.Unmarshaler for Decimal
+//
+// It is capable of unmarshaling 3 different representations supported by Cedar
+//   - { "__extn": { "fn": "decimal", "arg": "1234.5678" }}
+//   - { "fn": "decimal", "arg": "1234.5678" }
+//   - "1234.5678"
 func (d *Decimal) UnmarshalJSON(b []byte) error {
-	var arg string
-	if len(b) > 0 && b[0] == '"' {
-		if err := json.Unmarshal(b, &arg); err != nil {
-			return errors.Join(errJSONDecode, err)
-		}
-	} else {
-		// NOTE: cedar supports two other forms, for now we're only supporting the smallest implicit and explicit form.
-		// The following are not supported:
-		// "decimal(\"1234.5678\")"
-		// {"fn":"decimal","arg":"1234.5678"}
-		var res extValueJSON
-		if err := json.Unmarshal(b, &res); err != nil {
-			return errors.Join(errJSONDecode, err)
-		}
-		if res.Extn == nil {
-			return errJSONExtNotFound
-		}
-		if res.Extn.Fn != "decimal" {
-			return errJSONExtFnMatch
-		}
-		arg = res.Extn.Arg
-	}
-	vv, err := ParseDecimal(arg)
+	dd, err := unmarshalExtensionValue(b, "decimal", ParseDecimal)
 	if err != nil {
 		return err
 	}
-	*d = vv
+
+	*d = dd
 	return nil
 }
 

diff --git a/types/decimal_test.go b/types/decimal_test.go
@@ -1,6 +1,7 @@
 package types_test
 
 import (
+	"encoding/json"
 	"fmt"
 	"testing"
 
@@ -287,4 +288,28 @@ func TestDecimal(t *testing.T) {
 			string(testutil.Must(types.NewDecimal(42, 0)).MarshalCedar()),
 			`decimal("42.0")`)
 	})
+
+	t.Run("MarshalJSON", func(t *testing.T) {
+		t.Parallel()
+		expected := `{
+			"__extn": {
+				"fn": "decimal",
+				"arg": "1234.5678"
+			}
+		}`
+		d1 := testutil.Must(types.NewDecimal(12345678, -4))
+		testutil.JSONMarshalsTo(t, d1, expected)
+
+		var d2 types.Decimal
+		err := json.Unmarshal([]byte(expected), &d2)
+		testutil.OK(t, err)
+		testutil.Equals(t, d1, d2)
+	})
+
+	t.Run("UnmarshalJSON/error", func(t *testing.T) {
+		t.Parallel()
+		var dt2 types.Decimal
+		err := json.Unmarshal([]byte("{}"), &dt2)
+		testutil.Error(t, err)
+	})
 }
diff --git a/types/duration.go b/types/duration.go
@@ -3,7 +3,6 @@ package types
 import (
 	"bytes"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"math"
 	"strconv"
@@ -214,46 +213,16 @@ func (v Duration) String() string {
 
 // UnmarshalJSON implements encoding/json.Unmarshaler for Duration
 //
-// It is capable of unmarshaling 4 different representations supported by Cedar
-// - { "__extn": { "fn": "duration", "arg": "1h10m" }}
-// - { "fn": "duration", "arg": "1h10m" }
-// - "duration(\"1h10m\")"
-// - "1h10m"
+// It is capable of unmarshaling 3 different representations supported by Cedar
+//   - { "__extn": { "fn": "duration", "arg": "1h10m" }}
+//   - { "fn": "duration", "arg": "1h10m" }
+//   - "1h10m"
 func (v *Duration) UnmarshalJSON(b []byte) error {
-	var arg string
-	if bytes.HasPrefix(b, []byte(`"duration(\"`)) && bytes.HasSuffix(b, []byte(`\")"`)) {
-		arg = string(b[12 : len(b)-4])
-	} else if len(b) > 0 && b[0] == '"' {
-		if err := json.Unmarshal(b, &arg); err != nil {
-			return errors.Join(errJSONDecode, err)
-		}
-	} else {
-		var res extValueJSON
-		if err := json.Unmarshal(b, &res); err != nil {
-			return errors.Join(errJSONDecode, err)
-		}
-		if res.Extn == nil {
-			// If we didn't find an Extn, maybe it's just an extn.
-			var res2 extn
-			json.Unmarshal(b, &res2)
-			// We've tried Ext.Fn and Fn, so no good.
-			if res2.Fn == "" {
-				return errJSONExtNotFound
-			}
-			if res2.Fn != "duration" {
-				return errJSONExtFnMatch
-			}
-			arg = res2.Arg
-		} else if res.Extn.Fn != "duration" {
-			return errJSONExtFnMatch
-		} else {
-			arg = res.Extn.Arg
-		}
-	}
-	vv, err := ParseDuration(arg)
+	vv, err := unmarshalExtensionValue(b, "duration", ParseDuration)
 	if err != nil {
 		return err
 	}
+
 	*v = vv
 	return nil
 }

diff --git a/types/duration_test.go b/types/duration_test.go
@@ -1,6 +1,7 @@
 package types_test
 
 import (
+	"encoding/json"
 	"fmt"
 	"testing"
 	"time"
@@ -176,8 +177,25 @@ func TestDuration(t *testing.T) {
 
 	t.Run("MarshalJSON", func(t *testing.T) {
 		t.Parallel()
-		bs, err := types.NewDuration(42 * time.Millisecond).MarshalJSON()
+		expected := `{
+			"__extn": {
+				"fn": "duration",
+				"arg": "42ms"
+			}
+		}`
+		d1 := types.NewDuration(42 * time.Millisecond)
+		testutil.JSONMarshalsTo(t, d1, expected)
+
+		var d2 types.Duration
+		err := json.Unmarshal([]byte(expected), &d2)
 		testutil.OK(t, err)
-		testutil.Equals(t, string(bs), `{"__extn":{"fn":"duration","arg":"42ms"}}`)
+		testutil.Equals(t, d1, d2)
+	})
+
+	t.Run("UnmarshalJSON/error", func(t *testing.T) {
+		t.Parallel()
+		var dt2 types.Duration
+		err := json.Unmarshal([]byte("{}"), &dt2)
+		testutil.Error(t, err)
 	})
 }
diff --git a/types/ipaddr.go b/types/ipaddr.go
@@ -2,7 +2,6 @@ package types
 
 import (
 	"encoding/json"
-	"errors"
 	"fmt"
 	"hash/fnv"
 	"net/netip"
@@ -104,33 +103,18 @@ func (c IPAddr) Contains(o IPAddr) bool {
 	return c.Prefix().Contains(o.Addr()) && c.Prefix().Bits() <= o.Prefix().Bits()
 }
 
+// UnmarshalJSON implements encoding/json.Unmarshaler for IPAddr
+//
+// It is capable of unmarshaling 3 different representations supported by Cedar
+//   - { "__extn": { "fn": "ip", "arg": "12.34.56.78" }}
+//   - { "fn": "ip", "arg": "12.34.56.78" }
+//   - "12.34.56.78"
 func (v *IPAddr) UnmarshalJSON(b []byte) error {
-	var arg string
-	if len(b) > 0 && b[0] == '"' {
-		if err := json.Unmarshal(b, &arg); err != nil {
-			return errors.Join(errJSONDecode, err)
-		}
-	} else {
-		// NOTE: cedar supports two other forms, for now we're only supporting the smallest implicit explicit form.
-		// The following are not supported:
-		// "ip(\"192.168.0.42\")"
-		// {"fn":"ip","arg":"192.168.0.42"}
-		var res extValueJSON
-		if err := json.Unmarshal(b, &res); err != nil {
-			return errors.Join(errJSONDecode, err)
-		}
-		if res.Extn == nil {
-			return errJSONExtNotFound
-		}
-		if res.Extn.Fn != "ip" {
-			return errJSONExtFnMatch
-		}
-		arg = res.Extn.Arg
-	}
-	vv, err := ParseIPAddr(arg)
+	vv, err := unmarshalExtensionValue(b, "ip", ParseIPAddr)
 	if err != nil {
 		return err
 	}
+
 	*v = vv
 	return nil
 }

diff --git a/types/ipaddr_test.go b/types/ipaddr_test.go
@@ -1,6 +1,7 @@
 package types_test
 
 import (
+	"encoding/json"
 	"fmt"
 	"testing"
 
@@ -277,4 +278,27 @@ func TestIP(t *testing.T) {
 			`ip("10.0.0.42")`)
 	})
 
+	t.Run("MarshalJSON", func(t *testing.T) {
+		t.Parallel()
+		expected := `{
+			"__extn": {
+				"fn": "ip",
+				"arg": "12.34.56.78"
+			}
+		}`
+		i1 := testutil.Must(types.ParseIPAddr("12.34.56.78"))
+		testutil.JSONMarshalsTo(t, i1, expected)
+
+		var i2 types.IPAddr
+		err := json.Unmarshal([]byte(expected), &i2)
+		testutil.OK(t, err)
+		testutil.Equals(t, i1, i2)
+	})
+
+	t.Run("UnmarshalJSON/error", func(t *testing.T) {
+		t.Parallel()
+		var dt2 types.IPAddr
+		err := json.Unmarshal([]byte("{}"), &dt2)
+		testutil.Error(t, err)
+	})
 }