diff --git a/go.mod b/go.mod index 58b03c3..c9cf5b8 100644 --- a/go.mod +++ b/go.mod @@ -56,6 +56,7 @@ require ( github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-isatty v0.0.19 // indirect github.com/mattn/go-runewidth v0.0.14 // indirect + github.com/metalim/jsonmap v0.4.1 // indirect github.com/mholt/archiver/v3 v3.5.1 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect diff --git a/go.sum b/go.sum index ec274df..8795a75 100644 --- a/go.sum +++ b/go.sum @@ -133,6 +133,8 @@ github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRCM46jaSJTDAk= github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWVwUuU= github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/metalim/jsonmap v0.4.1 h1:kvj9Q5oj+xne2APsoe34LmJGTUlrY3D9WYw2zwbVuVM= +github.com/metalim/jsonmap v0.4.1/go.mod h1:Rlps8z72TXjyqKPAE7pttAsBfhiZ99FLn0qzxvT4jDs= github.com/mholt/archiver/v3 v3.5.1 h1:rDjOBX9JSF5BvoJGvjqK479aL70qh9DIpZCl+k7Clwo= github.com/mholt/archiver/v3 v3.5.1/go.mod h1:e3dqJ7H78uzsRSEACH1joayhuSyhnonssnDhppzS1L4= github.com/microcosm-cc/bluemonday v1.0.21/go.mod h1:ytNkv4RrDrLJ2pqlsSI46O6IVXmZOBBD4SaJyDwwTkM= diff --git a/maps/ordered_map.go b/maps/ordered_map.go index fe12b54..e4f29e8 100644 --- a/maps/ordered_map.go +++ b/maps/ordered_map.go @@ -1,11 +1,25 @@ package mapsutil import ( + "bytes" + "encoding/json" + "fmt" + "reflect" + + "github.com/projectdiscovery/utils/conversion" sliceutil "github.com/projectdiscovery/utils/slice" + "github.com/tidwall/gjson" "golang.org/x/exp/maps" ) +var ( + _ json.Marshaler = &OrderedMap[string, struct{}]{} + _ json.Unmarshaler = &OrderedMap[string, struct{}]{} +) + // OrderedMap is a map that preserves the order of elements +// Note: Order is only guaranteed for current level of OrderedMap +// nested values only have order preserved if they are also OrderedMap type OrderedMap[k comparable, v any] struct { keys []k m map[k]v @@ -84,6 +98,85 @@ func (o *OrderedMap[k, v]) Len() int { return len(o.keys) } +// MarshalJSON marshals the OrderedMap to JSON +func (o OrderedMap[k, v]) MarshalJSON() ([]byte, error) { + var buf bytes.Buffer + buf.WriteByte('{') + for i, key := range o.keys { + if i > 0 { + buf.WriteByte(',') + } + // marshal key + keyBin, err := json.Marshal(key) + if err != nil { + return nil, fmt.Errorf("marshal key: %w", err) + } + if len(keyBin) > 0 && keyBin[len(keyBin)-1] != '"' { + buf.WriteByte('"') + buf.Write(keyBin) + buf.WriteByte('"') + } else { + buf.Write(keyBin) + } + buf.WriteByte(':') + // marshal value + valueBin, err := json.Marshal(o.m[key]) + if err != nil { + return nil, fmt.Errorf("marshal value: %w", err) + } + buf.Write(valueBin) + } + buf.WriteByte('}') + return buf.Bytes(), nil +} + +type tempStruct[k comparable] struct { + Key k +} + +// UnmarshalJSON unmarshals the OrderedMap from JSON +func (o *OrderedMap[k, v]) UnmarshalJSON(data []byte) error { + // init + o.m = map[k]v{} + + // we are only concerned about current level of ordered map + // nested ordered maps are not supported or need to be supported + // via recursive use of OrderedMap + err := json.Unmarshal(data, &o.m) + if err != nil { + return err + } + + // get type of k + var tmpKey k + keyKind := reflect.TypeOf(tmpKey).Kind() + + o.keys = []k{} + // gjson is memory efficient and faster than encoding/json + // so it shouldn't have any performance impact ( might consume some cpu though ) + result := gjson.Parse(conversion.String(data)) + result.ForEach(func(key, value gjson.Result) bool { + if keyKind == reflect.Interface { + // heterogeneous keys use any and assign + o.keys = append(o.keys, any(key.Value()).(k)) + return true + } + if keyKind == reflect.String { + o.keys = append(o.keys, any(key.String()).(k)) + return true + } + // if not use tmpStruct to unmarshal + var temp tempStruct[k] + err = json.Unmarshal([]byte(`{"key":`+key.String()+`}`), &temp) + if err != nil { + return false + } + o.keys = append(o.keys, temp.Key) + return true + }) + return nil +} + // NewOrderedMap creates a new OrderedMap func NewOrderedMap[k comparable, v any]() OrderedMap[k, v] { return OrderedMap[k, v]{ diff --git a/maps/ordered_map_test.go b/maps/ordered_map_test.go index a5a8094..5f54c5b 100644 --- a/maps/ordered_map_test.go +++ b/maps/ordered_map_test.go @@ -1,7 +1,9 @@ package mapsutil import ( + "encoding/json" "fmt" + "reflect" "strconv" "testing" ) @@ -70,3 +72,72 @@ func TestOrderedMap(t *testing.T) { } } + +func TestOrderedMapMarshalUnmarshal(t *testing.T) { + t.Run("TestSimpleStringToStringMapping", func(t *testing.T) { + orderedMap1 := NewOrderedMap[string, string]() + orderedMap1.Set("name", "John Doe") + orderedMap1.Set("occupation", "Software Developer") + + marshaled1, err := json.Marshal(orderedMap1) + if err != nil { + t.Fatalf("Failed to marshal orderedMap1: %v", err) + } + + unmarshaled1 := NewOrderedMap[string, string]() + err = json.Unmarshal(marshaled1, &unmarshaled1) + if err != nil { + t.Fatalf("Failed to unmarshal orderedMap1: %v", err) + } + + if !reflect.DeepEqual(orderedMap1, unmarshaled1) { + t.Fatal("Unmarshaled map is not equal to the original map for orderedMap1") + } + }) + + t.Run("TestIntegerToStructMapping", func(t *testing.T) { + type Employee struct { + ID int `json:"id"` + Name string `json:"name"` + } + orderedMap2 := NewOrderedMap[int, Employee]() + orderedMap2.Set(1, Employee{ID: 1, Name: "Alice"}) + orderedMap2.Set(2, Employee{ID: 2, Name: "Bob"}) + + marshaled2, err := json.Marshal(orderedMap2) + if err != nil { + t.Fatalf("Failed to marshal orderedMap2: %v", err) + } + + unmarshaled2 := NewOrderedMap[int, Employee]() + err = json.Unmarshal(marshaled2, &unmarshaled2) + if err != nil { + t.Fatalf("Failed to unmarshal orderedMap2: %v", err) + } + + if !reflect.DeepEqual(orderedMap2, unmarshaled2) { + t.Fatal("Unmarshaled map is not equal to the original map for orderedMap2") + } + }) + + t.Run("TestStringToSliceOfStringsMapping", func(t *testing.T) { + orderedMap3 := NewOrderedMap[string, []string]() + orderedMap3.Set("fruits", []string{"apple", "banana", "cherry"}) + orderedMap3.Set("vegetables", []string{"tomato", "potato", "carrot"}) + + marshaled3, err := json.Marshal(orderedMap3) + if err != nil { + t.Fatalf("Failed to marshal orderedMap3: %v", err) + } + + unmarshaled3 := NewOrderedMap[string, []string]() + err = json.Unmarshal(marshaled3, &unmarshaled3) + if err != nil { + t.Fatalf("Failed to unmarshal orderedMap3: %v", err) + } + + if !reflect.DeepEqual(orderedMap3, unmarshaled3) { + t.Fatal("Unmarshaled map is not equal to the original map for orderedMap3") + } + }) +}