github.com/avenga/couper@v1.12.2/handler/ratelimit/config.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sort"
     7  	"time"
     8  
     9  	"github.com/avenga/couper/config"
    10  	"github.com/sirupsen/logrus"
    11  )
    12  
    13  const (
    14  	dummy = iota
    15  	modeBlock
    16  	modeWait
    17  	windowFixed
    18  	windowSliding
    19  )
    20  
    21  type RateLimit struct {
    22  	count       uint
    23  	logger      *logrus.Entry
    24  	mode        int
    25  	period      time.Duration
    26  	periodStart time.Time
    27  	perPeriod   uint
    28  	ringBuffer  *ringBuffer
    29  	window      int
    30  	quitCh      <-chan struct{}
    31  }
    32  
    33  type RateLimits []*RateLimit
    34  
    35  func ConfigureRateLimits(ctx context.Context, limits config.RateLimits, logger *logrus.Entry) (RateLimits, error) {
    36  	var (
    37  		mode       int
    38  		rateLimits RateLimits
    39  		window     int
    40  	)
    41  
    42  	uniqueDurations := make(map[time.Duration]struct{})
    43  
    44  	for _, limit := range limits {
    45  		d, err := config.ParseDuration("period", limit.Period, 0)
    46  		if err != nil {
    47  			return nil, err
    48  		}
    49  
    50  		if d == 0 {
    51  			return nil, fmt.Errorf("'period' must not be 0 (zero)")
    52  		}
    53  		if limit.PerPeriod == 0 {
    54  			return nil, fmt.Errorf("'per_period' must not be 0 (zero)")
    55  		}
    56  
    57  		if _, ok := uniqueDurations[time.Duration(d.Nanoseconds())]; ok {
    58  			return nil, fmt.Errorf("duplicate period (%q) found", limit.Period)
    59  		}
    60  
    61  		uniqueDurations[time.Duration(d.Nanoseconds())] = struct{}{}
    62  
    63  		switch limit.PeriodWindow {
    64  		case "":
    65  			fallthrough
    66  		case "sliding":
    67  			window = windowSliding
    68  		case "fixed":
    69  			window = windowFixed
    70  		default:
    71  			return nil, fmt.Errorf("unsupported 'period_window' (%q) given", limit.PeriodWindow)
    72  		}
    73  
    74  		switch limit.Mode {
    75  		case "":
    76  			fallthrough
    77  		case "wait":
    78  			mode = modeWait
    79  		case "block":
    80  			mode = modeBlock
    81  		default:
    82  			return nil, fmt.Errorf("unsupported 'mode' (%q) given", limit.Mode)
    83  		}
    84  
    85  		rateLimit := &RateLimit{
    86  			logger:    logger,
    87  			mode:      mode,
    88  			period:    time.Duration(d.Nanoseconds()),
    89  			perPeriod: limit.PerPeriod,
    90  			window:    window,
    91  			quitCh:    ctx.Done(),
    92  		}
    93  
    94  		switch rateLimit.window {
    95  		case windowSliding:
    96  			rateLimit.ringBuffer = newRingBuffer(rateLimit.perPeriod)
    97  		}
    98  
    99  		rateLimits = append(rateLimits, rateLimit)
   100  	}
   101  
   102  	// Sort 'rateLimits' by 'period' DESC.
   103  	sort.Slice(rateLimits, func(i, j int) bool {
   104  		return rateLimits[i].period > rateLimits[j].period
   105  	})
   106  
   107  	return rateLimits, nil
   108  }
   109  
   110  // countRequest MUST only be called after checkCapacity()
   111  func (rl *RateLimit) countRequest() {
   112  	switch rl.window {
   113  	case windowFixed:
   114  		rl.count++
   115  	case windowSliding:
   116  		rl.ringBuffer.put(time.Now())
   117  	}
   118  }