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.
|
// Directions supported by this package.
|
||||||
const (
|
const (
|
||||||
|
// Unknown direction is the default value and is not a valid Direction.
|
||||||
Unknown Direction = iota
|
Unknown Direction = iota
|
||||||
|
|
||||||
|
// Client initiated.
|
||||||
Client
|
Client
|
||||||
|
|
||||||
|
// Server initiated.
|
||||||
Server
|
Server
|
||||||
|
|
||||||
|
// Both is either client or server initiated.
|
||||||
Both
|
Both
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// Contains checks if the provided other direction is included in this direction.
|
||||||
func (dir Direction) Contains(other Direction) bool {
|
func (dir Direction) Contains(other Direction) bool {
|
||||||
switch dir {
|
switch dir {
|
||||||
case Client:
|
case Client:
|
||||||
@@ -49,6 +57,13 @@ var directionName = map[Direction]string{
|
|||||||
Both: "both",
|
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 {
|
func (dir Direction) String() string {
|
||||||
if s, ok := directionName[dir]; ok {
|
if s, ok := directionName[dir]; ok {
|
||||||
return s
|
return s
|
||||||
@@ -76,34 +91,26 @@ type detectResult struct {
|
|||||||
Confidence float64
|
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)
|
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) {
|
func Register(dir Direction, magic string, detect DetectFunc) {
|
||||||
|
if dir.IsValid() {
|
||||||
formatsMu.Lock()
|
formatsMu.Lock()
|
||||||
formats, _ := atomicFormats.Load().([]format)
|
formats, _ := atomicFormats.Load().([]format)
|
||||||
atomicFormats.Store(append(formats, format{dir, magic, detect}))
|
atomicFormats.Store(append(formats, format{dir, magic, detect}))
|
||||||
formatsMu.Unlock()
|
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.
|
// 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 {
|
for _, format := range formats {
|
||||||
if format.dir.Contains(dir) {
|
if format.dir.Contains(dir) {
|
||||||
// Check the buffer to see if we have sufficient bytes
|
// 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 {
|
if proto, confidence := format.detect(dir, data, srcPort, dstPort); proto != nil {
|
||||||
results = append(results, detectResult{proto, confidence})
|
results = append(results, detectResult{proto, confidence})
|
||||||
}
|
}
|
||||||
|
@@ -1,6 +1,6 @@
|
|||||||
package protocol
|
package protocol
|
||||||
|
|
||||||
// MatchPattern checks if the byte slice matches the magic string pattern.
|
// Match the input against the magic string pattern.
|
||||||
//
|
//
|
||||||
// '?' matches any single character
|
// '?' matches any single character
|
||||||
// '*' matches zero or more characters
|
// '*' matches zero or more characters
|
||||||
|
Reference in New Issue
Block a user