improve antiabuse and add testcase (#398)

* improve antiabuse and add testcase

* fix managers link and add managers_test

* make linter happy

* add ttl.cache

plugin/antiabuse/anti.go

@@ -1,4 +1,4 @@
-// Package antiabuse defines anti_abuse plugin ,support abuse words check and add/remove abuse words
+// Package antiabuse defines antiabuse plugin ,support abuse words check and add/remove abuse words
 package antiabuse
 
 import (
@@ -14,33 +14,24 @@ import (
 )
 
 func init() {
-	engine := control.Register("anti_abuse", &ctrl.Options[*zero.Ctx]{
+	engine := control.Register("antiabuse", &ctrl.Options[*zero.Ctx]{
 		DisableOnDefault:  false,
 		Help:              "违禁词检测",
 		PrivateDataFolder: "anti_abuse",
 	})
 	onceRule := fcext.DoOnceOnSuccess(func(ctx *zero.Ctx) bool {
+		managers = ctx.State["managers"].(*ctrl.Control[*zero.Ctx]).Manager
 		db.DBPath = engine.DataFolder() + "anti_abuse.db"
 		err := db.Open(time.Hour * 4)
 		if err != nil {
 			ctx.SendChain(message.Text("open db error: ", err))
 			return false
 		}
-		err = db.Create("banUser", &banUser{})
-		if err != nil {
-			ctx.SendChain(message.Text("create table error: ", err))
-			return false
-		}
 		err = db.Create("banWord", &banWord{})
 		if err != nil {
 			ctx.SendChain(message.Text("create table error: ", err))
 			return false
 		}
-		err = recoverUser()
-		if err != nil {
-			ctx.SendChain(message.Text("recover data error: ", err))
-			return false
-		}
 		err = recoverWord()
 		if err != nil {
 			ctx.SendChain(message.Text("recover data error: ", err))
@@ -48,29 +39,32 @@ func init() {
 		}
 		return true
 	})
-	engine.OnMessage(zero.OnlyGroup, onceRule, banRule)
+	engine.OnMessage(onceRule, zero.OnlyGroup, banRule)
 	engine.OnCommand("添加违禁词", zero.OnlyGroup, zero.AdminPermission, onceRule).Handle(
 		func(ctx *zero.Ctx) {
-			if err := insertWord(ctx.Event.GroupID, ctx.State["args"].(string)); err != nil {
+			args := ctx.State["args"].(string)
-				ctx.SendChain(message.Text("add ban word error:", err))
+			if err := insertWord(ctx.Event.GroupID, args); err != nil {
+				ctx.SendChain(message.Text("error:", err))
+			} else {
+				ctx.SendChain(message.Text(fmt.Sprintf("添加违禁词 %s 成功", args)))
 			}
 		})
 	engine.OnCommand("删除违禁词", zero.OnlyGroup, zero.AdminPermission, onceRule).Handle(
 		func(ctx *zero.Ctx) {
-			if err := deleteWord(ctx.Event.GroupID, ctx.State["args"].(string)); err != nil {
+			args := ctx.State["args"].(string)
-				ctx.SendChain(message.Text("add ban word error:", err))
+			if err := deleteWord(ctx.Event.GroupID, args); err != nil {
+				ctx.SendChain(message.Text("error:", err))
+			} else {
+				ctx.SendChain(message.Text(fmt.Sprintf("删除违禁词 %s 成功", args)))
 			}
 		})
 	engine.OnCommand("查看违禁词", zero.OnlyGroup, onceRule).Handle(
 		func(ctx *zero.Ctx) {
-			gidPrefix := fmt.Sprintf("%d-", ctx.Event.GroupID)
+			if set, ok := wordMap[ctx.Event.GroupID]; !ok {
-			var words []string
+				ctx.SendChain(message.Text("本群无违禁词"))
-			_ = wordSet.Iter(func(s string) error {
+			} else {
-				trueWord := strings.SplitN(s, gidPrefix, 1)[1]
+				ctx.SendChain(message.Text("本群违禁词有:", strings.Join(set.ToSlice(), " |")))
-				words = append(words, trueWord)
+			}
-				return nil
-			})
-			ctx.SendChain(message.Text("本群违禁词有:\n", strings.Join(words, " |")))
 		})
 
 }

plugin/antiabuse/database.go

@@ -1,81 +0,0 @@
-package antiabuse
-
-import (
-	"fmt"
-	"time"
-
-	sqlite "github.com/FloatTech/sqlite"
-)
-
-var db = &sqlite.Sqlite{}
-
-type banUser struct {
-	UUID    string `db:"uuid"`
-	DueTime int64  `db:"due_time"`
-}
-
-func insertUser(gid, uid int64) error {
-	obj := &banUser{fmt.Sprintf("%d-%d", gid, uid), time.Now().Add(4 * time.Hour).UnixNano()}
-	return db.Insert("banUser", obj)
-}
-
-func deleteUser(gid, uid int64) error {
-	sql := fmt.Sprintf("WHERE uuid=%d-%d", gid, uid)
-	return db.Del("banUser", sql)
-}
-
-func recoverUser() error {
-	if !db.CanFind("banUser", "") {
-		return nil
-	}
-	obj := &banUser{}
-	var uuids []string
-	err := db.FindFor("banUser", obj, "", func() error {
-		if time.Now().UnixNano() < obj.DueTime {
-			uuids = append(uuids, obj.UUID)
-		} else {
-			if err := db.Del("banUser", "WHERE uuid="+obj.UUID); err != nil {
-				return err
-			}
-		}
-		return nil
-	},
-	)
-	if err != nil {
-		return err
-	}
-	banSet.AddMany(uuids)
-	return nil
-}
-
-type banWord struct {
-	GroupWord string `db:"group_word"`
-}
-
-func insertWord(gid int64, word string) error {
-	obj := &banWord{fmt.Sprintf("%d-%s", gid, word)}
-	return db.Insert("banWord", obj)
-}
-
-func deleteWord(gid int64, word string) error {
-	sql := fmt.Sprintf("WHERE group_word = %d-%s", gid, word)
-	return db.Del("banWord", sql)
-}
-
-func recoverWord() error {
-	if !db.CanFind("banWord", "") {
-		return nil
-	}
-	obj := &banWord{}
-	var groupWords []string
-	err := db.FindFor("banWord", obj, "", func() error {
-		groupWords = append(groupWords, obj.GroupWord)
-		return nil
-	},
-	)
-	if err != nil {
-		return err
-	}
-	wordSet.AddMany(groupWords)
-	return nil
-}

plugin/antiabuse/set.go

@@ -3,23 +3,27 @@ package antiabuse
 import "sync"
 
 //Set defines HashSet structure
-type Set struct {
+type Set[T comparable] struct {
 	sync.RWMutex
-	m map[string]struct{}
+	m map[T]struct{}
 }
 
-var banSet = &Set{m: make(map[string]struct{})}
+// NewSet creates Set with optional key(s)
-var wordSet = &Set{m: make(map[string]struct{})}
+func NewSet[T comparable]() *Set[T] {
+	return &Set[T]{m: make(map[T]struct{})}
+}
 
-// Add adds element to Set
+// Add adds key(s) to Set
-func (s *Set) Add(key string) {
+func (s *Set[T]) Add(key ...T) {
 	s.Lock()
 	defer s.Unlock()
-	s.m[key] = struct{}{}
+	for _, k := range key {
+		s.m[k] = struct{}{}
+	}
 }
 
-// Include asserts element in Set
+// Include asserts key in Set
-func (s *Set) Include(key string) bool {
+func (s *Set[T]) Include(key T) bool {
 	s.RLock()
 	defer s.RUnlock()
 	_, ok := s.m[key]
@@ -27,9 +31,9 @@ func (s *Set) Include(key string) bool {
 }
 
 // Iter calls f when traversing Set
-func (s *Set) Iter(f func(string) error) error {
+func (s *Set[T]) Iter(f func(T) error) error {
-	s.Lock()
+	s.RLock()
-	defer s.Unlock()
+	defer s.RUnlock()
 	var err error
 	for key := range s.m {
 		err = f(key)
@@ -40,18 +44,19 @@ func (s *Set) Iter(f func(string) error) error {
 	return nil
 }
 
-// Remove removes element from Set
+// Remove removes key from Set
-func (s *Set) Remove(key string) {
+func (s *Set[T]) Remove(key T) {
 	s.Lock()
 	defer s.Unlock()
 	delete(s.m, key)
 }
 
-// AddMany adds multiple elements to Set
+// ToSlice convert Set to slice
-func (s *Set) AddMany(keys []string) {
+func (s *Set[T]) ToSlice() (res []T) {
-	s.Lock()
+	s.RLock()
-	defer s.Unlock()
+	defer s.RUnlock()
-	for _, k := range keys {
+	for key := range s.m {
-		s.m[k] = struct{}{}
+		res = append(res, key)
 	}
+	return res
 }

plugin/antiabuse/utils.go

@@ -1,47 +1,110 @@
 package antiabuse
 
 import (
+	"errors"
 	"fmt"
+	"hash/crc32"
 	"strings"
 	"time"
 
+	sqlite "github.com/FloatTech/sqlite"
+	"github.com/FloatTech/ttl"
+	ctrl "github.com/FloatTech/zbpctrl"
+	"github.com/sirupsen/logrus"
 	zero "github.com/wdvxdr1123/ZeroBot"
 	"github.com/wdvxdr1123/ZeroBot/message"
 )
 
+var managers *ctrl.Manager[*zero.Ctx] //managers lazy load
+var errBreak = errors.New("break")
+var db = &sqlite.Sqlite{}
+var crc32Table = crc32.MakeTable(crc32.IEEE)
+var wordMap = make(map[int64]*Set[string])
+
+func onDel(uid int64, _ struct{}) {
+	if managers == nil {
+		return
+	}
+	if err := managers.DoUnblock(uid); err != nil {
+		logrus.Error("do unblock error:", err)
+	}
+}
+
+var cache = ttl.NewCacheOn[int64, struct{}](4*time.Hour, [4]func(int64, struct{}){
+	nil, nil, onDel, nil})
+
+type banWord struct {
+	Crc32ID uint32 `db:"crc32_id"`
+	GroupID int64  `db:"group_id"`
+	Word    string `db:"word"`
+}
+
 func banRule(ctx *zero.Ctx) bool {
 	if !ctx.Event.IsToMe {
-		return false
+		return true
 	}
-	gid := ctx.Event.GroupID
 	uid := ctx.Event.UserID
-	uuid := fmt.Sprintf("%d-%d", gid, uid)
+	gid := ctx.Event.GroupID
-	if banSet.Include(uuid) {
+	wordSet := wordMap[gid]
-		return false
+	if wordSet == nil {
+		return true
+	}
+	err := wordSet.Iter(func(word string) error {
+		if strings.Contains(ctx.MessageString(), word) {
+			if err := managers.DoBlock(uid); err != nil {
+				return err
+			}
+			cache.Set(uid, struct{}{})
+			return errBreak
 		}
-	gidPrefix := fmt.Sprintf("%d-", ctx.Event.GroupID)
-	var words []string
-	_ = wordSet.Iter(func(s string) error {
-		trueWord := strings.SplitN(s, gidPrefix, 1)[1]
-		words = append(words, trueWord)
 		return nil
 	})
-	for _, word := range words {
+	if err != nil && err != errBreak {
-		if strings.Contains(ctx.MessageString(), word) {
+		ctx.SendChain(message.Text("block user error:", err))
-			if err := insertUser(gid, uid); err != nil {
-				ctx.SendChain(message.Text("ban error: ", err))
-			}
-			banSet.Add(uuid)
-			ctx.SetGroupBan(gid, uid, 4*3600)
-			time.AfterFunc(4*time.Hour, func() {
-				banSet.Remove(uuid)
-				if err := deleteUser(gid, uid); err != nil {
-					ctx.SendChain(message.Text("ban error: ", err))
-				}
-			})
-			ctx.SendChain(message.Text("检测到违禁词，已封禁/屏蔽4小时"))
-			return false
-		}
-	}
 		return true
+	}
+	ctx.SetGroupBan(gid, uid, 4*3600)
+	ctx.SendChain(message.Text("检测到违禁词,已封禁/屏蔽4小时"))
+	return false
+}
+
+func insertWord(gid int64, word string) error {
+	str := fmt.Sprintf("%d-%s", gid, word)
+	checksum := crc32.Checksum([]byte(str), crc32Table)
+	obj := &banWord{checksum, gid, word}
+	if _, ok := wordMap[gid]; !ok {
+		wordMap[gid] = NewSet[string]()
+	}
+	wordMap[gid].Add(word)
+	return db.Insert("banWord", obj)
+}
+
+func deleteWord(gid int64, word string) error {
+	if _, ok := wordMap[gid]; !ok {
+		return errors.New("本群还没有违禁词~")
+	}
+	if !wordMap[gid].Include(word) {
+		return errors.New(word + " 不在本群违禁词集合中")
+	}
+	wordMap[gid].Remove(word)
+	str := fmt.Sprintf("%d-%s", gid, word)
+	checksum := crc32.Checksum([]byte(str), crc32Table)
+	sql := fmt.Sprintf("WHERE crc32_id = %d", checksum)
+	return db.Del("banWord", sql)
+}
+
+func recoverWord() error {
+	if !db.CanFind("banWord", "") {
+		return nil
+	}
+	obj := &banWord{}
+	err := db.FindFor("banWord", obj, "", func() error {
+		if _, ok := wordMap[obj.GroupID]; !ok {
+			wordMap[obj.GroupID] = NewSet[string]()
+		}
+		wordMap[obj.GroupID].Add(obj.Word)
+		return nil
+	},
+	)
+	return err
 }

plugin/antiabuse/utils_test.go

@@ -0,0 +1,202 @@
+package antiabuse
+
+import (
+	"os"
+	"sort"
+	"testing"
+)
+
+func TestInsertWord(t *testing.T) {
+	wordMap = make(map[int64]*Set[string])
+	defer func() {
+		wordMap = make(map[int64]*Set[string])
+	}()
+	path := "test.db"
+	defer func() {
+		err := os.Remove("test.db")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	db.DBPath = path
+	err := db.Open(0)
+	defer func() {
+		err := db.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	err = db.Create("banWord", &banWord{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := db.Drop("banWord")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	err = insertWord(123, "one")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if ok := wordMap[123].Include("one"); !ok {
+		t.Fatal(`wordMap[123] should found "one" but not`)
+	}
+	if !db.CanFind("banWord", "WHERE group_id=123 AND word= 'one' ") {
+		t.Fatal(`db should found 123-one but not`)
+	}
+}
+
+func TestDeleteWord(t *testing.T) {
+	wordMap = make(map[int64]*Set[string])
+	defer func() {
+		wordMap = make(map[int64]*Set[string])
+	}()
+	path := "test.db"
+	defer func() {
+		err := os.Remove("test.db")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	db.DBPath = path
+	err := db.Open(0)
+	defer func() {
+		err := db.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	err = db.Create("banWord", &banWord{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := db.Drop("banWord")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	err = insertWord(123, "one")
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = deleteWord(123, "one")
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestShowWord(t *testing.T) {
+	wordMap = make(map[int64]*Set[string])
+	defer func() {
+		wordMap = make(map[int64]*Set[string])
+	}()
+	path := "test.db"
+	defer func() {
+		err := os.Remove("test.db")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	db.DBPath = path
+	err := db.Open(0)
+	defer func() {
+		err := db.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	err = db.Create("banWord", &banWord{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := db.Drop("banWord")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	err = insertWord(123, "one")
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = insertWord(123, "one")
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = insertWord(123, "two")
+	if err != nil {
+		t.Fatal(err)
+	}
+	var db123 []string
+	var map123 []string
+	obj := &banWord{}
+	err = db.FindFor("banWord", obj, "WHERE group_id=123", func() error {
+		db123 = append(db123, obj.Word)
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	sort.Strings(db123)
+	if len(db123) != 2 || db123[0] != "one" || db123[1] != "two" {
+		t.Fatal("db should found 123-one and 123-two but not")
+	}
+	map123 = wordMap[123].ToSlice()
+	sort.Strings(map123)
+	if len(map123) != 2 || map123[0] != "one" || map123[1] != "two" {
+		t.Fatal("wordMap[123] should found 123-one and 123-two but not")
+	}
+}
+
+func TestRecoverWord(t *testing.T) {
+	wordMap = make(map[int64]*Set[string])
+	defer func() {
+		wordMap = make(map[int64]*Set[string])
+	}()
+	path := "test.db"
+	defer func() {
+		err := os.Remove("test.db")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	db.DBPath = path
+	err := db.Open(0)
+	defer func() {
+		err := db.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	err = db.Create("banWord", &banWord{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := db.Drop("banWord")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	err = insertWord(123, "one")
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = insertWord(123, "two")
+	if err != nil {
+		t.Fatal(err)
+	}
+	wordMap = make(map[int64]*Set[string])
+	err = recoverWord()
+	if err != nil {
+		t.Fatal(err)
+	}
+	map123 := wordMap[123].ToSlice()
+	sort.Strings(map123)
+	if len(map123) != 2 || map123[0] != "one" || map123[1] != "two" {
+		t.Fatal("wordMap[123] should found 123-one and 123-two but not")
+	}
+}