feat: Limit API usage

.github/workflows/integration-tests.yml

@@ -67,6 +67,8 @@ jobs:
                     ZWE_CONFIGS_APIML_SERVICE_ADDITIONALREGISTRATION_0_DISCOVERYSERVICEURLS: https://discovery-service-2:10011/eureka
                     SERVER_MAX_HTTP_REQUEST_HEADER_SIZE: 16348
                     SERVER_WEBSOCKET_REQUESTBUFFERSIZE: 16348
+                    APIML_GATEWAY_ROUTING_SERVICESTOLIMITREQUESTRATE: discoverableclient
+                    APIML_GATEWAY_ROUTING_COOKIENAMEFORRATELIMIT: apimlAuthenticationToken
             zaas-service:
                 image: ghcr.io/balhar-jakub/zaas-service:${{ github.run_id }}-${{ github.run_number }}
                 env:

gateway-package/src/main/resources/bin/start.sh

@@ -318,6 +318,11 @@ _BPX_JOBNAME=${ZWE_zowe_job_prefix}${GATEWAY_CODE} ${JAVA_BIN_DIR}java \
     -Dapiml.gateway.cachePeriodSec=${ZWE_configs_apiml_gateway_registry_cachePeriodSec:-120} \
     -Dapiml.gateway.registry.enabled=${ZWE_configs_apiml_gateway_registry_enabled:-false} \
     -Dapiml.gateway.maxSimultaneousRequests=${ZWE_configs_gateway_registry_maxSimultaneousRequests:-20} \
+    -Dapiml.gateway.rateLimiterCapacity=${ZWE_configs_apiml_gateway_routing_rateLimiterCapacity:-20} \
+    -Dapiml.gateway.rateLimiterTokens=${ZWE_configs_apiml_gateway_routing_rateLimiterTokens:-20} \
+    -Dapiml.gateway.rateLimiterRefillDuration=${ZWE_configs_apiml_gateway_routing_rateLimiterRefillDuration:-1} \
+    -Dapiml.gateway.servicesToLimitRequestRate=${ZWE_configs_apiml_gateway_routing_servicesToLimitRequestRate:-} \
+    -Dapiml.gateway.cookieNameForRateLimit=${cookieName:-apimlAuthenticationToken} \
     -Dapiml.gateway.registry.metadata-key-allow-list=${ZWE_configs_gateway_registry_metadataKeyAllowList:-} \
     -Dapiml.gateway.refresh-interval-ms=${ZWE_configs_gateway_registry_refreshIntervalMs:-30000} \
     -Dserver.address=${ZWE_configs_zowe_network_server_listenAddresses_0:-${ZWE_zowe_network_server_listenAddresses_0:-"0.0.0.0"}} \

gateway-service/build.gradle

