refactor(io): nft/ipt generator func

This commit is contained in:
Haruue 2024-03-22 02:04:25 +08:00
parent 57c818038c
commit ef1416274d
No known key found for this signature in database
GPG Key ID: F6083B28CBCBC148

View File

@ -27,83 +27,60 @@ const (
nftTable = "opengfw" nftTable = "opengfw"
) )
var nftRulesForward = fmt.Sprintf(` func generateNftRules(local, rst bool) (*nftTableSpec, error) {
define ACCEPT_CTMARK=%d if local && rst {
define DROP_CTMARK=%d return nil, errors.New("tcp rst is not supported in local mode")
define QUEUE_NUM=%d
table %s %s {
chain FORWARD {
type filter hook forward priority filter; policy accept;
ct mark $ACCEPT_CTMARK counter accept
ct mark $DROP_CTMARK counter drop
counter queue num $QUEUE_NUM bypass
} }
} table := &nftTableSpec{
`, nfqueueConnMarkAccept, nfqueueConnMarkDrop, nfqueueNum, nftFamily, nftTable) Family: nftFamily,
Table: nftTable,
var nftRulesForwardRST = fmt.Sprintf(`
define ACCEPT_CTMARK=%d
define DROP_CTMARK=%d
define QUEUE_NUM=%d
table %s %s {
chain FORWARD {
type filter hook forward priority filter; policy accept;
ct mark $ACCEPT_CTMARK counter accept
ip protocol tcp ct mark $DROP_CTMARK counter reject with tcp reset
ct mark $DROP_CTMARK counter drop
counter queue num $QUEUE_NUM bypass
} }
} table.Defines = append(table.Defines, fmt.Sprintf("define ACCEPT_CTMARK=%d", nfqueueConnMarkAccept))
`, nfqueueConnMarkAccept, nfqueueConnMarkDrop, nfqueueNum, nftFamily, nftTable) table.Defines = append(table.Defines, fmt.Sprintf("define DROP_CTMARK=%d", nfqueueConnMarkDrop))
table.Defines = append(table.Defines, fmt.Sprintf("define QUEUE_NUM=%d", nfqueueNum))
var nftRulesLocal = fmt.Sprintf(` if local {
define ACCEPT_CTMARK=%d table.Chains = []nftChainSpec{
define DROP_CTMARK=%d {Chain: "INPUT", Header: "type filter hook input priority filter; policy accept;"},
define QUEUE_NUM=%d {Chain: "OUTPUT", Header: "type filter hook output priority filter; policy accept;"},
table %s %s {
chain INPUT {
type filter hook input priority filter; policy accept;
ct mark $ACCEPT_CTMARK counter accept
ct mark $DROP_CTMARK counter drop
counter queue num $QUEUE_NUM bypass
} }
chain OUTPUT { } else {
type filter hook output priority filter; policy accept; table.Chains = []nftChainSpec{
{Chain: "FORWARD", Header: "type filter hook forward priority filter; policy accept;"},
ct mark $ACCEPT_CTMARK counter accept
ct mark $DROP_CTMARK counter drop
counter queue num $QUEUE_NUM bypass
} }
} }
`, nfqueueConnMarkAccept, nfqueueConnMarkDrop, nfqueueNum, nftFamily, nftTable) for i := range table.Chains {
c := &table.Chains[i]
var iptRulesForward = []iptRule{ c.Rules = append(c.Rules, "ct mark $ACCEPT_CTMARK counter accept")
{"filter", "FORWARD", []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkAccept), "-j", "ACCEPT"}}, if rst {
{"filter", "FORWARD", []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkDrop), "-j", "DROP"}}, c.Rules = append(c.Rules, "ip protocol tcp ct mark $DROP_CTMARK counter reject with tcp reset")
{"filter", "FORWARD", []string{"-j", "NFQUEUE", "--queue-num", strconv.Itoa(nfqueueNum), "--queue-bypass"}}, }
c.Rules = append(c.Rules, "ct mark $DROP_CTMARK counter drop")
c.Rules = append(c.Rules, "counter queue num $QUEUE_NUM bypass")
}
return table, nil
} }
var iptRulesForwardRST = []iptRule{ func generateIptRules(local, rst bool) ([]iptRule, error) {
{"filter", "FORWARD", []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkAccept), "-j", "ACCEPT"}}, if local && rst {
{"filter", "FORWARD", []string{"-p", "tcp", "-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkDrop), "-j", "REJECT", "--reject-with", "tcp-reset"}}, return nil, errors.New("tcp rst is not supported in local mode")
{"filter", "FORWARD", []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkDrop), "-j", "DROP"}}, }
{"filter", "FORWARD", []string{"-j", "NFQUEUE", "--queue-num", strconv.Itoa(nfqueueNum), "--queue-bypass"}}, var chains []string
} if local {
chains = []string{"INPUT", "OUTPUT"}
} else {
chains = []string{"FORWARD"}
}
rules := make([]iptRule, 0, 4*len(chains))
for _, chain := range chains {
rules = append(rules, iptRule{"filter", chain, []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkAccept), "-j", "ACCEPT"}})
if rst {
rules = append(rules, iptRule{"filter", chain, []string{"-p", "tcp", "-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkDrop), "-j", "REJECT", "--reject-with", "tcp-reset"}})
}
rules = append(rules, iptRule{"filter", chain, []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkDrop), "-j", "DROP"}})
rules = append(rules, iptRule{"filter", chain, []string{"-j", "NFQUEUE", "--queue-num", strconv.Itoa(nfqueueNum), "--queue-bypass"}})
}
var iptRulesLocal = []iptRule{ return rules, nil
{"filter", "INPUT", []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkAccept), "-j", "ACCEPT"}},
{"filter", "INPUT", []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkDrop), "-j", "DROP"}},
{"filter", "INPUT", []string{"-j", "NFQUEUE", "--queue-num", strconv.Itoa(nfqueueNum), "--queue-bypass"}},
{"filter", "OUTPUT", []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkAccept), "-j", "ACCEPT"}},
{"filter", "OUTPUT", []string{"-m", "connmark", "--mark", strconv.Itoa(nfqueueConnMarkDrop), "-j", "DROP"}},
{"filter", "OUTPUT", []string{"-j", "NFQUEUE", "--queue-num", strconv.Itoa(nfqueueNum), "--queue-bypass"}},
} }
var _ PacketIO = (*nfqueuePacketIO)(nil) var _ PacketIO = (*nfqueuePacketIO)(nil)
@ -275,23 +252,17 @@ func (n *nfqueuePacketIO) Close() error {
} }
func (n *nfqueuePacketIO) setupNft(local, rst, remove bool) error { func (n *nfqueuePacketIO) setupNft(local, rst, remove bool) error {
var rules string rules, err := generateNftRules(local, rst)
if local { if err != nil {
rules = nftRulesLocal return err
} else {
if rst {
rules = nftRulesForwardRST
} else {
rules = nftRulesForward
} }
} rulesText := rules.String()
var err error
if remove { if remove {
err = nftDelete(nftFamily, nftTable) err = nftDelete(nftFamily, nftTable)
} else { } else {
// Delete first to make sure no leftover rules // Delete first to make sure no leftover rules
_ = nftDelete(nftFamily, nftTable) _ = nftDelete(nftFamily, nftTable)
err = nftAdd(rules) err = nftAdd(rulesText)
} }
if err != nil { if err != nil {
return err return err
@ -300,17 +271,10 @@ func (n *nfqueuePacketIO) setupNft(local, rst, remove bool) error {
} }
func (n *nfqueuePacketIO) setupIpt(local, rst, remove bool) error { func (n *nfqueuePacketIO) setupIpt(local, rst, remove bool) error {
var rules []iptRule rules, err := generateIptRules(local, rst)
if local { if err != nil {
rules = iptRulesLocal return err
} else {
if rst {
rules = iptRulesForwardRST
} else {
rules = iptRulesForward
} }
}
var err error
if remove { if remove {
err = iptsBatchDeleteIfExists([]*iptables.IPTables{n.ipt4, n.ipt6}, rules) err = iptsBatchDeleteIfExists([]*iptables.IPTables{n.ipt4, n.ipt6}, rules)
} else { } else {
@ -365,6 +329,42 @@ func nftDelete(family, table string) error {
return cmd.Run() return cmd.Run()
} }
type nftTableSpec struct {
Defines []string
Family, Table string
Chains []nftChainSpec
}
func (t *nftTableSpec) String() string {
chains := make([]string, 0, len(t.Chains))
for _, c := range t.Chains {
chains = append(chains, c.String())
}
return fmt.Sprintf(`
%s
table %s %s {
%s
}
`, strings.Join(t.Defines, "\n"), t.Family, t.Table, strings.Join(chains, ""))
}
type nftChainSpec struct {
Chain string
Header string
Rules []string
}
func (c *nftChainSpec) String() string {
return fmt.Sprintf(`
chain %s {
%s
%s
}
`, c.Chain, c.Header, strings.Join(c.Rules, "\n\x20\x20\x20\x20"))
}
type iptRule struct { type iptRule struct {
Table, Chain string Table, Chain string
RuleSpec []string RuleSpec []string