support affinity and validate key and value properly

Signed-off-by: Chanwit Kaewkasi <chanwit@gmail.com>

scheduler/filter/README.md

@@ -149,11 +149,18 @@ As you can see here, the containers were only scheduled on nodes with the redis
 
 #### Constraint Expression Syntax
 
-Additionally, you can use a not (`!`) to negate and a regular expression in the form of `/regexp/` for specifying a constraint.
+As previously mentioned, a constraint consists of a `key` and a `value`.
+A `key` must conform the alpha-numeric pattern, with the leading alphabet or underscore.
+
+A `value` must be one of the following:
+* An alpha-numeric string, dots, hyphens, and underscores.
+* A globbing pattern, i.e., `abc*`.
+* A regular expression in the form of `/regexp/`. We support the Go's regular expression syntax.
+
+You can use a not (`!`) to negate and a regular expression in the form of `/regexp/` for specifying a constraint.
 Relative comparisons, `>=` and `<=` are also supported, but they are limited to `string` comparison only.
 
 For example,
-
 * `constraint:name=node1` will match nodes named with `node1`.
 * `constraint:name==node1` will also match nodes named with `node1`. Note that `==` also allowed.
 * `constraint:name!=node1` will match all nodes, except `node1`.

scheduler/filter/affinity.go

