From 3fabccaf711ea80040a1c4eaa3d31241a5cecc48 Mon Sep 17 00:00:00 2001
From: Qiu Jian <qiujian@yunionyun.com>
Date: Thu, 18 Oct 2018 17:52:08 +0800
Subject: [PATCH] parser improvement

---
 pkg/util/conditionparser/parser.go      | 34 +++++++++++++++++++------
 pkg/util/conditionparser/parser_test.go |  7 +++++
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/pkg/util/conditionparser/parser.go b/pkg/util/conditionparser/parser.go
index f8a7dba747..4a4e6b3596 100644
--- a/pkg/util/conditionparser/parser.go
+++ b/pkg/util/conditionparser/parser.go
@@ -20,6 +20,8 @@ var (
 	ErrInvalidOp     = errors.New("invalid operation")
 	ErrFieldNotFound = errors.New("field not found")
 	ErrOutOfIndex    = errors.New("out of index")
+	ErrFuncArgument  = errors.New("invalid function arguments")
+	ErrMethodNotFound = errors.New("method not found")
 )
 
 func IsValid(exprStr string) bool {
@@ -110,23 +112,22 @@ func evalIndex(expr *ast.IndexExpr, input interface{}) (interface{}, error) {
 	if err != nil {
 		return nil, err
 	}
-	var indexI int64
-	switch indexV.(type) {
-	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
-		indexI = reflect.ValueOf(indexV).Int()
-	default:
-		return nil, ErrInvalidOp
-	}
 	switch X.(type) {
 	case []interface{}, *jsonutils.JSONArray:
 		arrX := getArray(X)
+		indexI := getInt(indexV)
 		if indexI >= 0 && indexI < int64(len(arrX)) {
 			return arrX[indexI], nil
 		} else {
 			return nil, ErrOutOfIndex
 		}
+	case *jsonutils.JSONDict:
+		jsonX := X.(*jsonutils.JSONDict)
+		field := getString(indexV)
+		return getJSONProperty(jsonX, field)
 	case string:
 		strX := X.(string)
+		indexI := getInt(indexV)
 		if indexI >= 0 && indexI < int64(len(strX)) {
 			return strX[indexI], nil
 		} else {
@@ -192,6 +193,18 @@ func evalCallInternal(funcV interface{}, args []interface{}) (interface{}, error
 			default:
 				return nil, ErrInvalidOp
 			}
+		case *jsonutils.JSONDict:
+			jsonX := caller.caller.(*jsonutils.JSONDict)
+			switch caller.method {
+			case "contains":
+				strArgs, err := args2Strings(args)
+				if err != nil {
+					return nil, err
+				}
+				return jsonX.Contains(strArgs...), nil
+			default:
+				return nil, ErrInvalidOp
+			}
 		default:
 			return nil, ErrInvalidOp
 		}
@@ -277,7 +290,12 @@ func evalSelectorInternal(X interface{}, identStr string) (interface{}, error) {
 	switch X.(type) {
 	case *jsonutils.JSONDict:
 		json := X.(*jsonutils.JSONDict)
-		return getJSONProperty(json, identStr)
+		ret, err := getJSONProperty(json, identStr)
+		if err == ErrFieldNotFound {
+			return &funcCaller{caller: X, method: identStr}, nil
+		} else {
+			return ret, err
+		}
 	case []interface{}, *jsonutils.JSONArray:
 		arrX := getArray(X)
 		ret := make([]interface{}, len(arrX))
diff --git a/pkg/util/conditionparser/parser_test.go b/pkg/util/conditionparser/parser_test.go
index fac612e176..12fa185531 100644
--- a/pkg/util/conditionparser/parser_test.go
+++ b/pkg/util/conditionparser/parser_test.go
@@ -25,6 +25,10 @@ func TestAst(t *testing.T) {
 		{`server.disks[0].medium_type == "ssd"`, true},
 		{`server.disks[0].medium_type == "hdd"`, false},
 		{`server.os_type == "windows" && server.disks[0].medium_type == "ssd"`, true},
+		{`server.contains("os_type")`, true},
+		{`server.disks[0].contains("medium_type")`, true},
+		{`server.disk[0].contains("medium_type")`, true},
+		{`server.disks[0].contains("backend")`, false},
 	}
 
 	for _, c := range cases {
@@ -76,6 +80,7 @@ func TestEval2(t *testing.T) {
 		want bool
 	}{
 		{`server.os_type == "Linux"`, true},
+		{`server["os_type"] == "Linux"`, true},
 		{`server.vmem_size > 2048`, false},
 		{`server.hypervisor.in("kvm", "aliyun")`, true},
 		{`server.disable_delete`, false},
@@ -130,6 +135,8 @@ func TestEval3(t *testing.T) {
 		{`server.disk.medium == "hdd"`, true},
 		{`server.disk.medium == "hybrid"`, false},
 		{`server.disk.medium.contains("ssd")`, true},
+		{`server.disk[0]["medium"] == "hdd"`, true},
+		{`server["disk"][0]["medium"] == "hdd"`, true},
 	}
 
 	for _, c := range cases {