refactor: engine/tcp/worker perf improvements
This commit is contained in:
+106
-90
@@ -10,20 +10,13 @@ import (
|
||||
"github.com/bwmarrin/snowflake"
|
||||
"github.com/google/gopacket"
|
||||
"github.com/google/gopacket/layers"
|
||||
"github.com/google/gopacket/reassembly"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultChanSize = 64
|
||||
defaultTCPMaxBufferedPagesTotal = 4096
|
||||
defaultTCPMaxBufferedPagesPerConnection = 64
|
||||
defaultUDPMaxStreams = 4096
|
||||
)
|
||||
var _ Engine = (*engine)(nil)
|
||||
|
||||
type workerPacket struct {
|
||||
StreamID uint32
|
||||
Data []byte
|
||||
LayerType gopacket.LayerType
|
||||
SrcMAC net.HardwareAddr
|
||||
DstMAC net.HardwareAddr
|
||||
SetVerdict func(io.Verdict, []byte) error
|
||||
@@ -35,12 +28,8 @@ type worker struct {
|
||||
logger Logger
|
||||
macResolver *sourceMACResolver
|
||||
|
||||
tcpStreamFactory *tcpStreamFactory
|
||||
tcpStreamPool *reassembly.StreamPool
|
||||
tcpAssembler *reassembly.Assembler
|
||||
|
||||
udpStreamFactory *udpStreamFactory
|
||||
udpStreamManager *udpStreamManager
|
||||
tcpFlowMgr *tcpFlowManager
|
||||
udpSM *udpStreamManager
|
||||
|
||||
modSerializeBuffer gopacket.SerializeBuffer
|
||||
}
|
||||
@@ -51,23 +40,17 @@ type workerConfig struct {
|
||||
Logger Logger
|
||||
Ruleset ruleset.Ruleset
|
||||
MACResolver *sourceMACResolver
|
||||
TCPMaxBufferedPagesTotal int
|
||||
TCPMaxBufferedPagesPerConn int
|
||||
TCPMaxBufferedPagesTotal int // unused, kept for config compat
|
||||
TCPMaxBufferedPagesPerConn int // unused, kept for config compat
|
||||
UDPMaxStreams int
|
||||
}
|
||||
|
||||
func (c *workerConfig) fillDefaults() {
|
||||
if c.ChanSize <= 0 {
|
||||
c.ChanSize = defaultChanSize
|
||||
}
|
||||
if c.TCPMaxBufferedPagesTotal <= 0 {
|
||||
c.TCPMaxBufferedPagesTotal = defaultTCPMaxBufferedPagesTotal
|
||||
}
|
||||
if c.TCPMaxBufferedPagesPerConn <= 0 {
|
||||
c.TCPMaxBufferedPagesPerConn = defaultTCPMaxBufferedPagesPerConnection
|
||||
c.ChanSize = 64
|
||||
}
|
||||
if c.UDPMaxStreams <= 0 {
|
||||
c.UDPMaxStreams = defaultUDPMaxStreams
|
||||
c.UDPMaxStreams = 4096
|
||||
}
|
||||
}
|
||||
|
||||
@@ -77,16 +60,12 @@ func newWorker(config workerConfig) (*worker, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tcpSF := &tcpStreamFactory{
|
||||
WorkerID: config.ID,
|
||||
Logger: config.Logger,
|
||||
Node: sfNode,
|
||||
Ruleset: config.Ruleset,
|
||||
|
||||
tcpMgr := newTCPFlowManager(config.ID, config.Logger, config.MACResolver, sfNode)
|
||||
if config.Ruleset != nil {
|
||||
tcpMgr.updateRuleset(config.Ruleset, 0)
|
||||
}
|
||||
tcpStreamPool := reassembly.NewStreamPool(tcpSF)
|
||||
tcpAssembler := reassembly.NewAssembler(tcpStreamPool)
|
||||
tcpAssembler.MaxBufferedPagesTotal = config.TCPMaxBufferedPagesTotal
|
||||
tcpAssembler.MaxBufferedPagesPerConnection = config.TCPMaxBufferedPagesPerConn
|
||||
|
||||
udpSF := &udpStreamFactory{
|
||||
WorkerID: config.ID,
|
||||
Logger: config.Logger,
|
||||
@@ -97,25 +76,24 @@ func newWorker(config workerConfig) (*worker, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &worker{
|
||||
id: config.ID,
|
||||
packetChan: make(chan *workerPacket, config.ChanSize),
|
||||
logger: config.Logger,
|
||||
macResolver: config.MACResolver,
|
||||
tcpStreamFactory: tcpSF,
|
||||
tcpStreamPool: tcpStreamPool,
|
||||
tcpAssembler: tcpAssembler,
|
||||
udpStreamFactory: udpSF,
|
||||
udpStreamManager: udpSM,
|
||||
tcpFlowMgr: tcpMgr,
|
||||
udpSM: udpSM,
|
||||
modSerializeBuffer: gopacket.NewSerializeBuffer(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *worker) Feed(p *workerPacket) {
|
||||
func (w *worker) Feed(p *workerPacket) bool {
|
||||
select {
|
||||
case w.packetChan <- p:
|
||||
return true
|
||||
default:
|
||||
_ = p.SetVerdict(io.VerdictAccept, nil)
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
@@ -126,78 +104,116 @@ func (w *worker) Run(ctx context.Context) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case wPkt := <-w.packetChan:
|
||||
if wPkt == nil {
|
||||
case wp := <-w.packetChan:
|
||||
if wp == nil {
|
||||
return
|
||||
}
|
||||
pkt := gopacket.NewPacket(wPkt.Data, wPkt.LayerType, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
|
||||
v, b := w.handle(wPkt.StreamID, pkt, wPkt.SrcMAC, wPkt.DstMAC)
|
||||
_ = wPkt.SetVerdict(v, b)
|
||||
v, b := w.handle(wp)
|
||||
_ = wp.SetVerdict(v, b)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *worker) UpdateRuleset(r ruleset.Ruleset) error {
|
||||
if err := w.tcpStreamFactory.UpdateRuleset(r); err != nil {
|
||||
return err
|
||||
}
|
||||
return w.udpStreamFactory.UpdateRuleset(r)
|
||||
w.tcpFlowMgr.updateRuleset(r, 0)
|
||||
return w.udpSM.factory.UpdateRuleset(r)
|
||||
}
|
||||
|
||||
func (w *worker) handle(streamID uint32, p gopacket.Packet, srcMAC, dstMAC net.HardwareAddr) (io.Verdict, []byte) {
|
||||
netLayer, trLayer := p.NetworkLayer(), p.TransportLayer()
|
||||
if netLayer == nil || trLayer == nil {
|
||||
// Invalid packet
|
||||
func (w *worker) handle(wp *workerPacket) (io.Verdict, []byte) {
|
||||
data := wp.Data
|
||||
if len(data) == 0 {
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
ipFlow := netLayer.NetworkFlow()
|
||||
if len(srcMAC) == 0 && w.macResolver != nil {
|
||||
srcMAC = w.macResolver.Resolve(net.IP(ipFlow.Src().Raw()))
|
||||
}
|
||||
switch tr := trLayer.(type) {
|
||||
case *layers.TCP:
|
||||
return w.handleTCP(ipFlow, srcMAC, dstMAC, p.Metadata(), tr), nil
|
||||
case *layers.UDP:
|
||||
v, modPayload := w.handleUDP(streamID, ipFlow, srcMAC, dstMAC, tr)
|
||||
if v == io.VerdictAcceptModify && modPayload != nil {
|
||||
tr.Payload = modPayload
|
||||
_ = tr.SetNetworkLayerForChecksum(netLayer)
|
||||
_ = w.modSerializeBuffer.Clear()
|
||||
err := gopacket.SerializePacket(w.modSerializeBuffer,
|
||||
gopacket.SerializeOptions{
|
||||
FixLengths: true,
|
||||
ComputeChecksums: true,
|
||||
}, p)
|
||||
if err != nil {
|
||||
// Just accept without modification for now
|
||||
|
||||
ipVersion := data[0] >> 4
|
||||
if ipVersion == 4 {
|
||||
l3, transport, ok := ParseL3(data)
|
||||
if !ok {
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
switch l3.Protocol {
|
||||
case 6: // TCP
|
||||
tcp, payload, ok := ParseTCP(transport)
|
||||
if !ok {
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
return v, w.modSerializeBuffer.Bytes()
|
||||
verdict := w.tcpFlowMgr.handle(
|
||||
wp.StreamID, l3, tcp, payload,
|
||||
wp.SrcMAC, wp.DstMAC,
|
||||
)
|
||||
return verdict, nil
|
||||
|
||||
case 17: // UDP
|
||||
udp, payload, ok := ParseUDP(transport)
|
||||
if !ok {
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
v, modPayload := w.handleUDP(
|
||||
wp.StreamID, l3, udp, payload,
|
||||
wp.SrcMAC, wp.DstMAC,
|
||||
)
|
||||
if v == io.VerdictAcceptModify && modPayload != nil {
|
||||
return w.serializeModifiedUDP(data, l3, udp, transport, modPayload)
|
||||
}
|
||||
return v, nil
|
||||
|
||||
default:
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
return v, nil
|
||||
default:
|
||||
// Unsupported protocol
|
||||
}
|
||||
|
||||
// Ethernet frame path (for custom PacketIO)
|
||||
if ipVersion == 6 {
|
||||
// TODO: IPv6 support with raw parsing
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
|
||||
func (w *worker) handleTCP(ipFlow gopacket.Flow, srcMAC, dstMAC net.HardwareAddr, pMeta *gopacket.PacketMetadata, tcp *layers.TCP) io.Verdict {
|
||||
ctx := &tcpContext{
|
||||
PacketMetadata: pMeta,
|
||||
Verdict: tcpVerdictAccept,
|
||||
SrcMAC: srcMAC,
|
||||
DstMAC: dstMAC,
|
||||
func (w *worker) handleUDP(streamID uint32, l3 L3Info, udp UDPInfo, payload []byte, srcMAC, dstMAC net.HardwareAddr) (io.Verdict, []byte) {
|
||||
ipSrc := net.IP(l3.SrcIP[:])
|
||||
ipDst := net.IP(l3.DstIP[:])
|
||||
ipFlow := gopacket.NewFlow(layers.EndpointIPv4, ipSrc.To4(), ipDst.To4())
|
||||
udpFlow := gopacket.NewFlow(layers.EndpointUDPPort, []byte{byte(udp.SrcPort >> 8), byte(udp.SrcPort)}, []byte{byte(udp.DstPort >> 8), byte(udp.DstPort)})
|
||||
|
||||
if len(srcMAC) == 0 && w.macResolver != nil {
|
||||
srcMAC = w.macResolver.Resolve(ipSrc)
|
||||
}
|
||||
w.tcpAssembler.AssembleWithContext(ipFlow, tcp, ctx)
|
||||
return io.Verdict(ctx.Verdict)
|
||||
}
|
||||
|
||||
func (w *worker) handleUDP(streamID uint32, ipFlow gopacket.Flow, srcMAC, dstMAC net.HardwareAddr, udp *layers.UDP) (io.Verdict, []byte) {
|
||||
ctx := &udpContext{
|
||||
uc := &udpContext{
|
||||
Verdict: udpVerdictAccept,
|
||||
SrcMAC: srcMAC,
|
||||
DstMAC: dstMAC,
|
||||
}
|
||||
w.udpStreamManager.MatchWithContext(streamID, ipFlow, udp, ctx)
|
||||
return io.Verdict(ctx.Verdict), ctx.Packet
|
||||
// Temporarily set payload on a UDP layer so existing UDP handling works
|
||||
// We pass the payload through the context
|
||||
w.udpSM.MatchWithContext(streamID, ipFlow, &layers.UDP{
|
||||
BaseLayer: layers.BaseLayer{Payload: payload},
|
||||
SrcPort: layers.UDPPort(udp.SrcPort),
|
||||
DstPort: layers.UDPPort(udp.DstPort),
|
||||
}, uc)
|
||||
return io.Verdict(uc.Verdict), uc.Packet
|
||||
}
|
||||
|
||||
func (w *worker) serializeModifiedUDP(fullData []byte, l3 L3Info, udp UDPInfo, transport []byte, modPayload []byte) (io.Verdict, []byte) {
|
||||
ipPkt := gopacket.NewPacket(fullData, layers.LayerTypeIPv4, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
|
||||
netLayer := ipPkt.NetworkLayer()
|
||||
trLayer := ipPkt.TransportLayer()
|
||||
if netLayer == nil || trLayer == nil {
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
udpLayer, ok := trLayer.(*layers.UDP)
|
||||
if !ok {
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
udpLayer.Payload = modPayload
|
||||
_ = udpLayer.SetNetworkLayerForChecksum(netLayer)
|
||||
_ = w.modSerializeBuffer.Clear()
|
||||
err := gopacket.SerializePacket(w.modSerializeBuffer,
|
||||
gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: true}, ipPkt)
|
||||
if err != nil {
|
||||
return io.VerdictAccept, nil
|
||||
}
|
||||
return io.VerdictAcceptModify, w.modSerializeBuffer.Bytes()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user