From 6757d9ad995d32efc4dca68a02be4f000e8c376f Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Sun, 28 May 2023 11:51:43 +0800
Subject: [PATCH] add

---
 web/Web.go | 163 +++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 115 insertions(+), 48 deletions(-)

diff --git a/web/Web.go b/web/Web.go
index 65c4c71..50ff8d1 100644
--- a/web/Web.go
+++ b/web/Web.go
@@ -156,75 +156,145 @@ func (t *WebPath) Store(path string, f func(w http.ResponseWriter, r *http.Reque
 	}
 }
 
-type CountLimits struct {
-	g []countLimit
+type Limits struct {
+	g []*limitItem
 	l sync.RWMutex
 }
 
-type countLimit struct {
-	cidr      *net.IPNet
-	available int
-}
-
-func (t *CountLimits) SetMaxCount(cidr string, max int) {
-	if _, cidrx, err := net.ParseCIDR(cidr); err != nil {
-		panic(err)
-	} else {
-		t.g = append(t.g, countLimit{cidrx, max})
-	}
+func (t *Limits) AddLimitItem(item *limitItem) {
+	t.g = append(t.g, item)
 }
 
-func (t *CountLimits) ReachMax(r *http.Request) (isOverflow bool) {
+func (t *Limits) ReachMax(r *http.Request) (isOverflow bool) {
 	if len(t.g) == 0 {
 		return
 	}
 	ip := net.ParseIP(strings.Split(r.RemoteAddr, ":")[0])
-	t.l.RLock()
-	defer t.l.RUnlock()
-	for i := 0; i < len(t.g); i++ {
-		if !t.g[i].cidr.Contains(ip) {
-			continue
+	for i := 0; !isOverflow && i < len(t.g); i++ {
+		var match bool
+		t.g[i].l.RLock()
+		for b := 0; !match && b < len(t.g[i].matchfs); b++ {
+			switch t.g[i].matchfs[b].k {
+			case rcidr:
+				match = t.g[i].matchfs[b].f(ip)
+			case rreq:
+				match = t.g[i].matchfs[b].f(r)
+			default:
+			}
 		}
-		if t.g[i].available == 0 {
+		if match && t.g[i].available == 0 {
 			isOverflow = true
-			break
 		}
+		t.g[i].l.RUnlock()
 	}
 	return
 }
 
-func (t *CountLimits) AddCount(r *http.Request) (isOverflow bool) {
+func (t *Limits) AddCount(r *http.Request) (isOverflow bool) {
 	if len(t.g) == 0 {
 		return
 	}
 	ip := net.ParseIP(strings.Split(r.RemoteAddr, ":")[0])
-	t.l.Lock()
-	defer t.l.Unlock()
-	var match []int
-	for i := 0; i < len(t.g); i++ {
-		if !t.g[i].cidr.Contains(ip) {
-			continue
+	var matchs []int
+
+	for i := 0; !isOverflow && i < len(t.g); i++ {
+		var match bool
+		t.g[i].l.RLock()
+		for b := 0; !match && b < len(t.g[i].matchfs); b++ {
+			switch t.g[i].matchfs[b].k {
+			case rcidr:
+				match = t.g[i].matchfs[b].f(ip)
+			case rreq:
+				match = t.g[i].matchfs[b].f(r)
+			default:
+			}
 		}
-		if t.g[i].available == 0 {
-			isOverflow = true
-			return
+		if match {
+			if t.g[i].available == 0 {
+				isOverflow = true
+			} else {
+				matchs = append(matchs, i)
+			}
 		}
-		match = append(match, i)
-	}
-	for i := 0; i < len(match); i++ {
-		t.g[match[i]].available -= 1
+		t.g[i].l.RUnlock()
 	}
-	go func() {
-		<-r.Context().Done()
+	if !isOverflow && len(matchs) != 0 {
 		t.l.Lock()
-		defer t.l.Unlock()
-		for i := 0; i < len(match); i++ {
-			t.g[match[i]].available += 1
+		for i := 0; !isOverflow && i < len(matchs); i++ {
+			t.g[matchs[i]].l.RLock()
+			if t.g[matchs[i]].available == 0 {
+				isOverflow = true
+			}
+			t.g[matchs[i]].l.RUnlock()
+		}
+		for i := 0; !isOverflow && i < len(matchs); i++ {
+			t.g[matchs[i]].l.Lock()
+			t.g[matchs[i]].available -= 1
+			t.g[matchs[i]].l.Unlock()
+		}
+		t.l.Unlock()
+		if !isOverflow {
+			go func() {
+				<-r.Context().Done()
+				t.l.Lock()
+				for i := 0; i < len(matchs); i++ {
+					t.g[matchs[i]].l.Lock()
+					t.g[matchs[i]].available += 1
+					t.g[matchs[i]].l.Unlock()
+				}
+				t.l.Unlock()
+			}()
 		}
-	}()
+	}
 	return
 }
 
+const (
+	rcidr = iota
+	rreq
+)
+
+type limitItem struct {
+	matchfs   []matchFunc
+	available int
+	l         sync.RWMutex
+}
+
+type matchFunc struct {
+	k int
+	f func(any) (match bool)
+}
+
+func NewLimitItem(max int) *limitItem {
+	return &limitItem{
+		available: max,
+	}
+}
+
+func (t *limitItem) Cidr(cidr string) *limitItem {
+	if _, cidrx, err := net.ParseCIDR(cidr); err != nil {
+		panic(err)
+	} else {
+		t.matchfs = append(t.matchfs, matchFunc{
+			rcidr,
+			func(a any) (match bool) {
+				return cidrx.Contains(a.(net.IP))
+			},
+		})
+	}
+	return t
+}
+
+func (t *limitItem) Request(matchf func(req *http.Request) (match bool)) *limitItem {
+	t.matchfs = append(t.matchfs, matchFunc{
+		rreq,
+		func(a any) (match bool) {
+			return matchf(a.(*http.Request))
+		},
+	})
+	return t
+}
+
 type Cache struct {
 	g   psync.MapExceeded[string, []byte]
 	gcL atomic.Int64
@@ -246,7 +316,7 @@ func (t cacheRes) WriteHeader(statusCode int) {
 	t.writeHeaderf(statusCode)
 }
 
-func (t *Cache) IsCache(key string) (res *[]byte, isCache bool) {
+func (t *Cache) IsCache(key string) (data *[]byte, isCache bool) {
 	return t.g.Load(key)
 }
 
@@ -274,12 +344,9 @@ func (t *Cache) Cache(key string, aliveDur time.Duration, w http.ResponseWriter)
 	return res
 }
 
-func (t *Cache) Store(key string, aliveDur time.Duration, data *[]byte) {
-	t.g.Store(key, data, aliveDur)
-	if s := int64(t.g.Len()); s > 10 && t.gcL.Load() <= s {
-		t.gcL.Store(s * 2)
-		t.g.GC()
-	}
+func WithStatusCode(w http.ResponseWriter, code int) {
+	w.WriteHeader(code)
+	_, _ = w.Write([]byte(http.StatusText(code)))
 }
 
 func Easy_boot() *Web {
-- 
2.39.2