114 lines
2.2 KiB
Go
114 lines
2.2 KiB
Go
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
|
|
}
|