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