add more unit tests and deserializer/serializer to regex

5 files changed, 124 insertions, 54 deletions

diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index 1c75205..7c6bce5 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -2,7 +2,6 @@ package clients
 
 import (
 	"context"
-	"regexp"
 	"sync"
 	"time"
 
@@ -10,6 +9,7 @@ import (
 	"github.com/mimecast/dtail/internal/discovery"
 	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/regex"
 	"github.com/mimecast/dtail/internal/ssh/client"
 
 	gossh "golang.org/x/crypto/ssh"
@@ -34,6 +34,8 @@ type baseClient struct {
 	retry bool
 	// Connection maker helper.
 	maker maker
+	// Regex is the regular expresion object for line filtering
+	Regex regex.Regex
 }
 
 func (c *baseClient) init(maker maker) {
@@ -48,9 +50,11 @@ func (c *baseClient) init(maker maker) {
 		c.connections = append(c.connections, c.makeConnection(server, c.sshAuthMethods, c.hostKeyCallback))
 	}
 
-	if _, err := regexp.Compile(c.Regex); err != nil {
-		logger.FatalExit(c.Regex, "Can't test compile regex", err)
+	regex, err := regex.New(c.Args.Regex, regex.Default)
+	if err != nil {
+		logger.FatalExit(c.Regex, "invalid regex!", err, regex)
 	}
+	c.Regex = regex
 
 	c.stats = newTailStats(len(c.connections))
 }
diff --git a/internal/regex/regex.go b/internal/regex/regex.go
index 35685a2..800bc31 100644
--- a/internal/regex/regex.go
+++ b/internal/regex/regex.go
@@ -1,69 +1,101 @@
 package regex
 
 import (
+	"fmt"
 	"regexp"
-)
-
-type Flag int
+	"strings"
 
-const (
-    // Default is the default regex mode (positive matching)
-    Default = iota
-    // Negative negates the regex 
-    Negative = iota
-    // Noop means no regex matching enabled, all defaults to true
-    Noop = iota
+	"github.com/mimecast/dtail/internal/io/logger"
 )
 
 type Regex struct {
-    str string
-    re *regexp.Regexp
-    flag Flag
+	// The original regex string
+	regexStr string
+	// The Golang regexp object
+	re *regexp.Regexp
+	// For now only use the first flag at flags[0], but in the future we can
+	// set and use multiple flags.
+	flags []Flag
 }
 
-func (r Regex) Noop() Regex {
-    return Regex{
-        flag: Noop,
-    }
+func (r Regex) String() string {
+	return fmt.Sprintf("Regex(regexStr:%s,flags:%s,re==nil:%t)", r.regexStr, r.flags, r.re == nil)
 }
 
-func New(str string, flag Flag) (Regex, error) {
-    r := Regex{
-        str:  str,
-        flag: flag,
-    }
+func NewNoop() Regex {
+	return Regex{flags: []Flag{Noop}}
+}
 
-	re, err := regexp.Compile(str)
+func New(regexStr string, flag Flag) (Regex, error) {
+	return new(regexStr, []Flag{flag})
+}
+
+func new(regexStr string, flags []Flag) (Regex, error) {
+	r := Regex{
+		regexStr: regexStr,
+		flags:    flags,
+	}
+
+	re, err := regexp.Compile(regexStr)
 
 	if err != nil {
-	    return r, err
+		return r, err
 	}
 
 	r.re = re
-    return r, nil
+	return r, nil
 }
 
 func (r Regex) MatchString(str string) bool {
-    switch r.flag {
-    case Default:
-        return r.re.MatchString(str)
-    case Negative:
-        return !r.re.MatchString(str)
-    case Noop:
-        return true
-    default:
-        return false
-    }
+	switch r.flags[0] {
+	case Default:
+		return r.re.MatchString(str)
+	case Negate:
+		return !r.re.MatchString(str)
+	case Noop:
+		return true
+	default:
+		return false
+	}
 }
 
-/*
 func (r Regex) Serialize() string {
-    // TODO: Serialize to hex encoded str
-    return fmt.Sprintf("%b,%s",r.negate,r,str)
+	var flags []string
+	for _, flag := range r.flags {
+		flags = append(flags, flag.String())
+	}
+
+	return fmt.Sprintf("regex:%s %s", strings.Join(flags, ","), r.regexStr)
 }
 
-func (r Regex) Deserialize(input string) (Regex, error) {
+func Deserialize(str string) (Regex, error) {
+	// Get regex string
+	s := strings.SplitN(str, " ", 2)
+	if len(s) < 2 {
+		return Regex{}, fmt.Errorf("unable to deserialize regex '%s'", str)
+	}
+
+	flagsStr := s[0]
+	regexStr := s[1]
 
-}
+	if !strings.HasPrefix(flagsStr, "regex") {
+		return Regex{}, fmt.Errorf("unable to deserialize regex '%s': should start with string 'regex'", str)
+	}
+
+	// Parse regex flags, e.g. "regex:flag1,flag2,flag3..."
+	var flags []Flag
+	if strings.Contains(flagsStr, ":") {
+		s := strings.SplitN(flagsStr, ":", 2)
+		for _, flagStr := range strings.Split(s[1], ",") {
+			flag, err := NewFlag(flagStr)
+			if err != nil {
+				logger.Error("ignoring flag", err)
+				continue
+			}
+			logger.Debug("Adding regex flag", flag)
+			flags = append(flags, flag)
+		}
+	}
 
-*/
+	return new(regexStr, flags)
+}
diff --git a/internal/regex/regex_test.go b/internal/regex/regex_test.go
index a90aee2..acdd518 100644
--- a/internal/regex/regex_test.go
+++ b/internal/regex/regex_test.go
@@ -1,16 +1,48 @@
 package regex
 
-import "testing"
+import (
+	"testing"
+)
 
 func TestRegex(t *testing.T) {
+	input := "hello"
 
-    input := "hello"
-    r, err := New(".hello", Default)
-    if err != nil {
-		t.Errorf("error: unable to create regex: %v\n", err)
-    }
+	r := NewNoop()
+	if !r.MatchString(input) {
+		t.Errorf("expected to match string '%s' with noop regex '%v' but didn't\n", input, r)
+	}
 
-    if r.MatchString(input) {
-		t.Errorf("error: expected to match string '%s' with regex '%v' but didnt\n", input, r)
-    }
+	r, err := New(".hello", Default)
+	if err != nil {
+		t.Errorf("unable to create regex: %v\n", err)
+	}
+	if r.MatchString(input) {
+		t.Errorf("expected to match string '%s' with regex '%v' but didn't\n", input, r)
+	}
+
+	r2, err := Deserialize(r.Serialize())
+	if err != nil {
+		t.Errorf("unable to serialize deserialized regex: %v: %v\n", r.Serialize(), err)
+	}
+	if r.String() != r2.String() {
+		t.Errorf("regex should be the same after deserialize(serialize(..)), got '%s' but expected '%s'.\n",
+			r2.String(), r.String())
+	}
+
+	r, err = New(".hello", Negate)
+	if err != nil {
+		t.Errorf("unable to create regex: %v\n", err)
+	}
+	if !r.MatchString(input) {
+		t.Errorf("expected to not match string '%s' with regex '%v' but matched\n", input, r)
+	}
+
+	r2, err = Deserialize(r.Serialize())
+	if err != nil {
+		t.Errorf("unable to serialize deserialized regex: %v: %v\n", r.Serialize(), err)
+	}
+	if r.String() != r2.String() {
+		t.Errorf("regex should be the same after deserialize(serialize(..)), got '%s' but expected '%s'.\n",
+			r2.String(), r.String())
+	}
 }
diff --git a/internal/server/handlers/readcommand.go b/internal/server/handlers/readcommand.go
index 63158f6..d4c871c 100644
--- a/internal/server/handlers/readcommand.go
+++ b/internal/server/handlers/readcommand.go
@@ -27,6 +27,8 @@ func newReadCommand(server *ServerHandler, mode omode.Mode) *readCommand {
 func (r *readCommand) Start(ctx context.Context, argc int, args []string) {
 	regex := "."
 	if argc >= 4 {
+		switch args[2] {
+		}
 		regex = strings.Join(args[3:], " ")
 		logger.Debug("Joined regex", regex)
 	}
diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index 9b52c85..34d2e30 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -200,7 +200,7 @@ func (h *ServerHandler) handleProtocolVersion(args []string) ([]string, int, err
 }
 
 func (h *ServerHandler) handleBase64(args []string, argc int) ([]string, int, error) {
-	err := errors.New("Unable to decode client message, DTail server and client versions not compatible")
+	err := errors.New("Unable to decode client message, DTail server and client versions may not be compatible")
 
 	if argc != 2 || args[0] != "base64" {
 		return args, argc, err