Enable headers to be used in Ratelimit rules

springdog-project/springdog-agent/src/main/resources/templates/content/rate-limit/actions/setting.html

@@ -94,20 +94,28 @@
         label.innerText = isEnabled ? 'enabled' : 'disabled';
       });
 
+      const headerEnableInput = document.createElement("input");
+      headerEnableInput.type = "hidden";
+      headerEnableInput.name = "headerNamesToEnable";
+
       $("#headerOption").select2({
         placeholder: 'Select headers',
         "language": {
           noResults: () => "No headers found."
         },
         allowClear: true
-      })
+      }).on("change", function () {
+        const selectedData = $(this).select2("data");
+        headerEnableInput.value = selectedData.map(item => item.element.dataset.name).join(",");
+      }).trigger('change');
 
       const paramEnableInput = document.createElement("input");
       paramEnableInput.type = "hidden";
       paramEnableInput.name = "parameterNamesToEnable";
 
       const form = document.querySelector("form[name=editRule]");
       form.appendChild(paramEnableInput);
+      form.appendChild(headerEnableInput);
 
       $("#parametersOption").select2({
         placeholder: 'Select parameters',
@@ -263,7 +271,13 @@ <h5 class="card-title">Ban Ruleset combinations</h5>
         </div>
         <div class="mb-3 row">
           <label class="form-label" for="headerOption">Custom Header</label>
-          <select class="form-control" id="headerOption" multiple></select>
+          <select class="form-control" id="headerOption" multiple>
+            <option th:data-name="${header.name}"
+                    th:data-id="${header.id}"
+                    th:each="header : ${endpoint.headers}"
+                    th:selected="${header.isEnabled()}"
+                    th:text="${header.name}"></option>
+          </select>
         </div>
         <div class="form-check mb-3">
           <input class="form-check-input" id="ipBlock"

springdog-project/springdog-autoconfigure/src/main/java/org/easypeelsecurity/springdog/autoconfigure/controller/parser/ControllerParser.java

@@ -29,6 +29,7 @@
 
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.stereotype.Component;
+import org.springframework.web.bind.annotation.RequestHeader;
 import org.springframework.web.bind.annotation.RequestMethod;
 import org.springframework.web.method.HandlerMethod;
 import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
@@ -37,6 +38,7 @@
 import org.easypeelsecurity.springdog.domain.ratelimit.VersionControlService;
 import org.easypeelsecurity.springdog.shared.configuration.SpringdogProperties;
 import org.easypeelsecurity.springdog.shared.dto.EndpointDto;
+import org.easypeelsecurity.springdog.shared.dto.EndpointHeaderDto;
 import org.easypeelsecurity.springdog.shared.dto.EndpointParameterDto;
 import org.easypeelsecurity.springdog.shared.enums.EndpointParameterType;
 import org.easypeelsecurity.springdog.shared.enums.HttpMethod;
@@ -75,19 +77,45 @@ private EndpointDto getEndpointDto(HandlerMethod method, String endPoint, HttpMe
     String methodSignature =
         MethodSignatureParser.parse(method);
 
-    EndpointDto endpoint = new EndpointDto.Builder()
+    EndpointDto endpoint = EndpointDto.builder()
         .methodSignature(methodSignature)
         .path(endPoint)
         .httpMethod(httpMethod)
         .isPatternPath(isPatternPath).build();
+    Set<EndpointParameterDto> parameters = parseRequestParameters(method);
+    Set<EndpointHeaderDto> headers = parseRequestHeaders(method);
+    endpoint.addParameters(parameters);
+    endpoint.addHeaders(headers);
+    return endpoint;
+  }
+
+  private Set<EndpointHeaderDto> parseRequestHeaders(HandlerMethod method) {
+    Set<EndpointHeaderDto> headers = new HashSet<>();
+    Parameter[] parameters = method.getMethod().getParameters();
+    for (Parameter parameter : parameters) {
+      RequestHeader requestHeaderAnnotation = parameter.getAnnotation(RequestHeader.class);
+      if (requestHeaderAnnotation != null) {
+        String headerName = requestHeaderAnnotation.value();
+        if (headerName.isEmpty()) {
+          headerName = parameter.getName();
+        }
+        headers.add(new EndpointHeaderDto(headerName, false));
+      }
+    }
+    return headers;
+  }
+
+  private Set<EndpointParameterDto> parseRequestParameters(HandlerMethod method) {
     Set<EndpointParameterDto> parameters = new HashSet<>();
     try {
       String[] paramNames =
           ParameterNameExtractor.getParameterNames(method.getBeanType(), method.getMethod().getName(),
               method.getMethod().getParameterTypes());
-
       Parameter[] methodParameters = method.getMethod().getParameters();
       for (int i = 0; i < methodParameters.length; i++) {
+        if (methodParameters[i].getAnnotation(RequestHeader.class) != null) {
+          continue;
+        }
         String name =
             paramNames != null && i < paramNames.length ? paramNames[i] : methodParameters[i].getName();
         parameters.add(
@@ -97,8 +125,7 @@ private EndpointDto getEndpointDto(HandlerMethod method, String endPoint, HttpMe
     } catch (IOException | NoSuchMethodException e) {
       logger.error("Error while extracting parameter names.", e);
     }
-    endpoint.addParameters(parameters);
-    return endpoint;
+    return parameters;
   }
 
   /**

springdog-project/springdog-domain/src/main/java/org/easypeelsecurity/springdog/domain/ratelimit/EndpointService.java

@@ -80,7 +80,7 @@ public void updateRule(EndpointDto endpointDto) {
     endpoint.updateRule(endpointDto.getRuleStatus(), endpointDto.isRuleIpBased(),
         endpointDto.isRulePermanentBan(), endpointDto.getRuleRequestLimitCount(),
         endpointDto.getRuleTimeLimitInSeconds(), endpointDto.getRuleBanTimeInSeconds(),
-        endpointDto.getParameterNamesToEnable());
+        endpointDto.getParameterNamesToEnable(), endpointDto.getHeaderNamesToEnable());
 
     RuleCache.changeRuleCached(endpoint.getMethodSignature(), EndpointConverter.toDto(endpoint));
     context.commitChanges();

springdog-project/springdog-domain/src/main/java/org/easypeelsecurity/springdog/domain/ratelimit/EndpointVersionCompare.java

@@ -22,8 +22,10 @@
 import java.util.List;
 
 import org.easypeelsecurity.springdog.domain.ratelimit.model.Endpoint;
+import org.easypeelsecurity.springdog.domain.ratelimit.model.EndpointHeader;
 import org.easypeelsecurity.springdog.domain.ratelimit.model.EndpointParameter;
 import org.easypeelsecurity.springdog.shared.dto.EndpointDto;
+import org.easypeelsecurity.springdog.shared.dto.EndpointHeaderDto;
 import org.easypeelsecurity.springdog.shared.dto.EndpointParameterDto;
 
 /**
@@ -91,6 +93,12 @@ private String generateFullHashFromEntity(List<Endpoint> endpoints) {
           .forEach(p -> hash.append(p.getName())
               .append(PROPERTY_DELIMITER)
               .append(p.getType()));
+
+      endpoint.getEndpointHeaders()
+          .stream()
+          .sorted(Comparator.comparing(EndpointHeader::getName))
+          .forEach(p -> hash.append(p.getName())
+              .append(PROPERTY_DELIMITER));
     }
 
     return generateSHA256Hex(hash.toString());
@@ -123,6 +131,12 @@ public String generateFullHashFromDto(List<EndpointDto> endpoints) {
           .forEach(p -> hash.append(p.getName())
               .append(PROPERTY_DELIMITER)
               .append(p.getType()));
+
+      endpoint.getHeaders()
+          .stream()
+          .sorted(Comparator.comparing(EndpointHeaderDto::getName))
+          .forEach(p -> hash.append(p.getName())
+              .append(PROPERTY_DELIMITER));
     }
 
     return generateSHA256Hex(hash.toString());

springdog-project/springdog-domain/src/main/java/org/easypeelsecurity/springdog/domain/ratelimit/converter/EndpointConverter.java

@@ -21,8 +21,10 @@
 import java.util.stream.Collectors;
 
 import org.easypeelsecurity.springdog.domain.ratelimit.model.Endpoint;
+import org.easypeelsecurity.springdog.domain.ratelimit.model.EndpointHeader;
 import org.easypeelsecurity.springdog.domain.ratelimit.model.EndpointParameter;
 import org.easypeelsecurity.springdog.shared.dto.EndpointDto;
+import org.easypeelsecurity.springdog.shared.dto.EndpointHeaderDto;
 import org.easypeelsecurity.springdog.shared.dto.EndpointParameterDto;
 import org.easypeelsecurity.springdog.shared.enums.HttpMethod;
 import org.easypeelsecurity.springdog.shared.enums.RuleStatus;
@@ -63,9 +65,21 @@ public static Endpoint toEntity(ObjectContext context, EndpointDto endpointDto)
       parameter.setEndpoint(endpoint);
       endpoint.addToEndpointParameters(parameter);
     }
+
+    for (EndpointHeaderDto header : endpointDto.getHeaders()) {
+      EndpointHeader headerItem = toEntity(context, header);
+      headerItem.setEndpoint(endpoint);
+      endpoint.addToEndpointHeaders(headerItem);
+    }
     return endpoint;
   }
 
+  private static EndpointHeader toEntity(ObjectContext context, EndpointHeaderDto header) {
+    EndpointHeader endpointHeader = context.newObject(EndpointHeader.class);
+    endpointHeader.setName(header.getName());
+    return endpointHeader;
+  }
+
   private static EndpointParameter toEntity(ObjectContext context, EndpointParameterDto endpointParameterDto) {
     EndpointParameter endpointParameter = context.newObject(EndpointParameter.class);
     endpointParameter.setName(endpointParameterDto.getName());
@@ -99,12 +113,13 @@ private static Set<EndpointParameterDto> toDto(Set<EndpointParameter> endpointPa
    */
   public static EndpointDto toDto(Endpoint endpointEntity) {
     Assert.notNull(endpointEntity, "Endpoint must not be null");
-    return new EndpointDto.Builder()
+    return EndpointDto.builder()
         .id(endpointEntity.getId())
         .path(endpointEntity.getPath())
         .methodSignature(endpointEntity.getMethodSignature())
         .httpMethod(HttpMethod.valueOf(endpointEntity.getHttpMethod()))
         .parameters(toParameterDto(endpointEntity.getEndpointParameters()))
+        .headers(toHeaderDto(endpointEntity.getEndpointHeaders()))
         .isPatternPath(endpointEntity.isIsPatternPath())
         .ruleStatus(RuleStatus.of(endpointEntity.getRuleStatus()))
         .ruleIpBased(endpointEntity.isRuleIpBased())
@@ -115,6 +130,16 @@ public static EndpointDto toDto(Endpoint endpointEntity) {
         .build();
   }
 
+  private static Set<EndpointHeaderDto> toHeaderDto(List<EndpointHeader> endpointHeaders) {
+    if (endpointHeaders == null) {
+      return Set.of();
+    }
+    return endpointHeaders
+        .stream()
+        .map(header -> new EndpointHeaderDto(header.getName(), header.isEnabled()))
+        .collect(Collectors.toSet());
+  }
+
   private static Set<EndpointParameterDto> toParameterDto(List<EndpointParameter> endpointParameters) {
     if (endpointParameters == null) {
       return Set.of();

springdog-project/springdog-domain/src/main/java/org/easypeelsecurity/springdog/domain/ratelimit/model/Endpoint.java

@@ -17,9 +17,9 @@
 package org.easypeelsecurity.springdog.domain.ratelimit.model;
 
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 import org.easypeelsecurity.springdog.domain.ratelimit.model.auto._Endpoint;
 import org.easypeelsecurity.springdog.shared.enums.RuleStatus;
@@ -39,6 +39,7 @@ public Endpoint() {
     setRulePermanentBan(false);
     setRuleIpBased(false);
     this.endpointParameters = new ArrayList<>();
+    this.endpointHeaders = new ArrayList<>();
   }
 
   @Override
@@ -54,6 +55,19 @@ public void removeFromEndpointParameters(EndpointParameter obj) {
     super.removeFromEndpointParameters(obj);
   }
 
+  @Override
+  public void addToEndpointHeaders(EndpointHeader obj) {
+    Assert.notNull(obj, "EndpointHeader must not be null");
+    obj.setEndpoint(this);
+    super.addToEndpointHeaders(obj);
+  }
+
+  @Override
+  public void removeFromEndpointHeaders(EndpointHeader obj) {
+    Assert.notNull(obj, "EndpointHeader must not be null");
+    super.removeFromEndpointHeaders(obj);
+  }
+
   /**
    * Update rule.
    *
@@ -66,37 +80,62 @@ public void removeFromEndpointParameters(EndpointParameter obj) {
    */
   public void updateRule(RuleStatus ruleStatus, boolean ruleIpBased, boolean rulePermanentBan,
       int ruleRequestLimitCount, int ruleTimeLimitInSeconds, int ruleBanTimeInSeconds,
-      Set<String> enableParamNames) {
+      Set<String> enableParamNames, Set<String> enableHeaderNames) {
     Assert.isNotEqual(ruleStatus, RuleStatus.NOT_CONFIGURED, "Couldn't change ruleStatus to NOT_CONFIGURED");
     Assert.isTrue(enableParamNames != null, "enableParamNames must not be null");
+    Assert.isTrue(enableHeaderNames != null, "enableHeaderNames must not be null");
+
+    ruleValidate(
+        ruleStatus.name(), ruleIpBased, rulePermanentBan, ruleRequestLimitCount,
+        ruleTimeLimitInSeconds, ruleBanTimeInSeconds, enableParamNames, enableHeaderNames,
+        this.getEndpointParameters(), this.getEndpointHeaders());
+
     setRuleStatus(ruleStatus.name());
     setRuleIpBased(ruleIpBased);
     setRulePermanentBan(rulePermanentBan);
     setRuleRequestLimitCount(ruleRequestLimitCount);
     setRuleTimeLimitInSeconds(ruleTimeLimitInSeconds);
     setRuleBanTimeInSeconds(ruleBanTimeInSeconds);
     getEndpointParameters().forEach(param -> param.setEnabled(enableParamNames.contains(param.getName())));
-    ruleValidate();
+    getEndpointHeaders().forEach(header -> header.setEnabled(enableHeaderNames.contains(header.getName())));
   }
 
-  private void ruleValidate() {
-    if (RuleStatus.ACTIVE.name().equals(this.ruleStatus)) {
-      if (this.ruleRequestLimitCount <= 0) {
+  private void ruleValidate(String ruleStatus, boolean ruleIpBased, boolean rulePermanentBan,
+      int ruleRequestLimitCount, int ruleTimeLimitInSeconds, int ruleBanTimeInSeconds,
+      Set<String> enableParamNames, Set<String> enableHeaderNames,
+      List<EndpointParameter> endpointParameters,
+      List<EndpointHeader> endpointHeaders) {
+    if (RuleStatus.ACTIVE.name().equals(ruleStatus)) {
+      if (ruleRequestLimitCount <= 0) {
         throw new IllegalArgumentException("Request limit count must be greater than 0");
       }
-      if (this.ruleTimeLimitInSeconds <= 0) {
+      if (ruleTimeLimitInSeconds <= 0) {
         throw new IllegalArgumentException("Time limit must be greater than 0");
       }
-      if (this.ruleBanTimeInSeconds <= 0 && !this.rulePermanentBan) {
+      if (ruleBanTimeInSeconds <= 0 && !rulePermanentBan) {
         throw new IllegalArgumentException("Ban time must be greater than 0");
       }
-      if (!this.ruleIpBased && this.getEndpointParameters().stream().filter(EndpointParameter::isEnabled)
-          .collect(Collectors.toSet()).isEmpty()) {
-        throw new IllegalArgumentException("At least one combinations must be enabled");
+      if (!ruleIpBased && (isParamNotActivated(enableParamNames, endpointParameters)
+          && isHeaderNotActivated(enableHeaderNames, endpointHeaders))) {
+        throw new IllegalArgumentException("At least one combination must be enabled");
       }
     }
   }
 
+  private static boolean isHeaderNotActivated(Set<String> enableHeaderNames,
+      List<EndpointHeader> endpointHeaders) {
+    return endpointHeaders.stream()
+        .filter(header -> enableHeaderNames.contains(header.getName()))
+        .count() == 0;
+  }
+
+  private static boolean isParamNotActivated(Set<String> enableParamNames,
+      List<EndpointParameter> endpointParameters) {
+    return endpointParameters.stream()
+        .filter(param -> enableParamNames.contains(param.getName()))
+        .count() == 0;
+  }
+
   @Override
   public boolean equals(Object obj) {
     if (this == obj) {

springdog-project/springdog-domain/src/main/java/org/easypeelsecurity/springdog/domain/ratelimit/model/EndpointHeader.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.easypeelsecurity.springdog.domain.ratelimit.model;
+
+import java.util.Objects;
+
+import org.easypeelsecurity.springdog.domain.ratelimit.model.auto._EndpointHeader;
+
+@SuppressWarnings("all")
+public class EndpointHeader extends _EndpointHeader {
+
+  private static final long serialVersionUID = 1L;
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
+    if (obj == null || getClass() != obj.getClass()) {
+      return false;
+    }
+    EndpointParameter that = (EndpointParameter) obj;
+    return this.getName().equals(that.getName()) && Objects.equals(this.getEndpoint(), that.getEndpoint());
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(this.getName(), this.getEndpoint());
+  }
+}