feat: Limit API usage

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

@@ -318,6 +318,9 @@ _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_gateway_registry_rateLimiterCapacity:-3} \
+    -Dapiml.gateway.rateLimiterTokens=${ZWE_configs_gateway_registry_rateLimiterTokens:-3} \
+    -Dapiml.gateway.rateLimiterRefillDuration=${ZWE_configs_gateway_registry_rateLimiterRefillDuration:-1} \
     -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

@@ -65,6 +65,7 @@ dependencies {
     api project(':apiml-tomcat-common')
     api project(':apiml-security-common')
 
+    implementation group: 'com.bucket4j', name: 'bucket4j-core', version: '8.7.0'
     implementation libs.spring.boot.starter.security
     implementation libs.spring.cloud.circuit.breaker
     implementation libs.spring.cloud.starter.eureka.client

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.rateLimiterCapacity:20}")
+    int capacity;
+    @Value("${apiml.gateway.rateLimiterTokens:1}")
+    int tokens;
+    @Value("${apiml.gateway.rateLimiterRefillDuration:1}")
+    Long 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 Long refillDuration;
+    }
+}

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

@@ -0,0 +1,82 @@
+/*
+ * 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 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.MediaType;
+import org.springframework.stereotype.Component;
+import org.zowe.apiml.message.log.ApimlLogger;
+import org.zowe.apiml.product.logging.annotations.InjectApimlLogger;
+import reactor.core.publisher.Mono;
+
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+
+@Component
+public class InMemoryRateLimiterFilterFactory extends AbstractGatewayFilterFactory<InMemoryRateLimiterFilterFactory.Config> {
+
+    @InjectApimlLogger
+    private final ApimlLogger apimlLog = ApimlLogger.empty();
+    private final InMemoryRateLimiter rateLimiter;
+    private final KeyResolver keyResolver;
+    @Value(value = "${apiml.routing.servicesToLimitRequestRate}")
+    List<String> serviceIds;
+
+    public InMemoryRateLimiterFilterFactory(InMemoryRateLimiter rateLimiter, KeyResolver keyResolver) {
+        super(Config.class);
+        this.rateLimiter = rateLimiter;
+        this.keyResolver = keyResolver;
+    }
+
+    @Override
+    public GatewayFilter apply(Config config) {
+        return (exchange, chain) -> {
+            String requestPath = exchange.getRequest().getPath().elements().get(1).value();
+            if (serviceIds.contains(requestPath)) {
+                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);
+                            exchange.getResponse().getHeaders().setContentType(MediaType.TEXT_HTML);
+                            String errorHtml = "<html><body><h1>429 Too Many Requests</h1>" +
+                                "<p>The connection limit for the service '" + requestPath + "' has been exceeded. Please try again later.</p>" +
+                                "</body></html>";
+                            byte[] bytes = errorHtml.getBytes(StandardCharsets.UTF_8);
+                            return exchange.getResponse().writeWith(Mono.just(exchange.getResponse().bufferFactory().wrap(bytes)));
+                        }
+                    });
+                });
+            } else {
+                return chain.filter(exchange);
+            }
+        };
+    }
+
+    @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,37 @@
+/*
+ * 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.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 {
+
+    private final String cookieName;
+
+    public KeyResolver() {
+        this.cookieName = "apimlAuthenticationToken";
+    }
+
+    @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/java/org/zowe/apiml/gateway/service/RouteLocator.java

@@ -39,7 +39,8 @@
 import static org.zowe.apiml.constants.EurekaMetadataDefinition.*;
 
 @Service
-public class RouteLocator implements RouteDefinitionLocator {
+public class
+RouteLocator implements RouteDefinitionLocator {
 
     private static final EurekaMetadataParser metadataParser = new EurekaMetadataParser();
 

gateway-service/src/main/resources/application.yml

@@ -85,6 +85,7 @@ apiml:
     routing:
         instanceIdHeader: false
         ignoredServices: discovery,zaas # to disable routing to the Discovery and ZAAS service
+        servicesToLimitRequestRate: discoverableclient
     service:
         apimlId: apiml1
         corsEnabled: true

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."

gateway-service/src/test/java/org/zowe/apiml/gateway/filters/InMemoryRateLimiterFilterFactoryTest.java

@@ -0,0 +1,106 @@
+/*
+ * 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.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.springframework.cloud.gateway.filter.GatewayFilterChain;
+import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
+import org.springframework.mock.web.server.MockServerWebExchange;
+import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.*;
+
+public class InMemoryRateLimiterFilterFactoryTest {
+
+    private InMemoryRateLimiter rateLimiter;
+    private KeyResolver keyResolver;
+    private InMemoryRateLimiterFilterFactory filterFactory;
+    private ServerWebExchange exchange;
+    private GatewayFilterChain chain;
+    private String serviceId;
+    private MockServerHttpRequest request;
+    private InMemoryRateLimiterFilterFactory.Config config;
+
+    @BeforeEach
+    public void setUp() {
+        serviceId = "testService";
+        rateLimiter = mock(InMemoryRateLimiter.class);
+        keyResolver = mock(KeyResolver.class);
+        filterFactory = new InMemoryRateLimiterFilterFactory(rateLimiter, keyResolver);
+        filterFactory.serviceIds = List.of(serviceId);
+        request = MockServerHttpRequest.get("/" + serviceId).build();
+        exchange = MockServerWebExchange.from(request);
+        chain = mock(GatewayFilterChain.class);
+        config = mock(InMemoryRateLimiterFilterFactory.Config.class);
+        when(config.getRouteId()).thenReturn("testRoute");
+    }
+
+    @Test
+    public void apply_shouldAllowRequest_whenTokensAreAvailable() {
+        when(keyResolver.resolve(exchange)).thenReturn(Mono.just("testKey"));
+        when(rateLimiter.isAllowed(anyString(), anyString())).thenReturn(Mono.just(new RateLimiter.Response(true, Map.of())));
+        when(chain.filter(any(ServerWebExchange.class))).thenReturn(Mono.empty());
+
+        StepVerifier.create(filterFactory.apply(config).filter(exchange, chain))
+            .expectComplete()
+            .verify();
+        verify(chain, times(1)).filter(exchange);
+    }
+
+    @Test
+    public void apply_shouldReturn429_whenTokensAreExhausted() {
+        when(keyResolver.resolve(exchange)).thenReturn(Mono.just("testKey"));
+        when(rateLimiter.isAllowed(anyString(), anyString())).thenReturn(Mono.just(new InMemoryRateLimiter.Response(false, Map.of())));
+
+        StepVerifier.create(filterFactory.apply(config).filter(exchange, chain))
+            .expectComplete()
+            .verify();
+        ServerHttpResponse response = exchange.getResponse();
+        assertEquals(HttpStatus.TOO_MANY_REQUESTS, response.getStatusCode());
+    }
+
+    @Test
+    public void apply_shouldAllowRequest_whenKeyIsNull() {
+        when(keyResolver.resolve(exchange)).thenReturn(Mono.just(""));
+        when(chain.filter(any(ServerWebExchange.class))).thenReturn(Mono.empty());
+
+        StepVerifier.create(filterFactory.apply(config).filter(exchange, chain))
+            .expectComplete()
+            .verify();
+        verify(chain, times(1)).filter(exchange);
+    }
+
+    @Test
+    public void apply_shouldAllowRequest_whenServiceIdDoesNotMatch() {
+        String nonMatchingId = "nonMatchingId";
+        when(keyResolver.resolve(exchange)).thenReturn(Mono.just("testKey"));
+        request = MockServerHttpRequest.get("/" + nonMatchingId).build();
+        exchange = MockServerWebExchange.from(request);
+        when(chain.filter(any(ServerWebExchange.class))).thenReturn(Mono.empty());
+
+        StepVerifier.create(filterFactory.apply(config).filter(exchange, chain))
+            .expectComplete()
+            .verify();
+        verify(chain, times(1)).filter(exchange);
+    }
+
+}
+