feat: register custom sanitizers (#16)

* feat: option for custom sanitizers

* docs: update docs with custom sanitizer option

* refactor: code review recommendations

Author: shawalli

README.md

@@ -82,6 +82,64 @@ s := sanitizer.New(sanitizer.OptionDateFormat{
 })
 ```
 
+### Custom Sanitizers
+
+Use this option to register a custom sanitizer function. The sanitizer function is responsible for determining if the field's type is supported for that sanitizer.
+
+The `Name` field tells us what tag name corresponds to the sanitizer.
+
+The `Sanitizer` field tells us which sanitizer function to call when this tag is used. All sanitizers must have this signature: `func(s Sanitizer, structValue reflect.Value, idx int) error`.
+
+```go
+// exclaim adds punctuated enthusiasm to a string.
+func exclaim(s Sanitizer, structValue reflect.Value, idx int) error {
+    fieldValue := structValue.Field(idx)
+    if fieldValue.Kind() == reflect.Ptr && !fieldValue.IsNil() {
+        fieldValue = fieldValue.Elem()
+    }
+
+    if fieldValue.Kind() != reflect.String {
+        return fmt.Errorf("exclaim: invalid type %q", fieldValue.Kind().String())
+    }
+
+    v := fieldValue.Interface().(string)
+    if strings.HasSuffix(v, "...") {
+        v = v[:len(v)-3] + "!"
+    } else if strings.HasSuffix(v, ".") {
+        v = v[:len(v)-1] + "!"
+    } else {
+        v += "!"
+    }
+
+    fieldValue.SetString(v)
+
+    return nil
+}
+
+type BlogPost struct {
+    Title    string `san:"title"`
+    Byline   string `san:"cap,exclaim"`
+    Contents []byte
+}
+
+func main() {
+    bp := BlogPost{
+        Title:  "sanitizing go structs"
+        Byline: "clean up you structs"
+    }
+
+    s := sanitizer.New(
+        sanitizer.OptionSanitizerFunc{
+            Name:      "exclaim",
+            Sanitizer: exclaim,
+        },
+    )
+    s.Sanitize(&bp)
+
+    fmt.Printf("Title:%q Byline:%q", bp.Title, bp.Byline)
+    // Title: "Sanitizing Go Structs" Byline:"Clean up your structs!"
+}
+```
 
 ## Available tags
 

options.go

@@ -43,3 +43,21 @@ func (o OptionDateFormat) id() string {
 func (o OptionDateFormat) value() interface{} {
 	return o
 }
+
+// OptionSanitizerFunc allows users to use custom sanitizer functions
+type OptionSanitizerFunc struct {
+	Name      string
+	Sanitizer SanitizerFunc
+}
+
+var _ Option = OptionSanitizerFunc{}
+
+const optionSanitizerFuncID = "sanitizer-func"
+
+func (o OptionSanitizerFunc) id() string {
+	return optionSanitizerFuncID
+}
+
+func (o OptionSanitizerFunc) value() interface{} {
+	return o.Sanitizer
+}

options_test.go

@@ -3,6 +3,7 @@ package sanitize
 import (
 	"reflect"
 	"testing"
+	"unsafe"
 )
 
 type unknownOption struct{}
@@ -17,6 +18,72 @@ func (o unknownOption) value() interface{} {
 	return "very strange indeed"
 }
 
+// sanitizersEqual checks for equality between two sanitizers in a manner similar to reflect.DeepEqual.
+//
+// We have to check equality manually due to sanitizer functions.
+//
+//	From reflect docs: "Func values are deeply equal if both are nil; otherwise they are not deeply equal."
+func sanitizersEqual(s *Sanitizer, o *Sanitizer) bool {
+	if s == nil && o == nil {
+		return true
+	} else if (s != nil && o == nil) || (s == nil && o != nil) {
+		return false
+	}
+
+	if !reflect.DeepEqual(s.tagName, o.tagName) {
+		return false
+	}
+
+	if !reflect.DeepEqual(s.dateInput, o.dateInput) {
+		return false
+	}
+
+	if !reflect.DeepEqual(s.dateKeepFormat, o.dateKeepFormat) {
+		return false
+	}
+
+	if !reflect.DeepEqual(s.dateOutput, o.dateOutput) {
+		return false
+	}
+
+	if s.sanitizersByName == nil && o.sanitizersByName == nil {
+		return true
+	} else if (s.sanitizersByName != nil && o.sanitizersByName == nil) || (s.sanitizersByName == nil && o.sanitizersByName != nil) {
+		return false
+	}
+
+	var sKeys []string
+	for k, _ := range s.sanitizersByName {
+		sKeys = append(sKeys, k)
+	}
+
+	var oKeys []string
+	for k, _ := range o.sanitizersByName {
+		oKeys = append(oKeys, k)
+	}
+
+	if !reflect.DeepEqual(sKeys, oKeys) {
+		return false
+	}
+
+	for _, k := range sKeys {
+		psf := s.sanitizersByName[k]
+		osf := o.sanitizersByName[k]
+		upsf := *(*unsafe.Pointer)(unsafe.Pointer(&psf))
+		upof := *(*unsafe.Pointer)(unsafe.Pointer(&osf))
+
+		if !reflect.DeepEqual(upsf, upof) {
+			return false
+		}
+	}
+
+	return true
+}
+
+func doNothing(s Sanitizer, structValue reflect.Value, idx int) error {
+	return nil
+}
+
 func Test_New(t *testing.T) {
 	type args struct {
 		options []Option
@@ -89,6 +156,32 @@ func Test_New(t *testing.T) {
 			want:    nil,
 			wantErr: true,
 		},
+		{
+			name: "valid sanitizer func option",
+			args: args{
+				options: []Option{
+					OptionSanitizerFunc{Name: "capfirst", Sanitizer: capFirst},
+				},
+			},
+			want: &Sanitizer{
+				tagName: DefaultTagName,
+				sanitizersByName: map[string]SanitizerFunc{
+					"capfirst": capFirst,
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "duplicate sanitizer func option",
+			args: args{
+				options: []Option{
+					OptionSanitizerFunc{Name: "capfirst", Sanitizer: capFirst},
+					OptionSanitizerFunc{Name: "capfirst", Sanitizer: doNothing},
+				},
+			},
+			want:    nil,
+			wantErr: true,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -97,7 +190,8 @@ func Test_New(t *testing.T) {
 				t.Errorf("New() error = %v, wantErr %v", err, tt.wantErr)
 				return
 			}
-			if !reflect.DeepEqual(got, tt.want) {
+
+			if !sanitizersEqual(got, tt.want) {
 				t.Errorf("New() = %v, want %v", got, tt.want)
 			}
 		})

sanitize.go

@@ -7,16 +7,20 @@ import (
 	"reflect"
 )
 
-// DefaultTagName intance is the name of the tag that must be present on the string
+// DefaultTagName instance is the name of the tag that must be present on the string
 // fields of the structs to be sanitized. Defaults to "san".
 const DefaultTagName = "san"
 
-// Sanitizer intance
+type SanitizerFunc func(s Sanitizer, structValue reflect.Value, idx int) error
+
+// Sanitizer instance
 type Sanitizer struct {
 	tagName        string
 	dateInput      []string
 	dateKeepFormat bool
 	dateOutput     string
+
+	sanitizersByName map[string]SanitizerFunc
 }
 
 // New sanitizer instance
@@ -37,6 +41,16 @@ func New(options ...Option) (*Sanitizer, error) {
 			s.dateInput = v.Input
 			s.dateKeepFormat = v.KeepFormat
 			s.dateOutput = v.Output
+		case optionSanitizerFuncID:
+			if s.sanitizersByName == nil {
+				s.sanitizersByName = make(map[string]SanitizerFunc)
+			}
+
+			v := o.(OptionSanitizerFunc)
+			if _, ok := s.sanitizersByName[v.Name]; ok {
+				return nil, fmt.Errorf("sanitizer already registered with name %q", o.id())
+			}
+			s.sanitizersByName[v.Name] = o.value().(SanitizerFunc)
 		default:
 			return nil, fmt.Errorf("option %q is not valid", o.id())
 		}
@@ -161,6 +175,19 @@ func (s Sanitizer) sanitizeRec(v reflect.Value) error {
 		field := v.Field(i)
 		fkind := field.Kind()
 
+		// Prioritize custom sanitizers; tag's value can be re-resolved inside the sanitizer
+		tags := s.fieldTags(v.Type().Field(i).Tag)
+		for tag := range tags {
+			sanitizerFunc, ok := s.sanitizersByName[tag]
+			if !ok {
+				continue
+			}
+
+			if err := sanitizerFunc(s, v, i); err != nil {
+				return err
+			}
+		}
+
 		// If the field is a slice, sanitize it first
 		isPtrToSlice := fkind == reflect.Ptr && field.Elem().Kind() == reflect.Slice
 		isSlice := fkind == reflect.Slice

sanitize_test.go

@@ -1,11 +1,35 @@
 package sanitize
 
 import (
+	"fmt"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 )
 
+func capFirst(s Sanitizer, structValue reflect.Value, idx int) error {
+	fieldValue := structValue.Field(idx)
+
+	if fieldValue.Kind() == reflect.Ptr && !fieldValue.IsNil() {
+		fieldValue = fieldValue.Elem()
+	}
+
+	if fieldValue.Kind() != reflect.String {
+		return fmt.Errorf("capfirst: invalid type %q", fieldValue.Kind().String())
+	}
+
+	v := fieldValue.Interface().(string)
+
+	first := string(v[0])
+	first = strings.ToUpper(first)
+	v = first + v[1:]
+
+	fieldValue.SetString(v)
+
+	return nil
+}
+
 func Test_Sanitize_CodeSample(t *testing.T) {
 	type Dog struct {
 		Name  string  `san:"max=5,trim,lower"`
@@ -50,6 +74,7 @@ func Test_Sanitize_Options(t *testing.T) {
 		Name            string `abcde:"max=5,trim,lower"`
 		Birthday        string `abcde:"date"`
 		PersonalWebsite string `abcde:"xss,trim"`
+		Breed           string `abcde:"capfirst"`
 	}
 
 	now := time.Now()
@@ -58,12 +83,14 @@ func Test_Sanitize_Options(t *testing.T) {
 		Name:            "Borky Borkins",
 		Birthday:        now.Format(time.RFC3339),
 		PersonalWebsite: "<html>[head]1=1?;{/head}(/html)",
+		Breed:           "frenchBulldog",
 	}
 
 	expected := Dog{
 		Name:            "borky",
 		Birthday:        now.Format(time.RFC850),
 		PersonalWebsite: "html head 1 1 /head /html",
+		Breed:           "FrenchBulldog",
 	}
 
 	s, _ := New(
@@ -72,8 +99,49 @@ func Test_Sanitize_Options(t *testing.T) {
 			Input:  []string{time.RFC3339},
 			Output: time.RFC850,
 		},
+		OptionSanitizerFunc{
+			Name:      "capfirst",
+			Sanitizer: capFirst,
+		},
 	)
-	s.Sanitize(&d)
+
+	if err := s.Sanitize(&d); err != nil {
+		t.Fatalf("Sanitize() - got unexpected error %v", err)
+	}
+
+	if !reflect.DeepEqual(d, expected) {
+		t.Errorf("Sanitize() - got %+v but wanted %+v", d, expected)
+	}
+}
+
+func Test_Sanitize_InvalidCustomSanitizerType(t *testing.T) {
+	type Dog struct {
+		Name    string `san:"max=5,trim,lower"`
+		Adopted bool   `san:"capfirst"`
+	}
+
+	d := Dog{
+		Name:    "Borky Borkins",
+		Adopted: true,
+	}
+
+	expected := Dog{
+		Name:    "borky",
+		Adopted: true,
+	}
+
+	s, _ := New(
+		OptionSanitizerFunc{
+			Name:      "capfirst",
+			Sanitizer: capFirst,
+		},
+	)
+
+	if err := s.Sanitize(&d); err == nil {
+		t.Fatal("Sanitize() - did not receive expected error")
+	} else if !strings.Contains(err.Error(), "capfirst: invalid type") {
+		t.Fatalf("Sanitize() - received unexpected error %v", err)
+	}
 
 	if !reflect.DeepEqual(d, expected) {
 		t.Errorf("Sanitize() - got %+v but wanted %+v", d, expected)