common/helpers: make subnetmap work with struct as values

The way it was converted from a mapstruct made it not possible to have
struct as values. Fix that by checking if keys look like IP or not.
This commit is contained in:
Vincent Bernat
2022-08-14 00:54:33 +02:00
parent 3f1dc38052
commit cd535c915d
2 changed files with 95 additions and 5 deletions

View File

@@ -7,6 +7,7 @@ import (
"fmt" "fmt"
"net" "net"
"reflect" "reflect"
"regexp"
"strings" "strings"
"github.com/kentik/patricia" "github.com/kentik/patricia"
@@ -77,6 +78,8 @@ func MustNewSubnetMap[V any](from map[string]V) *SubnetMap[V] {
return trie return trie
} }
var subnetLookAlikeRegex = regexp.MustCompile("^([a-fA-F:.0-9]+)(/([0-9]+))?$")
// SubnetMapUnmarshallerHook decodes SubnetMap and notably check that // SubnetMapUnmarshallerHook decodes SubnetMap and notably check that
// valid networks are provided as key. It also accepts a single value // valid networks are provided as key. It also accepts a single value
// instead of a map for backward compatibility. // instead of a map for backward compatibility.
@@ -87,7 +90,23 @@ func SubnetMapUnmarshallerHook[V any]() mapstructure.DecodeHookFunc {
} }
output := map[string]interface{}{} output := map[string]interface{}{}
var zero V var zero V
var plausibleSubnetMap bool
if from.Kind() == reflect.Map { if from.Kind() == reflect.Map {
// When we have a map, we check if all keys look like a subnet.
plausibleSubnetMap = true
for _, key := range from.MapKeys() {
key = ElemOrIdentity(key)
if key.Kind() != reflect.String {
plausibleSubnetMap = false
break
}
if !subnetLookAlikeRegex.MatchString(key.String()) {
plausibleSubnetMap = false
break
}
}
}
if plausibleSubnetMap {
// First case, we have a map // First case, we have a map
iter := from.MapRange() iter := from.MapRange()
for i := 0; iter.Next(); i++ { for i := 0; iter.Next(); i++ {
@@ -130,11 +149,9 @@ func SubnetMapUnmarshallerHook[V any]() mapstructure.DecodeHookFunc {
} }
output[key] = v.Interface() output[key] = v.Interface()
} }
} else if from.Type() == reflect.TypeOf(zero) || from.Type().ConvertibleTo(reflect.TypeOf(zero)) {
// Second case, we have a single value
output["::/0"] = from.Interface()
} else { } else {
return from.Interface(), nil // Second case, we have a single value and we let mapstructure handles it
output["::/0"] = from.Interface()
} }
// We have to decode output map, then turn it into a SubnetMap[V] // We have to decode output map, then turn it into a SubnetMap[V]
@@ -164,5 +181,5 @@ func (sm SubnetMap[V]) MarshalYAML() (interface{}, error) {
func (sm SubnetMap[V]) String() string { func (sm SubnetMap[V]) String() string {
out := sm.ToMap() out := sm.ToMap()
return fmt.Sprintf("%v", out) return fmt.Sprintf("%+v", out)
} }

View File

@@ -97,6 +97,10 @@ func TestSubnetMapUnmarshalHook(t *testing.T) {
Description: "Invalid IP", Description: "Invalid IP",
Input: gin.H{"200.33.300.1": "customer"}, Input: gin.H{"200.33.300.1": "customer"},
Error: true, Error: true,
}, {
Description: "Random key",
Input: gin.H{"kfgdjgkfj": "customer"},
Error: true,
}, { }, {
Description: "Single value", Description: "Single value",
Input: "customer", Input: "customer",
@@ -158,3 +162,72 @@ func TestSubnetMapUnmarshalHook(t *testing.T) {
}) })
} }
} }
func TestSubnetMapUnmarshalHookWithMapValue(t *testing.T) {
type SomeStruct struct {
Blip string
Blop string
}
cases := []struct {
Description string
Input gin.H
Expected gin.H
}{
{
Description: "single value",
Input: gin.H{
"blip": "some",
"blop": "thing",
},
Expected: gin.H{
"::/0": gin.H{
"Blip": "some",
"Blop": "thing",
},
},
}, {
Description: "proper map",
Input: gin.H{
"::/0": gin.H{
"blip": "some",
"blop": "thing",
},
"203.0.113.14": gin.H{
"blip": "other",
"blop": "stuff",
},
},
Expected: gin.H{
"::/0": gin.H{
"Blip": "some",
"Blop": "thing",
},
"203.0.113.14/32": gin.H{
"Blip": "other",
"Blop": "stuff",
},
},
},
}
for _, tc := range cases {
t.Run(tc.Description, func(t *testing.T) {
var tree helpers.SubnetMap[SomeStruct]
decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
Result: &tree,
ErrorUnused: true,
Metadata: nil,
DecodeHook: helpers.SubnetMapUnmarshallerHook[SomeStruct](),
})
if err != nil {
t.Fatalf("NewDecoder() error:\n%+v", err)
}
err = decoder.Decode(tc.Input)
if err != nil {
t.Fatalf("Decode() error:\n%+v", err)
}
if diff := helpers.Diff(tree.ToMap(), tc.Expected); diff != "" {
t.Fatalf("Decode() (-got, +want):\n%s", diff)
}
})
}
}