Switch to our new Match function
Also improved some docstrings
This commit is contained in:
@@ -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})
|
||||
}
|
||||
|
Reference in New Issue
Block a user