实现WAF

internal/utils/get.go

@@ -0,0 +1,77 @@
+package utils
+
+import (
+	"github.com/iwind/TeaGo/types"
+	"reflect"
+	"regexp"
+)
+
+var RegexpDigitNumber = regexp.MustCompile("^\\d+$")
+
+func Get(object interface{}, keys []string) interface{} {
+	if len(keys) == 0 {
+		return object
+	}
+
+	if object == nil {
+		return nil
+	}
+
+	firstKey := keys[0]
+	keys = keys[1:]
+
+	value := reflect.ValueOf(object)
+
+	if !value.IsValid() {
+		return nil
+	}
+
+	if value.Kind() == reflect.Ptr {
+		value = value.Elem()
+	}
+
+	if value.Kind() == reflect.Struct {
+		field := value.FieldByName(firstKey)
+		if !field.IsValid() {
+			return nil
+		}
+
+		if len(keys) == 0 {
+			return field.Interface()
+		}
+
+		return Get(field.Interface(), keys)
+	}
+
+	if value.Kind() == reflect.Map {
+		mapKey := reflect.ValueOf(firstKey)
+		mapValue := value.MapIndex(mapKey)
+		if !mapValue.IsValid() {
+			return nil
+		}
+
+		if len(keys) == 0 {
+			return mapValue.Interface()
+		}
+
+		return Get(mapValue.Interface(), keys)
+	}
+
+	if value.Kind() == reflect.Slice {
+		if RegexpDigitNumber.MatchString(firstKey) {
+			firstKeyInt := types.Int(firstKey)
+			if value.Len() > firstKeyInt {
+				result := value.Index(firstKeyInt).Interface()
+				if len(keys) == 0 {
+					return result
+				}
+
+				return Get(result, keys)
+			}
+		}
+
+		return nil
+	}
+
+	return nil
+}

internal/utils/get_test.go

@@ -0,0 +1,79 @@
+package utils
+
+import "testing"
+
+func TestGetStruct(t *testing.T) {
+	object := struct {
+		Name  string
+		Age   int
+		Books []string
+		Extend struct {
+			Location struct {
+				City string
+			}
+		}
+	}{
+		Name:  "lu",
+		Age:   20,
+		Books: []string{"Golang"},
+		Extend: struct {
+			Location struct {
+				City string
+			}
+		}{
+			Location: struct {
+				City string
+			}{
+				City: "Beijing",
+			},
+		},
+	}
+
+	if Get(object, []string{"Name"}) != "lu" {
+		t.Fatal("[ERROR]Name != lu")
+	}
+
+	if Get(object, []string{"Age"}) != 20 {
+		t.Fatal("[ERROR]Age != 20")
+	}
+
+	if Get(object, []string{"Books", "0"}) != "Golang" {
+		t.Fatal("[ERROR]books.0 != Golang")
+	}
+
+	t.Log("Extend.Location:", Get(object, []string{"Extend", "Location"}))
+
+	if Get(object, []string{"Extend", "Location", "City"}) != "Beijing" {
+		t.Fatal("[ERROR]Extend.Location.City != Beijing")
+	}
+}
+
+func TestGetMap(t *testing.T) {
+	object := map[string]interface{}{
+		"Name": "lu",
+		"Age":  20,
+		"Extend": map[string]interface{}{
+			"Location": map[string]interface{}{
+				"City": "Beijing",
+			},
+		},
+	}
+
+	if Get(object, []string{"Name"}) != "lu" {
+		t.Fatal("[ERROR]Name != lu")
+	}
+
+	if Get(object, []string{"Age"}) != 20 {
+		t.Fatal("[ERROR]Age != 20")
+	}
+
+	if Get(object, []string{"Books", "0"}) != nil {
+		t.Fatal("[ERROR]books.0 != nil")
+	}
+
+	t.Log(Get(object, []string{"Extend", "Location"}))
+
+	if Get(object, []string{"Extend", "Location", "City"}) != "Beijing" {
+		t.Fatal("[ERROR]Extend.Location.City != Beijing")
+	}
+}

internal/utils/string.go

@@ -0,0 +1,37 @@
+package utils
+
+import (
+	"strings"
+	"unsafe"
+)
+
+// convert bytes to string
+func UnsafeBytesToString(bs []byte) string {
+	return *(*string)(unsafe.Pointer(&bs))
+}
+
+// convert string to bytes
+func UnsafeStringToBytes(s string) []byte {
+	return *(*[]byte)(unsafe.Pointer(&s))
+}
+
+// format address
+func FormatAddress(addr string) string {
+	if strings.HasSuffix(addr, "unix:") {
+		return addr
+	}
+	addr = strings.Replace(addr, " ", "", -1)
+	addr = strings.Replace(addr, "\t", "", -1)
+	addr = strings.Replace(addr, "：", ":", -1)
+	addr = strings.TrimSpace(addr)
+	return addr
+}
+
+// format address list
+func FormatAddressList(addrList []string) []string {
+	result := []string{}
+	for _, addr := range addrList {
+		result = append(result, FormatAddress(addr))
+	}
+	return result
+}

internal/utils/string_test.go

@@ -0,0 +1,56 @@
+package utils
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestBytesToString(t *testing.T) {
+	t.Log(UnsafeBytesToString([]byte("Hello,World")))
+}
+
+func TestStringToBytes(t *testing.T) {
+	t.Log(string(UnsafeStringToBytes("Hello,World")))
+}
+
+func BenchmarkBytesToString(b *testing.B) {
+	data := []byte("Hello,World")
+	for i := 0; i < b.N; i++ {
+		_ = UnsafeBytesToString(data)
+	}
+}
+
+func BenchmarkBytesToString2(b *testing.B) {
+	data := []byte("Hello,World")
+	for i := 0; i < b.N; i++ {
+		_ = string(data)
+	}
+}
+
+func BenchmarkStringToBytes(b *testing.B) {
+	s := strings.Repeat("Hello,World", 1024)
+	for i := 0; i < b.N; i++ {
+		_ = UnsafeStringToBytes(s)
+	}
+}
+
+func BenchmarkStringToBytes2(b *testing.B) {
+	s := strings.Repeat("Hello,World", 1024)
+	for i := 0; i < b.N; i++ {
+		_ = []byte(s)
+	}
+}
+
+func TestFormatAddress(t *testing.T) {
+	t.Log(FormatAddress("127.0.0.1:1234"))
+	t.Log(FormatAddress("127.0.0.1 : 1234"))
+	t.Log(FormatAddress("127.0.0.1：1234"))
+}
+
+func TestFormatAddressList(t *testing.T) {
+	t.Log(FormatAddressList([]string{
+		"127.0.0.1:1234",
+		"127.0.0.1 : 1234",
+		"127.0.0.1：1234",
+	}))
+}