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

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  //go:build linux
     3  
     4  package nftables_test
     5  
     6  import (
     7  	"errors"
     8  	"github.com/TeaOSLab/EdgeNode/internal/firewalls/nftables"
     9  	"github.com/iwind/TeaGo/types"
    10  	"github.com/mdlayher/netlink"
    11  	"net"
    12  	"testing"
    13  	"time"
    14  )
    15  
    16  func getIPv4Set(t *testing.T) *nftables.Set {
    17  	var table = getIPv4Table(t)
    18  	set, err := table.GetSet("test_ipv4_set")
    19  	if err != nil {
    20  		if err == nftables.ErrSetNotFound {
    21  			set, err = table.AddSet("test_ipv4_set", &nftables.SetOptions{
    22  				KeyType:    nftables.TypeIPAddr,
    23  				HasTimeout: true,
    24  			})
    25  			if err != nil {
    26  				t.Fatal(err)
    27  			}
    28  		} else {
    29  			t.Fatal(err)
    30  		}
    31  	}
    32  	return set
    33  }
    34  
    35  func TestSet_AddElement(t *testing.T) {
    36  	var set = getIPv4Set(t)
    37  	err := set.AddElement(net.ParseIP("192.168.2.31").To4(), &nftables.ElementOptions{Timeout: 86400 * time.Second}, false)
    38  	if err != nil {
    39  		t.Fatal(err)
    40  	}
    41  	t.Log("ok")
    42  }
    43  
    44  func TestSet_DeleteElement(t *testing.T) {
    45  	var set = getIPv4Set(t)
    46  	err := set.DeleteElement(net.ParseIP("192.168.2.31").To4())
    47  	if err != nil {
    48  		t.Fatal(err)
    49  	}
    50  	t.Log("ok")
    51  }
    52  
    53  func TestSet_Batch(t *testing.T) {
    54  	var batch = getIPv4Set(t).Batch()
    55  
    56  	for _, ip := range []string{"192.168.2.30", "192.168.2.31", "192.168.2.32", "192.168.2.33", "192.168.2.34"} {
    57  		var ipData = net.ParseIP(ip).To4()
    58  		//err := batch.DeleteElement(ipData)
    59  		//if err != nil {
    60  		//	t.Fatal(err)
    61  		//}
    62  		err := batch.AddElement(ipData, &nftables.ElementOptions{Timeout: 10 * time.Second})
    63  		if err != nil {
    64  			t.Fatal(err)
    65  		}
    66  	}
    67  
    68  	err := batch.Commit()
    69  	if err != nil {
    70  		t.Logf("%#v", errors.Unwrap(err).(*netlink.OpError))
    71  		t.Fatal(err)
    72  	}
    73  	t.Log("ok")
    74  }
    75  
    76  func TestSet_Add_Many(t *testing.T) {
    77  	var set = getIPv4Set(t)
    78  
    79  	for i := 0; i < 255; i++ {
    80  		t.Log(i)
    81  		for j := 0; j < 255; j++ {
    82  			var ip = "192.167." + types.String(i) + "." + types.String(j)
    83  			var ipData = net.ParseIP(ip).To4()
    84  			err := set.Batch().AddElement(ipData, &nftables.ElementOptions{Timeout: 3600 * time.Second})
    85  			if err != nil {
    86  				t.Fatal(err)
    87  			}
    88  
    89  			if j%10 == 0 {
    90  				err = set.Batch().Commit()
    91  				if err != nil {
    92  					t.Fatal(err)
    93  				}
    94  			}
    95  		}
    96  		err := set.Batch().Commit()
    97  		if err != nil {
    98  			t.Fatal(err)
    99  		}
   100  	}
   101  	t.Log("ok")
   102  }
   103  
   104  /**func TestSet_Flush(t *testing.T) {
   105  	var set = getIPv4Set(t)
   106  	err := set.Flush()
   107  	if err != nil {
   108  		t.Fatal(err)
   109  	}
   110  	t.Log("ok")
   111  }**/