Switch to our new Match function

Also improved some docstrings
This commit is contained in:
2025-10-09 12:39:49 +02:00
parent bc35ff63e7
commit 31774b961d
2 changed files with 33 additions and 26 deletions

View File

@@ -24,12 +24,20 @@ type Direction int
// Directions supported by this package.
const (
// Unknown direction is the default value and is not a valid Direction.
Unknown Direction = iota
// Client initiated.
Client
// Server initiated.
Server
// Both is either client or server initiated.
Both
)
// Contains checks if the provided other direction is included in this direction.
func (dir Direction) Contains(other Direction) bool {
switch dir {
case Client:
@@ -49,6 +57,13 @@ var directionName = map[Direction]string{
Both: "both",
}
// IsValid checks if dir has a value recognized by this library.
//
// Also Unknown direction is not considered valid.
func (dir Direction) IsValid() bool {
return dir > Unknown && dir <= Both
}
func (dir Direction) String() string {
if s, ok := directionName[dir]; ok {
return s
@@ -76,34 +91,26 @@ type detectResult struct {
Confidence float64
}
// DetectFunc is a function which runs the in-depth protcol detection logic.
//
// The confidence score should be between 0 and 0.99. Score boundaries are not hard enforced by this library.
type DetectFunc func(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol, confidence float64)
// Register a new protocol detector.
//
// The direction indicates in what direction we'll inspect the []byte slice passed to the [Detect]
// function. Passing an invalid direction or None will be discarded.
//
// The magic string is used to quickly analyze if the []byte slice passed to [Detect] qualifies
// for further inspection by the [DetectFunc]. See the [Match] function documentation for how
// magic strings are matched against the input.
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
if dir.IsValid() {
formatsMu.Lock()
formats, _ := atomicFormats.Load().([]format)
atomicFormats.Store(append(formats, format{dir, magic, detect}))
formatsMu.Unlock()
}
// 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.
@@ -115,7 +122,7 @@ func Detect(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol,
for _, format := range formats {
if format.dir.Contains(dir) {
// Check the buffer to see if we have sufficient bytes
if matchMagic(format.magic, data) {
if Match(format.magic, data) {
if proto, confidence := format.detect(dir, data, srcPort, dstPort); proto != nil {
results = append(results, detectResult{proto, confidence})
}