diff --git a/proto/proto.go b/proto/proto.go index 325bbc6..5c905d4 100644 --- a/proto/proto.go +++ b/proto/proto.go @@ -66,22 +66,57 @@ func (t *AccessKey) ValidateService(service *Service) bool { // MarshalJSON adds freeCU, softQuota and hardQuota fields to the json. // This keeps compatibility with older versions of QuotaControl. // TODO: remove this once all services have migrated to newer version. -func (l *Limit) MarshalJSON() ([]byte, error) { +func (l Limit) MarshalJSON() ([]byte, error) { + // Alias is used to avoid infinite recursion when marshaling Limit. type Alias Limit + var v = struct { - *Alias - CreditsIncluded int64 `json:"freeCU"` - SoftQuota int64 `json:"softQuota"` - HardQuota int64 `json:"hardQuota"` + Alias + FreeCU int64 `json:"freeCU"` + SoftQuota int64 `json:"softQuota"` + HardQuota int64 `json:"hardQuota"` }{ - Alias: (*Alias)(l), - CreditsIncluded: l.FreeWarn, - SoftQuota: l.OverWarn, - HardQuota: l.OverMax, + Alias: (Alias)(l), + FreeCU: l.FreeWarn, + SoftQuota: l.OverWarn, + HardQuota: l.OverMax, } return json.Marshal(v) } +// UnmarshalJSON parses the json and checks if the older version fields are set +// and the newer ones are empty. If so, it uses the older fields to populate the new ones. +func (l *Limit) UnmarshalJSON(b []byte) error { + // Alias is used to avoid infinite recursion when marshaling Limit. + type Alias Limit + + v := struct { + Alias + FreeCU int64 `json:"freeCU"` + SoftQuota int64 `json:"softQuota"` + HardQuota int64 `json:"hardQuota"` + }{} + if err := json.Unmarshal(b, &v); err != nil { + return err + } + *l = Limit(v.Alias) + if v.FreeCU != 0 { + if l.FreeWarn == 0 { + l.FreeWarn = v.FreeCU + } + if l.FreeMax == 0 { + l.FreeMax = v.FreeCU + } + } + if v.SoftQuota != 0 && l.OverWarn == 0 { + l.OverWarn = v.SoftQuota + } + if v.HardQuota != 0 && l.OverMax == 0 { + l.OverMax = v.HardQuota + } + return nil +} + func (l Limit) Validate() error { if l.RateLimit < 1 { return fmt.Errorf("rateLimit must be > 0") diff --git a/proto/proto_test.go b/proto/proto_test.go index 0d4d803..56536e8 100644 --- a/proto/proto_test.go +++ b/proto/proto_test.go @@ -240,4 +240,12 @@ func TestLimitJSON(t *testing.T) { err = json.Unmarshal(rawJson, &legacyLimit) require.NoError(t, err) assert.Equal(t, legacyExpected, legacyLimit) + + legacyJson, err := json.Marshal(&legacyExpected) + require.NoError(t, err) + + limit = proto.Limit{} + err = json.Unmarshal(legacyJson, &limit) + require.NoError(t, err) + assert.Equal(t, expected, limit) } diff --git a/quotacontrol_test.go b/quotacontrol_test.go index 302124d..edcb6f3 100644 --- a/quotacontrol_test.go +++ b/quotacontrol_test.go @@ -298,7 +298,7 @@ func TestDefaultKey(t *testing.T) { limit := proto.Limit{ RateLimit: 100, - FreeWarn: 5, + FreeMax: 5, OverWarn: 7, OverMax: 10, }