Initial import

protocol/detect.go

@@ -0,0 +1,113 @@
+package protocol
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+	"sync/atomic"
+)
+
+// Strict mode requires a full, compliant packet to be captured. This is only
+// implemented by some detectors.
+var Strict bool
+
+// Common errors.
+var (
+	ErrTimeout = errors.New("timeout")
+	ErrUnknown = errors.New("unknown protocol")
+)
+
+// Direction indicates the communcation direction.
+type Direction int
+
+// Directions supported by this package.
+const (
+	Unknown Direction = iota
+	Client
+	Server
+	Both
+)
+
+func (dir Direction) Contains(other Direction) bool {
+	switch dir {
+	case Client:
+		return other == Client || other == Both
+	case Server:
+		return other == Server || other == Both
+	case Both:
+		return other == Client || other == Server
+	default:
+		return false
+	}
+}
+
+var directionName = map[Direction]string{
+	Client: "client",
+	Server: "server",
+	Both:   "both",
+}
+
+func (dir Direction) String() string {
+	if s, ok := directionName[dir]; ok {
+		return s
+	}
+	return fmt.Sprintf("invalid (%d)", int(dir))
+}
+
+type format struct {
+	dir    Direction
+	magic  string
+	detect DetectFunc
+}
+
+// Formats is the list of registered formats.
+var (
+	formatsMu     sync.Mutex
+	atomicFormats atomic.Value
+)
+
+type DetectFunc func(Direction, []byte) *Protocol
+
+func Register(dir Direction, magic string, detect DetectFunc) {
+	formatsMu.Lock()
+	formats, _ := atomicFormats.Load().([]format)
+	atomicFormats.Store(append(formats, format{dir, magic, detect}))
+	formatsMu.Unlock()
+}
+
+func matchMagic(magic string, data []byte) bool {
+	// Empty magic means the detector will always run.
+	if len(magic) == 0 {
+		return true
+	}
+
+	// The buffer should contain at least the same number of bytes
+	// as our magic.
+	if len(data) < len(magic) {
+		return false
+	}
+
+	// Match bytes in magic with bytes in data.
+	for i, b := range []byte(magic) {
+		if b != '?' && data[i] != b {
+			return false
+		}
+	}
+	return true
+}
+
+// Detect a protocol based on the provided data.
+func Detect(dir Direction, data []byte) (*Protocol, error) {
+	formats, _ := atomicFormats.Load().([]format)
+	for _, f := range formats {
+		if f.dir.Contains(dir) {
+			// Check the buffer to see if we have sufficient bytes
+			if matchMagic(f.magic, data) {
+				if p := f.detect(dir, data); p != nil {
+					return p, nil
+				}
+			}
+		}
+	}
+	return nil, ErrUnknown
+}