diff --git a/limit_config_test.go b/limit_config_test.go index bdbf435..1330d4b 100644 --- a/limit_config_test.go +++ b/limit_config_test.go @@ -1,6 +1,8 @@ package rcmgr import ( + "bytes" + "encoding/json" "os" "testing" @@ -50,4 +52,11 @@ func TestLimitConfigParser(t *testing.T) { require.NoError(t, err) require.Contains(t, cfg.Peer, peerID) require.Equal(t, int64(4097), cfg.Peer[peerID].Memory) + + // Roundtrip + jsonBytes, err := json.Marshal(&cfg) + require.NoError(t, err) + cfgAfterRoundTrip, err := readLimiterConfigFromJSON(bytes.NewReader(jsonBytes), defaults) + require.NoError(t, err) + require.Equal(t, cfg, cfgAfterRoundTrip) } diff --git a/limit_defaults.go b/limit_defaults.go index a6ced25..bb10b7c 100644 --- a/limit_defaults.go +++ b/limit_defaults.go @@ -1,6 +1,7 @@ package rcmgr import ( + "encoding/json" "math" "github.com/libp2p/go-libp2p-core/peer" @@ -136,6 +137,23 @@ type LimitConfig struct { Stream BaseLimit `json:",omitempty"` } +func (cfg *LimitConfig) MarshalJSON() ([]byte, error) { + // we want to marshal the encoded peer id + encodedPeerMap := make(map[string]BaseLimit, len(cfg.Peer)) + for p, v := range cfg.Peer { + encodedPeerMap[peer.Encode(p)] = v + } + + type Alias LimitConfig + return json.Marshal(&struct { + *Alias + Peer map[string]BaseLimit `json:",omitempty"` + }{ + Alias: (*Alias)(cfg), + Peer: encodedPeerMap, + }) +} + func (cfg *LimitConfig) Apply(c LimitConfig) { cfg.System.Apply(c.System) cfg.Transient.Apply(c.Transient)