Merge pull request #31 from coreruleset/feat/var-enums

feat: add enums for variables and collections

Author: fzipi

types/collection_test.go

@@ -0,0 +1,85 @@
+package types
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.yaml.in/yaml/v4"
+)
+
+var (
+	colTests = []struct {
+		collection CollectionName
+		yamlStr    string
+	}{
+		{ARGS, "ARGS"},
+		{ARGS_GET, "ARGS_GET"},
+		{ARGS_GET_NAMES, "ARGS_GET_NAMES"},
+		{ARGS_NAMES, "ARGS_NAMES"},
+		{ARGS_POST_NAMES, "ARGS_POST_NAMES"},
+		{ARGS_POST, "ARGS_POST"},
+		{ENV, "ENV"},
+		{FILES, "FILES"},
+		{GEO, "GEO"},
+		{GLOBAL, "GLOBAL"},
+		{IP, "IP"},
+		{MATCHED_VARS_NAMES, "MATCHED_VARS_NAMES"},
+		{MATCHED_VARS, "MATCHED_VARS"},
+		{MULTIPART_PART_HEADERS, "MULTIPART_PART_HEADERS"},
+		{PERF_RULES, "PERF_RULES"},
+		{REQUEST_COOKIES_NAMES, "REQUEST_COOKIES_NAMES"},
+		{REQUEST_COOKIES, "REQUEST_COOKIES"},
+		{REQUEST_HEADERS_NAMES, "REQUEST_HEADERS_NAMES"},
+		{REQUEST_HEADERS, "REQUEST_HEADERS"},
+		{RESPONSE_HEADERS_NAMES, "RESPONSE_HEADERS_NAMES"},
+		{RESPONSE_HEADERS, "RESPONSE_HEADERS"},
+		{RULE, "RULE"},
+		{SESSION, "SESSION"},
+		{TX, "TX"},
+		{XML, "XML"},
+	}
+)
+
+func TestCollectionNameToString(t *testing.T) {
+	for _, tt := range colTests {
+		t.Run(tt.yamlStr, func(t *testing.T) {
+			if tt.collection.String() != tt.yamlStr {
+				t.Errorf("Expected %q, got %q", tt.yamlStr, tt.collection.String())
+			}
+		})
+	}
+}
+
+func TestStringToCollectionName(t *testing.T) {
+	for _, tt := range colTests {
+		t.Run(tt.yamlStr, func(t *testing.T) {
+			collection := stringToCollectionName(tt.yamlStr)
+			if collection != tt.collection {
+				t.Errorf("Expected %q, got %q", tt.collection, collection)
+			}
+		})
+	}
+}
+
+func TestMarshalCollectionName(t *testing.T) {
+	for _, tt := range colTests {
+		t.Run(tt.yamlStr, func(t *testing.T) {
+			data, err := yaml.Marshal(tt.collection)
+			if err != nil {
+				t.Fatalf("Failed to marshal: %v", err)
+			}
+			if string(data) != tt.yamlStr+"\n" {
+				t.Errorf("Expected %q, got %q", tt.yamlStr+"\n", data)
+			}
+		})
+	}
+}
+
+func TestUnknownCollectionName(t *testing.T) {
+	t.Run("marshal unknown", func(t *testing.T) {
+		unknown := UNKNOWN_COLLECTION
+		_, err := unknown.MarshalYAML()
+		assert.Error(t, err)
+		assert.Equal(t, "Unknown collection name", err.Error())
+	})
+}

types/collections.go

@@ -9,35 +9,36 @@ type Collection struct {
 	Count     bool           `yaml:"count,omitempty"`
 }
 
