bilibili-go-common/app/interface/main/upload/http/antispam/antispam.go

package antispam

import (
	"fmt"
	"strings"
	"time"

	"go-common/app/interface/main/upload/model"
	"go-common/library/cache/redis"
	"go-common/library/ecode"
	bm "go-common/library/net/http/blademaster"
	"go-common/library/net/http/blademaster/binding"
)

const (
	_prefixRate  = "r_%d_%s_%d"
	_prefixTotal = "t_%d_%s_%d"
)

// Antispam is a antispam instance.
type Antispam struct {
	redis     *redis.Pool
	limitFunc func(bucket, dir string) (model.DirRateConfig, bool)
	conf      *Config
}

// Config antispam config.
// On     bool // switch on/off
// Second int  // every N second allow N requests.
// N      int  // one unit allow N requests.
// Hour   int  // every N hour allow M requests.
// M      int  // one winodw allow M requests.
type Config struct {
	On     bool // switch on/off
	Second int  // every N second allow N requests.
	N      int  // one unit allow N requests.
	Hour   int  // every N hour allow M requests.
	M      int  // one winodw allow M requests.

	Redis *redis.Config
}

// New new a antispam service.
func New(c *Config, l func(bucket, dir string) (model.DirRateConfig, bool)) (s *Antispam) {
	if c == nil {
		panic("antispam config nil")
	}
	s = &Antispam{
		limitFunc: l,
		redis:     redis.NewPool(c.Redis),
	}
	s.conf = c
	return s
}

// NativeRate limit user + path second level
func (s *Antispam) NativeRate(c *bm.Context, path string, mid interface{}) (err error) {
	curSecond := int(time.Now().Unix())
	burst := curSecond - curSecond%s.conf.Second
	key := rateKey(mid.(int64), path, burst)
	return s.antispam(c, key, s.conf.Second, s.conf.N)
}

// Rate antispam by user + bucket + dir.
func (s *Antispam) Rate(c *bm.Context) (err error) {
	mid, ok := c.Get("mid")
	if !ok {
		return
	}
	ap := new(struct {
		Bucket string `form:"bucket" json:"bucket"`
		Dir    string `form:"dir" json:"dir"`
	})
	if err = c.BindWith(ap, binding.FormMultipart); err != nil {
		return s.NativeRate(c, c.Request.URL.Path, mid)
	}
	if ap.Bucket == "" || ap.Dir == "" { //not need dir limit
		return s.NativeRate(c, c.Request.URL.Path, mid)
	}
	limit, ok := s.limitFunc(ap.Bucket, ap.Dir)
	if !ok {
		return s.NativeRate(c, c.Request.URL.Path, mid)
	}
	if limit.SecondQPS == 0 || limit.CountQPS == 0 {
		return s.NativeRate(c, c.Request.URL.Path, mid)
	}
	path := strings.Join([]string{ap.Bucket, ap.Dir}, "_")
	curSecond := int(time.Now().Unix())
	burst := curSecond - curSecond%limit.SecondQPS
	key := rateKey(mid.(int64), path, burst)
	return s.antispam(c, key, limit.SecondQPS, limit.CountQPS)
}

func totalKey(mid int64, path string, burst int) string {
	return fmt.Sprintf(_prefixTotal, mid, path, burst)
}

// Total antispam by user + path hour level
func (s *Antispam) Total(c *bm.Context, hour, count int) (err error) {
	second := hour * 3600
	mid, ok := c.Get("mid")
	if !ok {
		return
	}
	curHour := int(time.Now().Unix() / 3600)
	burst := curHour - curHour%hour
	key := totalKey(mid.(int64), c.Request.URL.Path, burst)
	return s.antispam(c, key, second, count)
}

func (s *Antispam) antispam(c *bm.Context, key string, interval, count int) (err error) {
	conn := s.redis.Get(c)
	defer conn.Close()
	cur, err := redis.Int(conn.Do("GET", key))
	if err != nil && err != redis.ErrNil {
		err = nil
		return
	}
	if cur >= count {
		err = ecode.LimitExceed
		return
	}
	err = nil
	conn.Send("INCR", key)
	conn.Send("EXPIRE", key, interval)
	if err1 := conn.Flush(); err1 != nil {
		return
	}
	for i := 0; i < 2; i++ {
		if _, err1 := conn.Receive(); err1 != nil {
			return
		}
	}
	return
}

func rateKey(mid int64, path string, burst int) string {
	return fmt.Sprintf(_prefixRate, mid, path, burst)
}

func (s *Antispam) ServeHTTP(ctx *bm.Context) {
	// user + bucket + dir.
	if err := s.Rate(ctx); err != nil {
		ctx.JSON(nil, ecode.ServiceUnavailable)
		ctx.Abort()
		return
	}
	// user + path
	if err := s.Total(ctx, s.conf.Hour, s.conf.M); err != nil {
		ctx.JSON(nil, ecode.ServiceUnavailable)
		ctx.Abort()
		return
	}
}

// Handler is antispam handle.
func (s *Antispam) Handler() bm.HandlerFunc {
	return s.ServeHTTP
}