Merge pull request #414 from linhvo/364-argument-validation

364 argument validation

db.go

@@ -91,6 +91,12 @@ func (db *DB) SetColumnLabel(v string) error {
 		return nil
 	}
 
+	// Make sure columnLabel is valid name
+	err := ValidateName(v)
+	if err != nil {
+		return err
+	}
+
 	// Persist meta data to disk on change.
 	db.columnLabel = v
 	if err := db.saveMeta(); err != nil {

frame.go

@@ -115,6 +115,13 @@ func (f *Frame) SetRowLabel(v string) error {
 		return nil
 	}
 
+
+	// Make sure rowLabel is valid name
+	err := ValidateName(v)
+	if err != nil {
+		return err
+	}
+
 	// Persist meta data to disk on change.
 	f.rowLabel = v
 	if err := f.saveMeta(); err != nil {

handle_internal_test.go

@@ -0,0 +1,76 @@
+package pilosa
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+)
+
+// Test custom UnmarshalJSON for postDBRequest object
+func TestPostDBRequestUnmarshalJSON(t *testing.T) {
+	tests := []struct {
+		json     string
+		expected postDBRequest
+		err      string
+	}{
+		{json: `{"db": "d", "options": {}}`, expected: postDBRequest{DB: "d", Options: DBOptions{}}},
+		{json: `{"db": 1, "options": {}}`, err: "db required and must be a string"},
+		{json: `{"db": "d", "options": 4}`, err: "options is not map[string]interface{}"},
+		{json: `{"db": "d", "option": {}}`, err: "Unknown key: option:map[]"},
+		{json: `{"db": "d", "options": {"columnLabel": "test"}}`, expected: postDBRequest{DB: "d", Options: DBOptions{ColumnLabel: "test"}}},
+		{json: `{"db": "d", "options": {"columnLabl": "test"}}`, err: "invalid key for options {columnLabl:test}"},
+	}
+	for _, test := range tests {
+		actual := &postDBRequest{}
+		err := json.Unmarshal([]byte(test.json), actual)
+
+		if err != nil {
+			if test.err == "" || test.err != err.Error() {
+				t.Errorf("expected error: %v, but got result: %v", test.err, err)
+			}
+		} else {
+			if test.err != "" {
+				t.Errorf("expected error: %v, but got no error", test.err)
+			}
+		}
+
+		if test.err == "" {
+			if !reflect.DeepEqual(*actual, test.expected) {
+				t.Errorf("expected: %v, but got: %v", test.expected, *actual)
+			}
+		}
+
+	}
+}
+
+// Test custom UnmarshalJSON for postFrameRequest object
+func TestPostFrameRequestUnmarshalJSON(t *testing.T) {
+	tests := []struct {
+		json     string
+		expected postFrameRequest
+		err      string
+	}{
+		{json: `{"db": "d", "frame":"f",  "options": {}}`, expected: postFrameRequest{DB: "d", Frame: "f", Options: FrameOptions{}}},
+		{json: `{"db": "d", "options": {}}`, err: "frame required and must be a string"},
+		{json: `{"db": "d", "frame":"f", "options": 4}`, err: "options is not map[string]interface{}"},
+		{json: `{"db": "d", "frame":"f", "option": {}}`, err: "Unknown key: {option:map[]}"},
+		{json: `{"db": "d", "frame":"f", "options": {"rowLabel": "test"}}`, expected: postFrameRequest{DB: "d", Frame: "f", Options: FrameOptions{RowLabel: "test"}}},
+		{json: `{"db": "d", "frame":"f", "options": {"rowLabl": "test"}}`, err: "invalid key for options {rowLabl:test}"},
+	}
+	for _, test := range tests {
+		actual := &postFrameRequest{}
+		err := json.Unmarshal([]byte(test.json), actual)
+		if err != nil {
+			if test.err == "" || test.err != err.Error() {
+				t.Errorf("expected error: %v, but got result: %v", test.err, err)
+			}
+		}
+
+		if test.err == "" {
+			if !reflect.DeepEqual(*actual, test.expected) {
+				t.Errorf("expected: %v, but got: %v", test.expected, *actual)
+			}
+		}
+
+	}
+}

handler.go

@@ -344,6 +344,63 @@ func (h *Handler) handlePostDB(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// Custom Unmarshal JSON to validate request body when creating a new database
+func (p *postDBRequest) UnmarshalJSON(b []byte) error {
+	var data map[string]interface{}
+	if err := json.Unmarshal(b, &data); err != nil {
+		return err
+	}
+	for key, value := range data {
+		switch key {
+		case "db":
+			val, ok := data["db"].(string)
+			if !ok {
+				return errors.New("db required and must be a string")
+			}
+			p.DB = val
+		case "options":
+			value, err := validateOptions(data, "columnLabel")
+			if err != nil {
+				return err
+			}
+			if value == "" {
+				p.Options = DBOptions{}
+			} else {
+				p.Options = DBOptions{ColumnLabel: value}
+			}
+
+		default:
+			return fmt.Errorf("Unknown key: %v:%v", key, value)
+		}
+	}
+	return nil
+}
+
+func validateOptions(data map[string]interface{}, field string) (string, error) {
+	options, ok := data["options"].(map[string]interface{})
+	if !ok {
+		return "", errors.New("options is not map[string]interface{}")
+	}
+	var optionValue string
+	if len(options) == 0 {
+		optionValue = ""
+	} else {
+		for k, v := range options {
+			switch k {
+			case field:
+				val, ok := options[field].(string)
+				if !ok {
+					return "", fmt.Errorf("invalid option %v: {%v:%v}", field, k, v)
+				}
+				optionValue = val
+			default:
+				return "", fmt.Errorf("invalid key for options {%v:%v}", k, v)
+			}
+		}
+	}
+	return optionValue, nil
+}
+
 type postDBRequest struct {
 	DB      string    `json:"db"`
 	Options DBOptions `json:"options"`
@@ -478,6 +535,7 @@ type postDBAttrDiffResponse struct {
 
 // handlePostFrame handles POST /frame request.
 func (h *Handler) handlePostFrame(w http.ResponseWriter, r *http.Request) {
+
 	// Decode request.
 	var req postFrameRequest
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
@@ -508,6 +566,46 @@ func (h *Handler) handlePostFrame(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// Custom Unmarshal JSON to validate request body when creating a new frame
+func (p *postFrameRequest) UnmarshalJSON(b []byte) error {
+	var data map[string]interface{}
+	if err := json.Unmarshal(b, &data); err != nil {
+		return err
+	}
+	for key, value := range data {
+		switch key {
+		case "db":
+			val, ok := data["db"].(string)
+			if !ok {
+				return errors.New("db required and must be a string")
+			}
+			p.DB = val
+		case "frame":
+			val, ok := data["frame"].(string)
+			if !ok {
+				return errors.New("frame required and must be a string")
+			}
+			p.Frame = val
+
+		case "options":
+			value, err := validateOptions(data, "rowLabel")
+			if err != nil {
+				return err
+			}
+			if value == "" {
+				p.Options = FrameOptions{}
+			} else {
+				p.Options = FrameOptions{RowLabel: value}
+			}
+
+		default:
+			return fmt.Errorf("Unknown key: {%v:%v}", key, value)
+		}
+	}
+	return nil
+
+}
+
 type postFrameRequest struct {
 	DB      string       `json:"db"`
 	Frame   string       `json:"frame"`