@@ -80,6 +80,7 @@ dependencies {
     implementation libs.nimbus.jose.jwt
     implementation libs.bcpkix
     implementation libs.caffeine
+    implementation libs.bucket4j.core
 
     implementation libs.swagger2.parser
     implementation libs.swagger3.parser

gateway-service/src/main/java/org/zowe/apiml/gateway/config/RoutingConfig.java

@@ -62,6 +62,10 @@ public List<FilterDefinition> filters() {
         retryFilter.addArg("series", "");
         filters.add(retryFilter);
 
+        FilterDefinition rateLimiterFilter = new FilterDefinition();
+        rateLimiterFilter.setName("InMemoryRateLimiterFilterFactory");
+        filters.add(rateLimiterFilter);
+
         for (String headerName : ignoredHeadersWhenCorsEnabled.split(",")) {
             FilterDefinition removeHeaders = new FilterDefinition();
             removeHeaders.setName("RemoveRequestHeader");

gateway-service/src/main/java/org/zowe/apiml/gateway/filters/InMemoryRateLimiter.java

@@ -0,0 +1,91 @@
+/*
+ * This program and the accompanying materials are made available under the terms of the
+ * Eclipse Public License v2.0 which accompanies this distribution, and is available at
+ * https://www.eclipse.org/legal/epl-v20.html
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ * Copyright Contributors to the Zowe Project.
+ */
+
+package org.zowe.apiml.gateway.filters;
+
+import io.github.bucket4j.Bandwidth;
+import io.github.bucket4j.Bucket;
+import lombok.Getter;
+import lombok.Setter;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
+import org.springframework.stereotype.Component;
+import reactor.core.publisher.Mono;
+
+import java.time.Duration;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+@Component
+public class InMemoryRateLimiter implements RateLimiter<InMemoryRateLimiter.Config> {
+
+    private final Map<String, Bucket> cache = new ConcurrentHashMap<>();
+    @Value("${apiml.gateway.routing.rateLimiterCapacity:20}")
+    int capacity;
+    @Value("${apiml.gateway.routing.rateLimiterTokens:20}")
+    int tokens;
+    @Value("${apiml.gateway.routing.rateLimiterRefillDuration:1}")
+    Integer refillDuration;
+
+    @Override
+    public Mono<Response> isAllowed(String routeId, String id) {
+        Bucket bucket = cache.computeIfAbsent(id, this::newBucket);
+        if (bucket.tryConsume(1)) {
+            return Mono.just(new Response(true, getHeaders(bucket)));
+        } else {
+            return Mono.just(new Response(false, getHeaders(bucket)));
+        }
+    }
+
+    private Bucket newBucket(String id) {
+        Bandwidth limit = Bandwidth.builder().capacity(capacity).refillGreedy(tokens, Duration.ofMinutes(refillDuration)).build();
+        return Bucket.builder().addLimit(limit).build();
+    }
+
+    private Map<String, String> getHeaders(Bucket bucket) {
+        Map<String, String> headers = new ConcurrentHashMap<>();
+        headers.put("X-RateLimit-Remaining", String.valueOf(bucket.getAvailableTokens()));
+        return headers;
+    }
+
+    @Override
+    public Map<String, Config> getConfig() {
+        Config defaultConfig = new Config();
+        defaultConfig.setCapacity(capacity);
+        defaultConfig.setTokens(tokens);
+        defaultConfig.setRefillDuration(refillDuration);
+
+        Map<String, Config> configMap = new ConcurrentHashMap<>();
+        configMap.put("default", defaultConfig);
+        return configMap;
+    }
+
+    @Override
+    public Class<Config> getConfigClass() {
+        return Config.class;
+    }
+
+    @Override
+    public Config newConfig() {
+        Config config = new Config();
+        config.setCapacity(capacity);
+        config.setTokens(tokens);
+        config.setRefillDuration(refillDuration);
+        return config;
+    }
+
+    @Setter
+    @Getter
+    public static class Config {
+        private int capacity;
+        private int tokens;
+        private int refillDuration;
+    }
+}

gateway-service/src/main/java/org/zowe/apiml/gateway/filters/InMemoryRateLimiterFilterFactory.java

@@ -0,0 +1,95 @@
+/*
+ * This program and the accompanying materials are made available under the terms of the
+ * Eclipse Public License v2.0 which accompanies this distribution, and is available at
+ * https://www.eclipse.org/legal/epl-v20.html
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ * Copyright Contributors to the Zowe Project.
+ */
+
+package org.zowe.apiml.gateway.filters;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import lombok.Getter;
+import lombok.Setter;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.cloud.gateway.filter.GatewayFilter;
+import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.server.PathContainer;
+import org.springframework.stereotype.Component;
+import org.zowe.apiml.message.core.Message;
+import org.zowe.apiml.message.core.MessageService;
+import org.zowe.apiml.message.log.ApimlLogger;
+import org.zowe.apiml.product.logging.annotations.InjectApimlLogger;
+import reactor.core.publisher.Mono;
+
+import java.util.List;
+
+@Component
+public class InMemoryRateLimiterFilterFactory extends AbstractGatewayFilterFactory<InMemoryRateLimiterFilterFactory.Config> {
+
+    @InjectApimlLogger
+    private ApimlLogger apimlLog = ApimlLogger.empty();
+
+    private final InMemoryRateLimiter rateLimiter;
+
+    private final KeyResolver keyResolver;
+
+    @Value("${apiml.gateway.routing.servicesToLimitRequestRate:-}")
+    List<String> serviceIds;
+
+    private final ObjectMapper mapper;
+
+    private final MessageService messageService;
+
+    public InMemoryRateLimiterFilterFactory(InMemoryRateLimiter rateLimiter, KeyResolver keyResolver, ObjectMapper mapper, MessageService messageService) {
+        super(Config.class);
+        this.rateLimiter = rateLimiter;
+        this.keyResolver = keyResolver;
+        this.mapper = mapper;
+        this.messageService = messageService;
+    }
+
+    @Override
+    public GatewayFilter apply(Config config) {
+        return (exchange, chain) -> {
+            List<PathContainer.Element> pathElements = exchange.getRequest().getPath().elements();
+            String requestPath = (!pathElements.isEmpty() && pathElements.size() > 1) ? pathElements.get(1).value() : null;
+            if (requestPath == null || !serviceIds.contains(requestPath)) {
+                return chain.filter(exchange);
+            }
+            return keyResolver.resolve(exchange).flatMap(key -> {
+                if (key.isEmpty()) {
+                    return chain.filter(exchange);
+                }
+                return rateLimiter.isAllowed(config.getRouteId(), key).flatMap(response -> {
+                    if (response.isAllowed()) {
+                        return chain.filter(exchange);
+                    } else {
+                        apimlLog.log("org.zowe.apiml.gateway.connectionsLimitApproached", "Connections limit exceeded for service '{}'", requestPath);
+                        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
+                        Message message = messageService.createMessage("org.zowe.apiml.gateway.connectionsLimitApproached", "Connections limit exceeded for service '{}'", requestPath);
+                        try {
+                            return exchange.getResponse().writeWith(Mono.just(exchange.getResponse().bufferFactory().wrap(mapper.writeValueAsBytes(message.mapToView()))));
+                        } catch (JsonProcessingException e) {
+                            apimlLog.log("org.zowe.apiml.security.errorWritingResponse", e.getMessage());
+                            return Mono.error(e);
+                        }
+                    }
+                });
+            });
+        };
+    }
+
+    @Getter
+    @Setter
+    public static class Config {
+        private String routeId;
+        private Integer capacity;
+        private Integer tokens;
+        private Integer refillIntervalSeconds;
+    }
+}

gateway-service/src/main/java/org/zowe/apiml/gateway/filters/KeyResolver.java

@@ -0,0 +1,35 @@
+/*
+ * This program and the accompanying materials are made available under the terms of the
+ * Eclipse Public License v2.0 which accompanies this distribution, and is available at
+ * https://www.eclipse.org/legal/epl-v20.html
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ * Copyright Contributors to the Zowe Project.
+ */
+
+package org.zowe.apiml.gateway.filters;
+
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.http.HttpCookie;
+import org.springframework.stereotype.Component;
+import reactor.core.publisher.Mono;
+
+import java.util.Collections;
+
+@Component
+public class KeyResolver implements org.springframework.cloud.gateway.filter.ratelimit.KeyResolver {
+
+    @Value("${apiml.gateway.routing.cookieNameForRateLimit:apimlAuthenticationToken}")
+    private String cookieName;
+
+    @Override
+    public Mono<String> resolve(org.springframework.web.server.ServerWebExchange exchange) {
+        return Mono.just(exchange.getRequest().getCookies().getOrDefault(cookieName, Collections.emptyList())
+            .stream()
+            .findFirst()
+            .map(HttpCookie::getValue)
+            .orElse("")
+        );
+    }
+}

gateway-service/src/main/resources/gateway-log-messages.yml

@@ -118,3 +118,10 @@ messages:
       text: "Cannot receive information about services on API Gateway with apimlId '%s' because: %s"
       reason: "Cannot connect to the Gateway service."
       action: "Make sure that the external Gateway service is running and the truststore of the both Gateways contain the corresponding certificate."
+
+    - key: org.zowe.apiml.gateway.connectionsLimitApproached
+      number: ZWESG429
+      type: ERROR
+      text: "Request was denied access."
+      reason: "Connections limit exceeded."
+      action: "Wait for the number of active connections to decrease before retrying your request."