@@ -13,7 +13,11 @@ type AffinityFilter struct {
 }
 
 func (f *AffinityFilter) Filter(config *dockerclient.ContainerConfig, nodes []*cluster.Node) ([]*cluster.Node, error) {
-	affinities := extractEnv("affinity", config.Env)
+	affinities, err := extractEnv("affinity", config.Env)
+	if err != nil {
+		return nil, err
+	}
+
 	for k, v := range affinities {
 		log.Debugf("matching affinity: %s=%s", k, v)
 		candidates := []*cluster.Node{}
@@ -21,7 +25,7 @@ func (f *AffinityFilter) Filter(config *dockerclient.ContainerConfig, nodes []*c
 			switch k {
 			case "container":
 				for _, container := range node.Containers() {
-					if match(v, container.Id) || match(v, container.Names[0]) {
+					if match(v, container.Id, false) || match(v, container.Names[0], false) {
 						candidates = append(candidates, node)
 						break
 					}
@@ -29,12 +33,12 @@ func (f *AffinityFilter) Filter(config *dockerclient.ContainerConfig, nodes []*c
 			case "image":
 			done:
 				for _, image := range node.Images() {
-					if match(v, image.Id) {
+					if match(v, image.Id, false) {
 						candidates = append(candidates, node)
 						break
 					}
 					for _, t := range image.RepoTags {
-						if match(v, t) {
+						if match(v, t, false) {
 							candidates = append(candidates, node)
 							break done
 						}

scheduler/filter/constraint.go

@@ -13,7 +13,10 @@ type ConstraintFilter struct {
 }
 
 func (f *ConstraintFilter) Filter(config *dockerclient.ContainerConfig, nodes []*cluster.Node) ([]*cluster.Node, error) {
-	constraints := extractEnv("constraint", config.Env)
+	constraints, err := extractEnv("constraint", config.Env)
+	if err != nil {
+		return nil, err
+	}
 	for k, v := range constraints {
 		log.Debugf("matching constraint: %s=%s", k, v)
 

scheduler/filter/constraint_test.go

@@ -111,6 +111,12 @@ func TestConstrainteFilter(t *testing.T) {
 	}, nodes)
 	assert.NoError(t, err)
 	assert.Len(t, result, 2)
+
+	result, err = f.Filter(&dockerclient.ContainerConfig{
+		Env: []string{"constraint:region=*us*"},
+	}, nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 2)
 }
 
 func TestConstraintNotExpr(t *testing.T) {
@@ -228,41 +234,6 @@ func TestConstraintRegExp(t *testing.T) {
 	assert.Equal(t, result[0], nodes[2])
 }
 
-func TestFilterRegExpWithEscape(t *testing.T) {
-	var (
-		f      = ConstraintFilter{}
-		nodes  = testFixtures()
-		result []*cluster.Node
-		err    error
-	)
-
-	// Prepare node with a strange name
-	node3 := cluster.NewNode("node-3", 0)
-	node3.ID = "node-3-id"
-	node3.Name = "node-3-name"
-	node3.Labels = map[string]string{
-		"name":   "foo[bar]",
-		"group":  "2",
-		"region": "eu",
-	}
-	nodes = append(nodes, node3)
-
-	// Test filter with a strange name
-	result, err = f.Filter(&dockerclient.ContainerConfig{
-		Env: []string{`constraint:name=/foo\[bar\]/`},
-	}, nodes)
-	assert.NoError(t, err)
-	assert.Len(t, result, 1)
-	assert.Equal(t, result[0], nodes[3])
-
-	// Test ! filter with a strange name
-	result, err = f.Filter(&dockerclient.ContainerConfig{
-		Env: []string{`constraint:name=!/foo\[bar\]/`},
-	}, nodes)
-	assert.NoError(t, err)
-	assert.Len(t, result, 3)
-}
-
 func TestFilterRegExpCaseInsensitive(t *testing.T) {
 	var (
 		f      = ConstraintFilter{}

scheduler/filter/utils.go

@@ -1,6 +1,7 @@
 package filter
 
 import (
+	"fmt"
 	"regexp"
 	"strings"
 
@@ -61,20 +62,45 @@ func parse(k, v string) (string, string, comparison, bool) {
 	return k, v, mode, useRegex
 }
 
-func extractEnv(key string, env []string) map[string]string {
+func extractEnv(key string, env []string) (map[string]string, error) {
 	values := make(map[string]string)
 	for _, e := range env {
 		if strings.HasPrefix(e, key+":") {
 			value := strings.TrimPrefix(e, key+":")
 			parts := strings.SplitN(value, "=", 2)
+
+			// validate key
+			// allow alpha-numeric
+			// but also allow !,>,< operators as suffix
+			matched, err := regexp.MatchString(`^(?i)[a-z_][a-z0-9\-_]+[><!]?$`, parts[0])
+
+			if err != nil {
+				return nil, err
+			}
+			if matched == false {
+				return nil, fmt.Errorf("Key '%s' is not valid", parts[0])
+			}
+
 			if len(parts) == 2 {
+
+				// validate value
+				// allow leading = in case of using ==
+				// allow * for globbing
+				// allow regexp
+				matched, err := regexp.MatchString(`^(?i)[=!\/]?[a-z0-9:\-_\.\*/\(\)\?\+\[\]\\\^\$]+$`, parts[1])
+				if err != nil {
+					return nil, err
+				}
+				if matched == false {
+					return nil, fmt.Errorf("Value '%s' is not valid", parts[1])
+				}
 				values[strings.ToLower(parts[0])] = parts[1]
 			} else {
 				values[strings.ToLower(parts[0])] = ""
 			}
 		}
 	}
-	return values
+	return values, nil
 }
 
 // Create the regex for globbing (ex: ub*t* -> ^ub.*t.*$) and match.

scheduler/filter/utils_test.go

@@ -0,0 +1,38 @@
+package filter
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestValidatingEnvExtraction(t *testing.T) {
+	err := error(nil)
+
+	// Cannot use the leading digit for key
+	_, err = extractEnv("constraint", []string{"constraint:1node"})
+	assert.Error(t, err)
+
+	// Cannot use space in key
+	_, err = extractEnv("constraint", []string{"constraint:node ==node1"})
+	assert.Error(t, err)
+
+	// Cannot use dot in key
+	_, err = extractEnv("constraint", []string{"constraint:no.de==node1"})
+	assert.Error(t, err)
+
+	// Cannot use * in key
+	_, err = extractEnv("constraint", []string{"constraint:no*de==node1"})
+	assert.Error(t, err)
+
+	// Allow leading underscore
+	_, err = extractEnv("constraint", []string{"constraint:_node==_node1"})
+	assert.NoError(t, err)
+
+	// Allow globbing
+	_, err = extractEnv("constraint", []string{"constraint:node==*node*"})
+	assert.NoError(t, err)
+
+	// Allow regexp in value
+	_, err = extractEnv("constraint", []string{"constraint:node==/(?i)^[a-b]+c*$/"})
+	assert.NoError(t, err)
+}