github.com/TeaOSLab/EdgeNode@v1.3.8/internal/waf/ip_list_test.go (about)

     1  // Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package waf_test
     4  
     5  import (
     6  	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
     7  	"github.com/TeaOSLab/EdgeNode/internal/utils/testutils"
     8  	"github.com/TeaOSLab/EdgeNode/internal/waf"
     9  	"github.com/iwind/TeaGo/Tea"
    10  	"github.com/iwind/TeaGo/assert"
    11  	_ "github.com/iwind/TeaGo/bootstrap"
    12  	"github.com/iwind/TeaGo/logs"
    13  	timeutil "github.com/iwind/TeaGo/utils/time"
    14  	"os"
    15  	"runtime"
    16  	"strconv"
    17  	"testing"
    18  	"time"
    19  )
    20  
    21  func TestNewIPList(t *testing.T) {
    22  	var list = waf.NewIPList(waf.IPListTypeDeny)
    23  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.1", time.Now().Unix())
    24  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.2", time.Now().Unix()+1)
    25  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.1", time.Now().Unix()+2)
    26  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeServer, 1, "127.0.0.3", time.Now().Unix()+3)
    27  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.10", time.Now().Unix()+10)
    28  
    29  	list.RemoveIP("127.0.0.1", 1, false)
    30  
    31  	logs.PrintAsJSON(list.IPMap(), t)
    32  	logs.PrintAsJSON(list.IdMap(), t)
    33  }
    34  
    35  func TestIPList_Expire(t *testing.T) {
    36  	if !testutils.IsSingleTesting() {
    37  		return
    38  	}
    39  
    40  	var list = waf.NewIPList(waf.IPListTypeDeny)
    41  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.1", time.Now().Unix())
    42  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.2", time.Now().Unix()+1)
    43  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.1", time.Now().Unix()+2)
    44  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.3", time.Now().Unix()+3)
    45  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.10", time.Now().Unix()+6)
    46  
    47  	var ticker = time.NewTicker(1 * time.Second)
    48  	for range ticker.C {
    49  		t.Log("====")
    50  		logs.PrintAsJSON(list.IPMap(), t)
    51  		logs.PrintAsJSON(list.IdMap(), t)
    52  		if len(list.IdMap()) == 0 {
    53  			break
    54  		}
    55  	}
    56  }
    57  
    58  func TestIPList_Contains(t *testing.T) {
    59  	var a = assert.NewAssertion(t)
    60  
    61  	var list = waf.NewIPList(waf.IPListTypeDeny)
    62  
    63  	for i := 0; i < 1_0000; i++ {
    64  		list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
    65  	}
    66  	//list.RemoveIP("192.168.1.100")
    67  	{
    68  		a.IsTrue(list.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1.100"))
    69  	}
    70  	{
    71  		a.IsFalse(list.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.2.100"))
    72  	}
    73  }
    74  
    75  func TestIPList_ContainsExpires(t *testing.T) {
    76  	var list = waf.NewIPList(waf.IPListTypeDeny)
    77  
    78  	for i := 0; i < 1_0000; i++ {
    79  		list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
    80  	}
    81  	// list.RemoveIP("192.168.1.100", 1, false)
    82  	for _, ip := range []string{"192.168.1.100", "192.168.2.100"} {
    83  		expiresAt, ok := list.ContainsExpires(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, ip)
    84  		t.Log(ok, expiresAt, timeutil.FormatTime("Y-m-d H:i:s", expiresAt))
    85  	}
    86  }
    87  
    88  func TestIPList_Save(t *testing.T) {
    89  	var a = assert.NewAssertion(t)
    90  
    91  	var list = waf.NewIPList(waf.IPListTypeAllow)
    92  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1.100", time.Now().Unix()+3600)
    93  	list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 2, "192.168.1.101", time.Now().Unix()+3600)
    94  
    95  	var file = Tea.Root + "/data/waf.iplist.json"
    96  	err := list.Save(file)
    97  	if err != nil {
    98  		t.Fatal(err)
    99  	}
   100  
   101  	var newList = waf.NewIPList(waf.IPListTypeAllow)
   102  	err = newList.Load(file)
   103  	if err != nil {
   104  		t.Fatal(err)
   105  	}
   106  
   107  	a.IsTrue(newList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1.100"))
   108  
   109  	_ = os.Remove(file)
   110  }
   111  
   112  func BenchmarkIPList_Add(b *testing.B) {
   113  	runtime.GOMAXPROCS(1)
   114  
   115  	var list = waf.NewIPList(waf.IPListTypeDeny)
   116  	for i := 0; i < b.N; i++ {
   117  		list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
   118  	}
   119  	b.Log(len(list.IPMap()))
   120  }
   121  
   122  func BenchmarkIPList_Has(b *testing.B) {
   123  	runtime.GOMAXPROCS(1)
   124  
   125  	var list = waf.NewIPList(waf.IPListTypeDeny)
   126  	b.ResetTimer()
   127  
   128  	for i := 0; i < 1_0000; i++ {
   129  		list.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
   130  	}
   131  
   132  	for i := 0; i < b.N; i++ {
   133  		list.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1.100")
   134  	}
   135  }