-type CollectionName string
+type CollectionName int
 
 const (
 	// Collections
-	ARGS                   CollectionName = "ARGS"
-	ARGS_GET               CollectionName = "ARGS_GET"
-	ARGS_GET_NAMES         CollectionName = "ARGS_GET_NAMES"
-	ARGS_NAMES             CollectionName = "ARGS_NAMES"
-	ARGS_POST_NAMES        CollectionName = "ARGS_POST_NAMES"
-	ARGS_POST              CollectionName = "ARGS_POST"
-	ENV                    CollectionName = "ENV"
-	FILES                  CollectionName = "FILES"
-	GEO                    CollectionName = "GEO"
-	GLOBAL                 CollectionName = "GLOBAL"
-	IP                     CollectionName = "IP"
-	MATCHED_VARS_NAMES     CollectionName = "MATCHED_VARS_NAMES"
-	MATCHED_VARS           CollectionName = "MATCHED_VARS"
-	MULTIPART_PART_HEADERS CollectionName = "MULTIPART_PART_HEADERS"
-	PERF_RULES             CollectionName = "PERF_RULES"
-	REQUEST_COOKIES_NAMES  CollectionName = "REQUEST_COOKIES_NAMES"
-	REQUEST_COOKIES        CollectionName = "REQUEST_COOKIES"
-	REQUEST_HEADERS_NAMES  CollectionName = "REQUEST_HEADERS_NAMES"
-	REQUEST_HEADERS        CollectionName = "REQUEST_HEADERS"
-	RESPONSE_HEADERS_NAMES CollectionName = "RESPONSE_HEADERS_NAMES"
-	RESPONSE_HEADERS       CollectionName = "RESPONSE_HEADERS"
-	RULE                   CollectionName = "RULE"
-	SESSION                CollectionName = "SESSION"
-	TX                     CollectionName = "TX"
-	XML                    CollectionName = "XML"
+	UNKNOWN_COLLECTION CollectionName = iota
+	ARGS
+	ARGS_GET
+	ARGS_GET_NAMES
+	ARGS_NAMES
+	ARGS_POST_NAMES
+	ARGS_POST
+	ENV
+	FILES
+	GEO
+	GLOBAL
+	IP
+	MATCHED_VARS_NAMES
+	MATCHED_VARS
+	MULTIPART_PART_HEADERS
+	PERF_RULES
+	REQUEST_COOKIES_NAMES
+	REQUEST_COOKIES
+	REQUEST_HEADERS_NAMES
+	REQUEST_HEADERS
+	RESPONSE_HEADERS_NAMES
+	RESPONSE_HEADERS
+	RULE
+	SESSION
+	TX
+	XML
 )
 
 var (
@@ -70,26 +71,102 @@ var (
 	}
 )
 
