github.com/Tyktechnologies/tyk@v2.9.5+incompatible/gateway/session_manager.go (about)

     1  package gateway
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"time"
     7  
     8  	"github.com/TykTechnologies/leakybucket"
     9  	"github.com/TykTechnologies/leakybucket/memorycache"
    10  	"github.com/TykTechnologies/tyk/config"
    11  	"github.com/TykTechnologies/tyk/storage"
    12  	"github.com/TykTechnologies/tyk/user"
    13  )
    14  
    15  type PublicSession struct {
    16  	Quota struct {
    17  		QuotaMax       int64 `json:"quota_max"`
    18  		QuotaRemaining int64 `json:"quota_remaining"`
    19  		QuotaRenews    int64 `json:"quota_renews"`
    20  	} `json:"quota"`
    21  	RateLimit struct {
    22  		Rate float64 `json:"requests"`
    23  		Per  float64 `json:"per_unit"`
    24  	} `json:"rate_limit"`
    25  }
    26  
    27  const (
    28  	QuotaKeyPrefix     = "quota-"
    29  	RateLimitKeyPrefix = "rate-limit-"
    30  )
    31  
    32  // SessionLimiter is the rate limiter for the API, use ForwardMessage() to
    33  // check if a message should pass through or not
    34  type SessionLimiter struct {
    35  	bucketStore leakybucket.Storage
    36  }
    37  
    38  func (l *SessionLimiter) doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey string,
    39  	currentSession *user.SessionState,
    40  	store storage.Handler,
    41  	globalConf *config.Config,
    42  	apiLimit *user.APILimit, dryRun bool) bool {
    43  
    44  	var per, rate float64
    45  
    46  	if apiLimit != nil { // respect limit on API level
    47  		per = apiLimit.Per
    48  		rate = apiLimit.Rate
    49  	} else {
    50  		per = currentSession.Per
    51  		rate = currentSession.Rate
    52  	}
    53  
    54  	log.Debug("[RATELIMIT] Inbound raw key is: ", key)
    55  	log.Debug("[RATELIMIT] Rate limiter key is: ", rateLimiterKey)
    56  	pipeline := globalConf.EnableNonTransactionalRateLimiter
    57  
    58  	var ratePerPeriodNow int
    59  	if dryRun {
    60  		ratePerPeriodNow, _ = store.GetRollingWindow(rateLimiterKey, int64(per), pipeline)
    61  	} else {
    62  		ratePerPeriodNow, _ = store.SetRollingWindow(rateLimiterKey, int64(per), "-1", pipeline)
    63  	}
    64  
    65  	//log.Info("Num Requests: ", ratePerPeriodNow)
    66  
    67  	// Subtract by 1 because of the delayed add in the window
    68  	subtractor := 1
    69  	if globalConf.EnableSentinelRateLimiter || globalConf.DRLEnableSentinelRateLimiter {
    70  		// and another subtraction because of the preemptive limit
    71  		subtractor = 2
    72  	}
    73  	// The test TestRateLimitForAPIAndRateLimitAndQuotaCheck
    74  	// will only work with ththese two lines here
    75  	//log.Info("break: ", (int(currentSession.Rate) - subtractor))
    76  	if ratePerPeriodNow > int(rate)-subtractor {
    77  		// Set a sentinel value with expire
    78  		if globalConf.EnableSentinelRateLimiter || globalConf.DRLEnableSentinelRateLimiter {
    79  			if !dryRun {
    80  				store.SetRawKey(rateLimiterSentinelKey, "1", int64(per))
    81  			}
    82  		}
    83  		return true
    84  	}
    85  
    86  	return false
    87  }
    88  
    89  type sessionFailReason uint
    90  
    91  const (
    92  	sessionFailNone sessionFailReason = iota
    93  	sessionFailRateLimit
    94  	sessionFailQuota
    95  )
    96  
    97  func (l *SessionLimiter) limitSentinel(
    98  	currentSession *user.SessionState,
    99  	key string,
   100  	rateScope string,
   101  	store storage.Handler,
   102  	globalConf *config.Config,
   103  	apiLimit *user.APILimit,
   104  	dryRun bool,
   105  ) bool {
   106  	rateLimiterKey := RateLimitKeyPrefix + rateScope + currentSession.GetKeyHash()
   107  	rateLimiterSentinelKey := RateLimitKeyPrefix + rateScope + currentSession.GetKeyHash() + ".BLOCKED"
   108  
   109  	go l.doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey, currentSession, store, globalConf, apiLimit, dryRun)
   110  
   111  	// Check sentinel
   112  	_, sentinelActive := store.GetRawKey(rateLimiterSentinelKey)
   113  	if sentinelActive == nil {
   114  		// Sentinel is set, fail
   115  		return true
   116  	}
   117  	return false
   118  }
   119  
   120  func (l *SessionLimiter) limitRedis(
   121  	currentSession *user.SessionState,
   122  	key string,
   123  	rateScope string,
   124  	store storage.Handler,
   125  	globalConf *config.Config,
   126  	apiLimit *user.APILimit,
   127  	dryRun bool,
   128  ) bool {
   129  	rateLimiterKey := RateLimitKeyPrefix + rateScope + currentSession.GetKeyHash()
   130  	rateLimiterSentinelKey := RateLimitKeyPrefix + rateScope + currentSession.GetKeyHash() + ".BLOCKED"
   131  
   132  	if l.doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey, currentSession, store, globalConf, apiLimit, dryRun) {
   133  		return true
   134  	}
   135  	return false
   136  }
   137  func (l *SessionLimiter) limitDRL(
   138  	currentSession *user.SessionState,
   139  	key string,
   140  	rateScope string,
   141  	apiLimit *user.APILimit,
   142  	dryRun bool,
   143  ) bool {
   144  	// In-memory limiter
   145  	if l.bucketStore == nil {
   146  		l.bucketStore = memorycache.New()
   147  	}
   148  
   149  	bucketKey := key + ":" + rateScope + currentSession.LastUpdated
   150  	currRate := apiLimit.Rate
   151  	per := apiLimit.Per
   152  
   153  	// DRL will always overflow with more servers on low rates
   154  	rate := uint(currRate * float64(DRLManager.RequestTokenValue))
   155  	if rate < uint(DRLManager.CurrentTokenValue()) {
   156  		rate = uint(DRLManager.CurrentTokenValue())
   157  	}
   158  	userBucket, err := l.bucketStore.Create(bucketKey, rate, time.Duration(per)*time.Second)
   159  	if err != nil {
   160  		log.Error("Failed to create bucket!")
   161  		return true
   162  	}
   163  
   164  	if dryRun {
   165  		// if userBucket is empty and not expired.
   166  		if userBucket.Remaining() == 0 && time.Now().Before(userBucket.Reset()) {
   167  			return true
   168  		}
   169  	} else {
   170  		_, errF := userBucket.Add(uint(DRLManager.CurrentTokenValue()))
   171  		if errF != nil {
   172  			return true
   173  		}
   174  	}
   175  	return false
   176  }
   177  
   178  func (sfr sessionFailReason) String() string {
   179  	switch sfr {
   180  	case sessionFailNone:
   181  		return "sessionFailNone"
   182  	case sessionFailRateLimit:
   183  		return "sessionFailRateLimit"
   184  	case sessionFailQuota:
   185  		return "sessionFailQuota"
   186  	default:
   187  		return fmt.Sprintf("%d", uint(sfr))
   188  	}
   189  }
   190  
   191  // ForwardMessage will enforce rate limiting, returning a non-zero
   192  // sessionFailReason if session limits have been exceeded.
   193  // Key values to manage rate are Rate and Per, e.g. Rate of 10 messages
   194  // Per 10 seconds
   195  func (l *SessionLimiter) ForwardMessage(r *http.Request, currentSession *user.SessionState, key string, store storage.Handler, enableRL, enableQ bool, globalConf *config.Config, apiID string, dryRun bool) sessionFailReason {
   196  	// check for limit on API level (set to session by ApplyPolicies)
   197  	var apiLimit *user.APILimit
   198  	var allowanceScope string
   199  
   200  	if len(currentSession.GetAccessRights()) > 0 {
   201  		if rights, ok := currentSession.GetAccessRightByAPIID(apiID); !ok {
   202  			log.WithField("apiID", apiID).Debug("[RATE] unexpected apiID")
   203  
   204  			return sessionFailRateLimit
   205  		} else {
   206  			apiLimit = rights.Limit
   207  			allowanceScope = rights.AllowanceScope
   208  		}
   209  	}
   210  
   211  	if apiLimit == nil {
   212  		apiLimit = &user.APILimit{
   213  			QuotaMax:           currentSession.QuotaMax,
   214  			QuotaRenewalRate:   currentSession.QuotaRenewalRate,
   215  			QuotaRenews:        currentSession.QuotaRenews,
   216  			Rate:               currentSession.Rate,
   217  			Per:                currentSession.Per,
   218  			ThrottleInterval:   currentSession.ThrottleInterval,
   219  			ThrottleRetryLimit: currentSession.ThrottleRetryLimit,
   220  		}
   221  	}
   222  
   223  	// If rate is -1 or 0, it means unlimited and no need for rate limiting.
   224  	if enableRL && apiLimit.Rate > 0 {
   225  		rateScope := ""
   226  		if allowanceScope != "" {
   227  			rateScope = allowanceScope + "-"
   228  		}
   229  		if globalConf.EnableSentinelRateLimiter {
   230  			if l.limitSentinel(currentSession, key, rateScope, store, globalConf, apiLimit, dryRun) {
   231  				return sessionFailRateLimit
   232  			}
   233  		} else if globalConf.EnableRedisRollingLimiter {
   234  			if l.limitRedis(currentSession, key, rateScope, store, globalConf, apiLimit, dryRun) {
   235  				return sessionFailRateLimit
   236  			}
   237  		} else {
   238  			var n float64
   239  			if DRLManager.Servers != nil {
   240  				n = float64(DRLManager.Servers.Count())
   241  			}
   242  			rate := apiLimit.Rate / apiLimit.Per
   243  			c := globalConf.DRLThreshold
   244  			if c == 0 {
   245  				// defaults to 5
   246  				c = 5
   247  			}
   248  
   249  			if n <= 1 || n*c < rate {
   250  				// If we have 1 server, there is no need to strain redis at all the leaky
   251  				// bucket algorithm will suffice.
   252  				if l.limitDRL(currentSession, key, rateScope, apiLimit, dryRun) {
   253  					return sessionFailRateLimit
   254  				}
   255  			} else {
   256  				if l.limitRedis(currentSession, key, rateScope, store, globalConf, apiLimit, dryRun) {
   257  					return sessionFailRateLimit
   258  				}
   259  			}
   260  		}
   261  	}
   262  
   263  	if enableQ {
   264  		if globalConf.LegacyEnableAllowanceCountdown {
   265  			currentSession.Allowance = currentSession.Allowance - 1
   266  		}
   267  
   268  		if l.RedisQuotaExceeded(r, currentSession, allowanceScope, apiLimit, store) {
   269  			return sessionFailQuota
   270  		}
   271  	}
   272  
   273  	return sessionFailNone
   274  
   275  }
   276  
   277  func (l *SessionLimiter) RedisQuotaExceeded(r *http.Request, currentSession *user.SessionState, scope string, limit *user.APILimit, store storage.Handler) bool {
   278  	// Unlimited?
   279  	if limit.QuotaMax == -1 || limit.QuotaMax == 0 {
   280  		// No quota set
   281  		return false
   282  	}
   283  
   284  	quotaScope := ""
   285  	if scope != "" {
   286  		quotaScope = scope + "-"
   287  	}
   288  
   289  	rawKey := QuotaKeyPrefix + quotaScope + currentSession.GetKeyHash()
   290  	quotaRenewalRate := limit.QuotaRenewalRate
   291  	quotaRenews := limit.QuotaRenews
   292  	quotaMax := limit.QuotaMax
   293  
   294  	log.Debug("[QUOTA] Quota limiter key is: ", rawKey)
   295  	log.Debug("Renewing with TTL: ", quotaRenewalRate)
   296  	// INCR the key (If it equals 1 - set EXPIRE)
   297  	qInt := store.IncrememntWithExpire(rawKey, quotaRenewalRate)
   298  
   299  	// if the returned val is >= quota: block
   300  	if qInt-1 >= quotaMax {
   301  		renewalDate := time.Unix(quotaRenews, 0)
   302  		log.Debug("Renewal Date is: ", renewalDate)
   303  		log.Debug("As epoch: ", quotaRenews)
   304  		log.Debug("Session: ", currentSession)
   305  		log.Debug("Now:", time.Now())
   306  		if time.Now().After(renewalDate) {
   307  			// The renewal date is in the past, we should update the quota!
   308  			// Also, this fixes legacy issues where there is no TTL on quota buckets
   309  			log.Debug("Incorrect key expiry setting detected, correcting")
   310  			go store.DeleteRawKey(rawKey)
   311  			qInt = 1
   312  		} else {
   313  			// Renewal date is in the future and the quota is exceeded
   314  			return true
   315  		}
   316  
   317  	}
   318  
   319  	// If this is a new Quota period, ensure we let the end user know
   320  	if qInt == 1 {
   321  		quotaRenews = time.Now().Unix() + quotaRenewalRate
   322  		ctxScheduleSessionUpdate(r)
   323  	}
   324  
   325  	// If not, pass and set the values of the session to quotamax - counter
   326  	remaining := quotaMax - qInt
   327  	if remaining < 0 {
   328  		remaining = 0
   329  	}
   330  
   331  	for k, v := range currentSession.GetAccessRights() {
   332  		if v.Limit == nil {
   333  			continue
   334  		}
   335  
   336  		if v.AllowanceScope == scope {
   337  			v.Limit.QuotaRemaining = remaining
   338  			v.Limit.QuotaRenews = quotaRenews
   339  		}
   340  		currentSession.SetAccessRight(k, v)
   341  	}
   342  
   343  	if scope == "" {
   344  		currentSession.QuotaRemaining = remaining
   345  		currentSession.QuotaRenews = quotaRenews
   346  	}
   347  
   348  	return false
   349  }