feat(web): Adding rate limit headers to all 429 errors (#365)

spinnaker · Mar 27, 2017 · 319d93d · 319d93d
1 parent 6ed2b99
commit 319d93d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 7 deletions.
diff --git a/gate-web/src/main/groovy/com/netflix/spinnaker/gate/ratelimit/Rate.java b/gate-web/src/main/groovy/com/netflix/spinnaker/gate/ratelimit/Rate.java
@@ -19,9 +19,10 @@
 
 public class Rate {
 
-  private static final String CAPACITY_HEADER = "X-RateLimit-Capacity";
-  private static final String REMAINING_HEADER = "X-RateLimit-Remaining";
-  private static final String RESET_HEADER = "X-RateLimit-Reset";
+  static final String CAPACITY_HEADER = "X-RateLimit-Capacity";
+  static final String REMAINING_HEADER = "X-RateLimit-Remaining";
+  static final String RESET_HEADER = "X-RateLimit-Reset";
+  static final String LEARNING_HEADER = "X-RateLimit-Learning";
 
   Integer capacity;
   Integer rateSeconds;
@@ -33,9 +34,10 @@ public boolean isThrottled() {
     return throttled;
   }
 
-  public void assignHttpHeaders(HttpServletResponse response) {
+  public void assignHttpHeaders(HttpServletResponse response, Boolean learning) {
     response.setIntHeader(CAPACITY_HEADER, capacity);
     response.setIntHeader(REMAINING_HEADER, remaining);
     response.setDateHeader(RESET_HEADER, reset);
+    response.setHeader(LEARNING_HEADER, learning.toString());
   }
 }
diff --git a/gate-web/src/main/groovy/com/netflix/spinnaker/gate/ratelimit/RateLimitingInterceptor.java b/gate-web/src/main/groovy/com/netflix/spinnaker/gate/ratelimit/RateLimitingInterceptor.java
@@ -23,10 +23,12 @@
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.web.servlet.ModelAndView;
 import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;
 
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
+import java.time.ZonedDateTime;
 
 public class RateLimitingInterceptor extends HandlerInterceptorAdapter {
 
@@ -56,26 +58,38 @@ public boolean preHandle(HttpServletRequest request, HttpServletResponse respons
 
     Rate rate = rateLimiter.incrementAndGetRate(principal);
 
+    rate.assignHttpHeaders(response, learning);
+
     if (learning) {
       if (rate.isThrottled()) {
         throttlingCounter.increment();
-        log.warn("Rate limiting principal (principal: {}, learning: true)", principal);
+        log.warn("Rate limiting principal (principal: {}, rateSeconds: {}, capacity: {}, learning: true)", principal, rate.rateSeconds, rate.capacity);
       }
       return true;
     }
 
-    rate.assignHttpHeaders(response);
 
     if (rate.isThrottled()) {
       throttlingCounter.increment();
-      log.warn("Rate limiting principal (principal: {}, rateSeconds: {}, capacity: {})", principal, rate.rateSeconds, rate.capacity);
+      log.warn("Rate limiting principal (principal: {}, rateSeconds: {}, capacity: {}, learning: false)", principal, rate.rateSeconds, rate.capacity);
       response.sendError(429, "Rate capacity exceeded");
       return false;
     }
 
     return true;
   }
 
+  @Override
+  public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
+    // Hystrix et-al can return 429's, which we'll want to intercept to provide a reset header
+    if (response.getStatus() == 429 && !response.getHeaderNames().contains(Rate.RESET_HEADER)) {
+      response.setIntHeader(Rate.CAPACITY_HEADER, -1);
+      response.setIntHeader(Rate.REMAINING_HEADER, 0);
+      response.setDateHeader(Rate.RESET_HEADER, ZonedDateTime.now().plusSeconds(5).toEpochSecond());
+      response.setHeader(Rate.LEARNING_HEADER, "false");
+    }
+  }
+
   private Object getPrincipal() {
     SecurityContext context = SecurityContextHolder.getContext();
     Authentication authentication = context.getAuthentication();