OpenGFW/engine/engine.go

119 lines
2.7 KiB
Go
Raw Normal View History

2024-01-20 08:45:01 +08:00
package engine
import (
"context"
"runtime"
"github.com/apernet/OpenGFW/io"
"github.com/apernet/OpenGFW/ruleset"
"github.com/google/gopacket"
"github.com/google/gopacket/layers"
)
var _ Engine = (*engine)(nil)
type engine struct {
logger Logger
io io.PacketIO
2024-01-20 08:45:01 +08:00
workers []*worker
}
func NewEngine(config Config) (Engine, error) {
workerCount := config.Workers
if workerCount <= 0 {
workerCount = runtime.NumCPU()
}
var err error
workers := make([]*worker, workerCount)
for i := range workers {
workers[i], err = newWorker(workerConfig{
ID: i,
ChanSize: config.WorkerQueueSize,
Logger: config.Logger,
Ruleset: config.Ruleset,
TCPMaxBufferedPagesTotal: config.WorkerTCPMaxBufferedPagesTotal,
TCPMaxBufferedPagesPerConn: config.WorkerTCPMaxBufferedPagesPerConn,
2024-04-09 02:54:35 +08:00
TCPTimeout: config.WorkerTCPTimeout,
2024-01-20 08:45:01 +08:00
UDPMaxStreams: config.WorkerUDPMaxStreams,
})
if err != nil {
return nil, err
}
}
return &engine{
logger: config.Logger,
io: config.IO,
2024-01-20 08:45:01 +08:00
workers: workers,
}, nil
}
func (e *engine) UpdateRuleset(r ruleset.Ruleset) error {
for _, w := range e.workers {
if err := w.UpdateRuleset(r); err != nil {
return err
}
}
return nil
}
func (e *engine) Run(ctx context.Context) error {
ioCtx, ioCancel := context.WithCancel(ctx)
defer ioCancel() // Stop workers & IO
2024-01-20 08:45:01 +08:00
// Start workers
for _, w := range e.workers {
go w.Run(ioCtx)
}
// Register IO callback
errChan := make(chan error, 1)
err := e.io.Register(ioCtx, func(p io.Packet, err error) bool {
2024-01-20 08:45:01 +08:00
if err != nil {
errChan <- err
return false
2024-01-20 08:45:01 +08:00
}
return e.dispatch(p)
})
if err != nil {
return err
2024-01-20 08:45:01 +08:00
}
// Block until IO errors or context is cancelled
select {
case err := <-errChan:
return err
case <-ctx.Done():
return nil
}
}
// dispatch dispatches a packet to a worker.
func (e *engine) dispatch(p io.Packet) bool {
2024-01-20 08:45:01 +08:00
data := p.Data()
ipVersion := data[0] >> 4
var layerType gopacket.LayerType
if ipVersion == 4 {
layerType = layers.LayerTypeIPv4
} else if ipVersion == 6 {
layerType = layers.LayerTypeIPv6
} else {
// Unsupported network layer
_ = e.io.SetVerdict(p, io.VerdictAcceptStream, nil)
2024-01-20 08:45:01 +08:00
return true
}
// Convert to gopacket.Packet
packet := gopacket.NewPacket(data, layerType, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
packet.Metadata().Timestamp = p.Timestamp()
2024-01-20 08:45:01 +08:00
// Load balance by stream ID
index := p.StreamID() % uint32(len(e.workers))
e.workers[index].Feed(&workerPacket{
StreamID: p.StreamID(),
Packet: packet,
SetVerdict: func(v io.Verdict, b []byte) error {
return e.io.SetVerdict(p, v, b)
2024-01-20 08:45:01 +08:00
},
})
return true
}