github.com/TeaOSLab/EdgeNode@v1.3.8/internal/firewalls/nftables/table.go (about)

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  //go:build linux
     3  
     4  package nftables
     5  
     6  import (
     7  	"errors"
     8  	nft "github.com/google/nftables"
     9  	"github.com/iwind/TeaGo/types"
    10  	"strings"
    11  )
    12  
    13  type Table struct {
    14  	conn     *Conn
    15  	rawTable *nft.Table
    16  }
    17  
    18  func NewTable(conn *Conn, rawTable *nft.Table) *Table {
    19  	return &Table{
    20  		conn:     conn,
    21  		rawTable: rawTable,
    22  	}
    23  }
    24  
    25  func (this *Table) Raw() *nft.Table {
    26  	return this.rawTable
    27  }
    28  
    29  func (this *Table) Name() string {
    30  	return this.rawTable.Name
    31  }
    32  
    33  func (this *Table) Family() TableFamily {
    34  	return this.rawTable.Family
    35  }
    36  
    37  func (this *Table) GetChain(name string) (*Chain, error) {
    38  	rawChains, err := this.conn.Raw().ListChains()
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	for _, rawChain := range rawChains {
    43  		// must compare table name
    44  		if rawChain.Name == name && rawChain.Table.Name == this.rawTable.Name {
    45  			return NewChain(this.conn, this.rawTable, rawChain), nil
    46  		}
    47  	}
    48  	return nil, ErrChainNotFound
    49  }
    50  
    51  func (this *Table) AddChain(name string, chainPolicy *ChainPolicy) (*Chain, error) {
    52  	if len(name) > MaxChainNameLength {
    53  		return nil, errors.New("chain name too long (max " + types.String(MaxChainNameLength) + ")")
    54  	}
    55  
    56  	var rawChain = this.conn.Raw().AddChain(&nft.Chain{
    57  		Name:     name,
    58  		Table:    this.rawTable,
    59  		Hooknum:  nft.ChainHookInput,
    60  		Priority: nft.ChainPriorityFilter,
    61  		Type:     nft.ChainTypeFilter,
    62  		Policy:   chainPolicy,
    63  	})
    64  
    65  	err := this.conn.Commit()
    66  	if err != nil {
    67  		return nil, err
    68  	}
    69  	return NewChain(this.conn, this.rawTable, rawChain), nil
    70  }
    71  
    72  func (this *Table) AddAcceptChain(name string) (*Chain, error) {
    73  	var policy = ChainPolicyAccept
    74  	return this.AddChain(name, &policy)
    75  }
    76  
    77  func (this *Table) AddDropChain(name string) (*Chain, error) {
    78  	var policy = ChainPolicyDrop
    79  	return this.AddChain(name, &policy)
    80  }
    81  
    82  func (this *Table) DeleteChain(name string) error {
    83  	chain, err := this.GetChain(name)
    84  	if err != nil {
    85  		if err == ErrChainNotFound {
    86  			return nil
    87  		}
    88  		return err
    89  	}
    90  	this.conn.Raw().DelChain(chain.Raw())
    91  	return this.conn.Commit()
    92  }
    93  
    94  func (this *Table) GetSet(name string) (*Set, error) {
    95  	rawSet, err := this.conn.Raw().GetSetByName(this.rawTable, name)
    96  	if err != nil {
    97  		if strings.Contains(err.Error(), "no such file or directory") {
    98  			return nil, ErrSetNotFound
    99  		}
   100  		return nil, err
   101  	}
   102  
   103  	return NewSet(this.conn, rawSet), nil
   104  }
   105  
   106  func (this *Table) AddSet(name string, options *SetOptions) (*Set, error) {
   107  	if len(name) > MaxSetNameLength {
   108  		return nil, errors.New("set name too long (max " + types.String(MaxSetNameLength) + ")")
   109  	}
   110  
   111  	if options == nil {
   112  		options = &SetOptions{}
   113  	}
   114  	var rawSet = &nft.Set{
   115  		Table:      this.rawTable,
   116  		ID:         options.Id,
   117  		Name:       name,
   118  		Anonymous:  options.Anonymous,
   119  		Constant:   options.Constant,
   120  		Interval:   options.Interval,
   121  		IsMap:      options.IsMap,
   122  		HasTimeout: options.HasTimeout,
   123  		Timeout:    options.Timeout,
   124  		KeyType:    options.KeyType,
   125  		DataType:   options.DataType,
   126  	}
   127  	err := this.conn.Raw().AddSet(rawSet, nil)
   128  	if err != nil {
   129  		return nil, err
   130  	}
   131  
   132  	err = this.conn.Commit()
   133  	if err != nil {
   134  		return nil, err
   135  	}
   136  
   137  	return NewSet(this.conn, rawSet), nil
   138  }
   139  
   140  func (this *Table) DeleteSet(name string) error {
   141  	set, err := this.GetSet(name)
   142  	if err != nil {
   143  		if err == ErrSetNotFound {
   144  			return nil
   145  		}
   146  		return err
   147  	}
   148  
   149  	this.conn.Raw().DelSet(set.Raw())
   150  	return this.conn.Commit()
   151  }
   152  
   153  func (this *Table) Flush() error {
   154  	this.conn.Raw().FlushTable(this.rawTable)
   155  	return this.conn.Commit()
   156  }