github.com/sohaha/zlsgo@v1.7.13-0.20240501141223-10dd1a906f76/znet/limiter/limiter.go (about)

     1  package limiter
     2  
     3  import (
     4  	"net/http"
     5  	"sort"
     6  	"time"
     7  
     8  	"github.com/sohaha/zlsgo/znet"
     9  )
    10  
    11  // Rule user access control strategy
    12  type Rule struct {
    13  	rules []*singleRule
    14  }
    15  
    16  // New Newlimiter
    17  func New(allowed uint64, overflow ...func(c *znet.Context)) znet.HandlerFunc {
    18  	r := NewRule()
    19  	f := func(c *znet.Context) {
    20  		c.String(http.StatusTooManyRequests, http.StatusText(http.StatusTooManyRequests))
    21  	}
    22  	if len(overflow) > 0 {
    23  		f = overflow[0]
    24  	}
    25  	r.AddRule(time.Second, int(allowed))
    26  	return func(c *znet.Context) {
    27  		if !r.AllowVisitByIP(c.GetClientIP()) {
    28  			f(c)
    29  			c.Abort()
    30  			return
    31  		}
    32  		c.Next()
    33  	}
    34  }
    35  
    36  // NewRule Custom limiter rule
    37  func NewRule() *Rule {
    38  	return &Rule{}
    39  }
    40  
    41  // AddRule increase user access control strategy
    42  // If less than 1s, please use golang.org/x/time/rate
    43  func (r *Rule) AddRule(exp time.Duration, allowed int, estimated ...int) {
    44  	r.rules = append(r.rules, newRule(exp, allowed, estimated...))
    45  	sort.Slice(r.rules, func(i int, j int) bool {
    46  		return r.rules[i].defaultExpiration < r.rules[j].defaultExpiration
    47  	})
    48  }
    49  
    50  // AllowVisit Is access allowed
    51  func (r *Rule) AllowVisit(keys ...interface{}) bool {
    52  	if len(r.rules) == 0 {
    53  		return true
    54  	}
    55  	for i := range r.rules {
    56  		for _, key := range keys {
    57  			if !r.rules[i].allowVisit(key) {
    58  				return false
    59  			}
    60  		}
    61  	}
    62  	return true
    63  }
    64  
    65  // AllowVisitByIP AllowVisit IP
    66  func (r *Rule) AllowVisitByIP(ip string) bool {
    67  	i, _ := znet.IPToLong(ip)
    68  	if i == 0 {
    69  		return false
    70  	}
    71  	return r.AllowVisit(i)
    72  }