com.founder.core.web.RateLimitExceptionHandler Maven / Gradle / Ivy
package com.founder.core.web;
import com.founder.core.exception.RateLimitException;
import com.founder.core.model.Mode;
import com.founder.core.starter.properties.RateLimiterProperties;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.ResponseBody;
@ControllerAdvice
@ConditionalOnProperty(prefix = RateLimiterProperties.PREFIX, name = "exception-handler.enable", havingValue = "true", matchIfMissing = true)
public class RateLimitExceptionHandler {
private final RateLimiterProperties limiterProperties;
public static final String REMAINING_HEADER = "X-RateLimit-Remaining";
public RateLimitExceptionHandler(RateLimiterProperties limiterProperties) {
this.limiterProperties = limiterProperties;
}
@ExceptionHandler(value = RateLimitException.class)
@ResponseBody
public ResponseEntity exceptionHandler(RateLimitException e) {
HttpHeaders headers = new HttpHeaders();
if (e.getMode().equals(Mode.TIME_WINDOW)){
headers.add(HttpHeaders.RETRY_AFTER, String.valueOf(e.getExtra()));
}else {
headers.add(REMAINING_HEADER, String.valueOf(e.getExtra()));
}
return ResponseEntity.status(limiterProperties.getStatusCode())
.headers(headers)
.contentType(MediaType.APPLICATION_JSON)
.body(limiterProperties.getResponseBody());
}
}