Skip to content

Commit

Permalink
rules init refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
glaslos committed Jan 2, 2025
1 parent 84b8373 commit f99e578
Showing 5 changed files with 20 additions and 49 deletions.
File renamed without changes.
File renamed without changes.
8 changes: 1 addition & 7 deletions glutton.go
Original file line number Diff line number Diff line change
@@ -97,17 +97,11 @@ func New(ctx context.Context) (*Glutton, error) {
}

var err error
g.rules, err = rules.ParseRuleSpec(rulesFile)
g.rules, err = rules.Init(rulesFile)
if err != nil {
return nil, err
}

for idx, rule := range g.rules {
if err := rules.InitRule(idx, rule); err != nil {
return nil, fmt.Errorf("failed to initialize rule: %w", err)
}
}

return g, nil
}

47 changes: 18 additions & 29 deletions rules/rules.go
Original file line number Diff line number Diff line change
@@ -16,14 +16,13 @@ import (
type RuleType int

const (
Rewrite RuleType = iota
UserConnHandler
UserConnHandler RuleType = iota
Drop
)

type Config struct {
Version int `yaml:"version"`
Rules []*Rule `yaml:"rules"`
Version int `yaml:"version"`
Rules Rules `yaml:"rules"`
}

type Rule struct {
@@ -36,39 +35,29 @@ type Rule struct {
ruleType RuleType
index int
matcher *pcap.BPF
port int
}

func (r *Rule) String() string {
return fmt.Sprintf("Rule: %s", r.Match)
}

func ParseRuleSpec(file io.Reader) (Rules, error) {
func Init(file io.Reader) (Rules, error) {
config := &Config{}
if err := yaml.NewDecoder(file).Decode(config); err != nil {
return nil, err
}

if config.Version == 0 {
// TODO: log warning
config.Version = 1
}

if config.Version != 1 {
return nil, fmt.Errorf("unsupported rules version: %v", config.Version)
if err := config.Rules.init(); err != nil {
return nil, err
}

return config.Rules, nil
}

func InitRule(idx int, rule *Rule) error {
func (rule *Rule) init(idx int) error {
if rule.isInit {
return nil
}

switch rule.Type {
case "rewrite":
rule.ruleType = Rewrite
case "conn_handler":
rule.ruleType = UserConnHandler
case "drop":
@@ -85,17 +74,6 @@ func InitRule(idx int, rule *Rule) error {
}
}

if rule.Target != "" {
var err error

if rule.ruleType == Rewrite {
rule.port, err = strconv.Atoi(rule.Target)
if err != nil {
return err
}
}
}

rule.index = idx
rule.isInit = true

@@ -197,3 +175,14 @@ func (rs Rules) Match(network string, srcAddr, dstAddr net.Addr) (*Rule, error)

return nil, nil
}

// Init initializes the rules
func (rs Rules) init() error {
for i, rule := range rs {
if err := rule.init(i); err != nil {
return err
}
rs[i] = rule
}
return nil
}
14 changes: 1 addition & 13 deletions rules/rules_test.go
Original file line number Diff line number Diff line change
@@ -15,7 +15,7 @@ import (
func parseRules(t *testing.T) Rules {
fh, err := os.Open("test.yaml")
require.NoError(t, err)
rules, err := ParseRuleSpec(fh)
rules, err := Init(fh)
require.NoError(t, err)
return rules
}
@@ -28,10 +28,6 @@ func TestParseRuleSpec(t *testing.T) {
func TestInitRule(t *testing.T) {
rules := parseRules(t)
require.NotEmpty(t, rules)
for i := range rules {
err := InitRule(i, rules[i])
require.NoError(t, err)
}

for _, rule := range rules {
require.True(t, rule.isInit)
@@ -71,10 +67,6 @@ func TestFakePacketBytes(t *testing.T) {
func TestRunMatchTCP(t *testing.T) {
rules := parseRules(t)
require.NotEmpty(t, rules)
for i := range rules {
err := InitRule(i, rules[i])
require.NoError(t, err)
}
conn, ln := testConn(t)
defer func() {
conn.Close()
@@ -94,10 +86,6 @@ func TestRunMatchTCP(t *testing.T) {
func TestRunMatchUDP(t *testing.T) {
rules := parseRules(t)
require.NotEmpty(t, rules)
for i := range rules {
err := InitRule(i, rules[i])
require.NoError(t, err)
}
conn, ln := testConn(t)
defer func() {
conn.Close()

0 comments on commit f99e578

Please sign in to comment.