Implementation of UserData struct encode/decode (#49)

This commit is contained in:
Serguei Bezverkhi 2019-08-15 02:25:30 -04:00 committed by Michael Stapelberg
parent 579fe47a77
commit 0b3d8b56f6
2 changed files with 4 additions and 124 deletions

View File

@ -846,80 +846,6 @@ func TestAddRuleWithPosition(t *testing.T) {
}
}
func TestAddRuleRuleID(t *testing.T) {
want := [][]byte{
// batch begin
[]byte("\x00\x00\x00\x0a"),
// nft add rule ip ipv4table ipv4chain-1 ip version 6 comment \0x01
[]byte("\x02\x00\x00\x00\x0e\x00\x01\x00\x69\x70\x76\x34\x74\x61\x62\x6c\x65\x00\x00\x00\x10\x00\x02\x00\x69\x70\x76\x34\x63\x68\x61\x69\x6e\x2d\x31\x00\xa8\x00\x04\x80\x34\x00\x01\x80\x0c\x00\x01\x00\x70\x61\x79\x6c\x6f\x61\x64\x00\x24\x00\x02\x80\x08\x00\x01\x00\x00\x00\x00\x01\x08\x00\x02\x00\x00\x00\x00\x01\x08\x00\x03\x00\x00\x00\x00\x00\x08\x00\x04\x00\x00\x00\x00\x01\x44\x00\x01\x80\x0c\x00\x01\x00\x62\x69\x74\x77\x69\x73\x65\x00\x34\x00\x02\x80\x08\x00\x01\x00\x00\x00\x00\x01\x08\x00\x02\x00\x00\x00\x00\x01\x08\x00\x03\x00\x00\x00\x00\x01\x0c\x00\x04\x80\x05\x00\x01\x00\xf0\x00\x00\x00\x0c\x00\x05\x80\x05\x00\x01\x00\x00\x00\x00\x00\x2c\x00\x01\x80\x08\x00\x01\x00\x63\x6d\x70\x00\x20\x00\x02\x80\x08\x00\x01\x00\x00\x00\x00\x01\x08\x00\x02\x00\x00\x00\x00\x00\x0c\x00\x03\x80\x05\x00\x01\x00\x60\x00\x00\x00\x08\x00\x07\x00\x00\x00\x00\x01"),
// batch end
[]byte("\x00\x00\x00\x0a"),
}
c := &nftables.Conn{
TestDial: func(req []netlink.Message) ([]netlink.Message, error) {
for idx, msg := range req {
b, err := msg.MarshalBinary()
if err != nil {
t.Fatal(err)
}
if len(b) < 16 {
continue
}
b = b[16:]
if len(want) == 0 {
t.Errorf("no want entry for message %d: %x", idx, b)
continue
}
if got, want := b, want[0]; !bytes.Equal(got, want) {
t.Errorf("message %d: %s", idx, linediff(nfdump(got), nfdump(want)))
}
want = want[1:]
}
return req, nil
},
}
c.AddRule(&nftables.Rule{
RuleID: 1,
Table: &nftables.Table{Name: "ipv4table", Family: nftables.TableFamilyIPv4},
Chain: &nftables.Chain{
Name: "ipv4chain-1",
Type: nftables.ChainTypeFilter,
Hooknum: nftables.ChainHookPrerouting,
Priority: 0,
},
Exprs: []expr.Any{
// [ payload load 1b @ network header + 0 => reg 1 ]
&expr.Payload{
DestRegister: 1,
Base: expr.PayloadBaseNetworkHeader,
Offset: 0, // Offset for a transport protocol header
Len: 1, // 1 bytes for port
},
// [ bitwise reg 1 = (reg=1 & 0x000000f0 ) ^ 0x00000000 ]
&expr.Bitwise{
SourceRegister: 1,
DestRegister: 1,
Len: 1,
Mask: []byte{0xf0},
Xor: []byte{0x0},
},
// [ cmp eq reg 1 0x00000060 ]
&expr.Cmp{
Op: expr.CmpOpEq,
Register: 1,
Data: []byte{(0x6 << 4)},
},
},
})
if err := c.Flush(); err != nil {
t.Fatal(err)
}
}
func TestAddChain(t *testing.T) {
tests := []struct {
name string

54
rule.go
View File

@ -31,10 +31,10 @@ var ruleHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.
type Rule struct {
Table *Table
Chain *Chain
RuleID uint32
Position uint64
Handle uint64
Exprs []expr.Any
UserData []byte
}
// GetRule returns the rules in the specified table and chain.
@ -81,52 +81,6 @@ func (cc *Conn) GetRule(t *Table, c *Chain) ([]*Rule, error) {
return rules, nil
}
// GetRuleHandle returns a specific rule's handle. Rule is identified by Table, Chain and RuleID.
func (cc *Conn) GetRuleHandle(t *Table, c *Chain, ruleID uint32) (uint64, error) {
conn, err := cc.dialNetlink()
if err != nil {
return 0, err
}
defer conn.Close()
if ruleID == 0 {
return 0, fmt.Errorf("rule's id cannot be 0")
}
data, err := netlink.MarshalAttributes([]netlink.Attribute{
{Type: unix.NFTA_RULE_TABLE, Data: []byte(t.Name + "\x00")},
{Type: unix.NFTA_RULE_CHAIN, Data: []byte(c.Name + "\x00")},
{Type: unix.NFTA_RULE_USERDATA, Data: binaryutil.BigEndian.PutUint32(ruleID)},
})
if err != nil {
return 0, err
}
message := netlink.Message{
Header: netlink.Header{
Type: netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_GETRULE),
Flags: netlink.Request | netlink.Acknowledge | netlink.Dump | unix.NLM_F_ECHO,
},
Data: append(extraHeader(uint8(t.Family), 0), data...),
}
if _, err := conn.SendMessages([]netlink.Message{message}); err != nil {
return 0, fmt.Errorf("SendMessages: %v", err)
}
reply, err := conn.Receive()
if err != nil {
return 0, fmt.Errorf("Receive: %v", err)
}
for _, msg := range reply {
rr, err := ruleFromMsg(msg)
if err != nil {
return 0, err
}
if rr.RuleID == ruleID {
return rr.Handle, nil
}
}
return 0, fmt.Errorf("rule with id %d is not found", ruleID)
}
// AddRule adds the specified Rule
func (cc *Conn) AddRule(r *Rule) *Rule {
exprAttrs := make([]netlink.Attribute, len(r.Exprs))
@ -144,9 +98,9 @@ func (cc *Conn) AddRule(r *Rule) *Rule {
msgData := []byte{}
msgData = append(msgData, data...)
var flags netlink.HeaderFlags
if r.RuleID != 0 {
if r.UserData != nil {
msgData = append(msgData, cc.marshalAttr([]netlink.Attribute{
{Type: unix.NFTA_RULE_USERDATA, Data: binaryutil.BigEndian.PutUint32(r.RuleID)},
{Type: unix.NFTA_RULE_USERDATA, Data: r.UserData},
})...)
}
if r.Position != 0 {
@ -286,7 +240,7 @@ func ruleFromMsg(msg netlink.Message) (*Rule, error) {
case unix.NFTA_RULE_HANDLE:
r.Handle = ad.Uint64()
case unix.NFTA_RULE_USERDATA:
r.RuleID = ad.Uint32()
r.UserData = ad.Bytes()
}
}
return &r, ad.Err()