+func (c CollectionName) String() string {
+	switch c {
+	case ARGS:
+		return "ARGS"
+	case ARGS_GET:
+		return "ARGS_GET"
+	case ARGS_GET_NAMES:
+		return "ARGS_GET_NAMES"
+	case ARGS_NAMES:
+		return "ARGS_NAMES"
+	case ARGS_POST_NAMES:
+		return "ARGS_POST_NAMES"
+	case ARGS_POST:
+		return "ARGS_POST"
+	case ENV:
+		return "ENV"
+	case FILES:
+		return "FILES"
+	case GEO:
+		return "GEO"
+	case GLOBAL:
+		return "GLOBAL"
+	case IP:
+		return "IP"
+	case MATCHED_VARS_NAMES:
+		return "MATCHED_VARS_NAMES"
+	case MATCHED_VARS:
+		return "MATCHED_VARS"
+	case MULTIPART_PART_HEADERS:
+		return "MULTIPART_PART_HEADERS"
+	case PERF_RULES:
+		return "PERF_RULES"
+	case REQUEST_COOKIES_NAMES:
+		return "REQUEST_COOKIES_NAMES"
+	case REQUEST_COOKIES:
+		return "REQUEST_COOKIES"
+	case REQUEST_HEADERS_NAMES:
+		return "REQUEST_HEADERS_NAMES"
+	case REQUEST_HEADERS:
+		return "REQUEST_HEADERS"
+	case RESPONSE_HEADERS_NAMES:
+		return "RESPONSE_HEADERS_NAMES"
+	case RESPONSE_HEADERS:
+		return "RESPONSE_HEADERS"
+	case RULE:
+		return "RULE"
+	case SESSION:
+		return "SESSION"
+	case TX:
+		return "TX"
+	case XML:
+		return "XML"
+	default:
+		return "unknown"
+	}
+}
+
+func (c CollectionName) MarshalYAML() (interface{}, error) {
+	if c == UNKNOWN_COLLECTION {
+		return nil, fmt.Errorf("Unknown collection name")
+	}
+	return c.String(), nil
+}
+
+func (c *CollectionName) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var name string
+	if err := unmarshal(&name); err != nil {
+		return err
+	}
+	*c = stringToCollectionName(name)
+	if *c == UNKNOWN_COLLECTION {
+		return fmt.Errorf("Collection name %s is not valid", name)
+	}
+	return nil
+}
+
 func CollectionsToString(collections []Collection, separator string) string {
 	result := ""
 	for i, collection := range collections {
 		if len(collection.Arguments) == 0 && len(collection.Excluded) == 0 {
 			if collection.Count {
 				result += "&"
 			}
-			result += string(collection.Name)
+			result += collection.Name.String()
 		} else {
 			for j, arg := range collection.Arguments {
 				if collection.Count {
 					result += "&"
 				}
-				result += string(collection.Name) + ":" + arg
+				result += collection.Name.String() + ":" + arg
 				if j != len(collection.Arguments)-1 || len(collection.Excluded) > 0 {
 					result += separator
 				}
 			}
 			for j, excluded := range collection.Excluded {
-				result += "!" + string(collection.Name) + ":" + excluded
+				result += "!" + collection.Name.String() + ":" + excluded
 				if j != len(collection.Excluded)-1 {
 					result += separator
 				}
@@ -102,10 +179,10 @@ func CollectionsToString(collections []Collection, separator string) string {
 	return result
 }
 
-func GetCollection(name string) (CollectionName, error) {
+func stringToCollectionName(name string) CollectionName {
 	col, exists := allCollections[name]
 	if !exists {
-		return "", fmt.Errorf("Invalid collection name: %s", name)
+		return UNKNOWN_COLLECTION
 	}
-	return col, nil
+	return col
 }

types/secrule.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"fmt"
 	"slices"
 )
 
@@ -34,9 +35,9 @@ func (d SecRule) GetTransformations() Transformations {
 }
 
 func (s *SecRule) AddVariable(name string, excluded bool) error {
-	variable, err := GetVariable(name)
-	if err != nil {
-		return err
+	variable := stringToVariableName(name)
+	if variable == UNKNOWN_VAR {
+		return fmt.Errorf("Invalid variable name: %s", name)
 	}
 	if excluded {
 		vars := []Variable{}
@@ -53,9 +54,9 @@ func (s *SecRule) AddVariable(name string, excluded bool) error {
 }
 
 func (s *SecRule) AddCollection(name, value string, excluded, asCount bool) error {
-	col, err := GetCollection(name)
-	if err != nil {
-		return err
+	col := stringToCollectionName(name)
+	if col == UNKNOWN_COLLECTION {
+		return fmt.Errorf("Invalid collection name: %s", name)
 	}
 	if excluded && !asCount {
 		results := []Collection{}

types/update_target.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"fmt"
 	"strconv"
 )
 
@@ -22,9 +23,9 @@ func NewUpdateTargetDirective() *UpdateTargetDirective {
 }
 
 func (d *UpdateTargetDirective) AddVariable(name string, excluded bool) error {
-	variable, err := GetVariable(name)
-	if err != nil {
-		return err
+	variable := stringToVariableName(name)
+	if variable == UNKNOWN_VAR {
+		return fmt.Errorf("Invalid variable name: %s", name)
 	}
 	if excluded {
 		vars := []Variable{}
@@ -41,9 +42,9 @@ func (d *UpdateTargetDirective) AddVariable(name string, excluded bool) error {
 }
 
 func (d *UpdateTargetDirective) AddCollection(name, value string, excluded, asCount bool) error {
-	col, err := GetCollection(name)
-	if err != nil {
-		return err
+	col := stringToCollectionName(name)
+	if col == UNKNOWN_COLLECTION {
+		return fmt.Errorf("Invalid collection name: %s", name)
 	}
 	if excluded && !asCount {
 		results := []Collection{}