Broken Promise response header propagation

.github/workflows/maven.yml

@@ -23,21 +23,21 @@ jobs:
         distribution: temurin
         java-version: 17
     - name: Cache Maven Dependencies (~/.m2/repository)
-      uses: actions/cache@v3.0.11
+      uses: actions/cache@v3.2.6
       with:
         path: ~/.m2/repository
         key: ${{ runner.os }}-maven-${{ hashFiles('**/pom.xml') }}
         restore-keys: |
           ${{ runner.os }}-maven-:
     - name: Cache NPM Dependencies (core/com.b2international.snowowl.core.rest/snow-owl-api-docs/node_modules)
-      uses: actions/cache@v3.0.11
+      uses: actions/cache@v3.2.6
       with:
         path: core/com.b2international.snowowl.core.rest/snow-owl-api-docs/node_modules
         key: ${{ runner.os }}-npm-${{ hashFiles('core/com.b2international.snowowl.core.rest/snow-owl-api-docs/package-lock.json') }}
         restore-keys: |
           ${{ runner.os }}-npm-:
     - name: Setup Maven settings.xml
-      uses: whelk-io/maven-settings-xml-action@v14
+      uses: whelk-io/maven-settings-xml-action@v21
       with:
         servers: '[{ "id": "b2i-releases", "username": "${env.NEXUS_DEPLOYMENT_USER}", "password": "${env.NEXUS_DEPLOYMENT_PASS}" }, { "id": "b2i-snapshots", "username": "${env.NEXUS_DEPLOYMENT_USER}", "password": "${env.NEXUS_DEPLOYMENT_PASS}" }, { "id": "nexus_deployment", "username": "${env.NEXUS_DEPLOYMENT_USER}", "password": "${env.NEXUS_DEPLOYMENT_PASS}" }]'
     # Initializes the CodeQL tools for scanning.

core/com.b2international.snowowl.core.rest.tests/src/com/b2international/snowowl/core/rest/AllSnowOwlApiTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2011-2021 B2i Healthcare Pte Ltd, http://b2i.sg
+ * Copyright 2011-2023 B2i Healthcare Pte Ltd, http://b2i.sg
  * 
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@
 import com.b2international.snowowl.core.rest.auth.BasicAuthenticationTest;
 import com.b2international.snowowl.core.rest.bundle.BundleRestApiTest;
 import com.b2international.snowowl.core.rest.codesystem.CodeSystemApiTest;
+import com.b2international.snowowl.core.rest.rate.RateLimitTest;
 import com.b2international.snowowl.core.rest.resource.ResourceApiTest;
 import com.b2international.snowowl.test.commons.BundleStartRule;
 import com.b2international.snowowl.test.commons.SnowOwlAppRule;
@@ -37,10 +38,11 @@
 @SuiteClasses({ 
 	BasicAuthenticationTest.class,
 	AuthorizationTest.class,
+	RateLimitTest.class,
 	CodeSystemApiTest.class,
 	ResourceApiTest.class,
 	BundleApiTest.class,
-	BundleRestApiTest.class
+	BundleRestApiTest.class,
 })
 public class AllSnowOwlApiTests {
 

core/com.b2international.snowowl.core.rest.tests/src/com/b2international/snowowl/core/rest/rate/RateLimitTest.java

@@ -0,0 +1,142 @@
+/*
+ * Copyright 2023 B2i Healthcare Pte Ltd, http://b2i.sg
+ * 
+ * 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 com.b2international.snowowl.core.rest.rate;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.hamcrest.CoreMatchers.anyOf;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.junit.Assert.fail;
+
+import java.time.LocalDateTime;
+
+import org.elasticsearch.core.Map;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.FixMethodOrder;
+import org.junit.Test;
+import org.junit.runners.MethodSorters;
+
+import com.b2international.commons.exceptions.TooManyRequestsException;
+import com.b2international.snowowl.core.ApplicationContext;
+import com.b2international.snowowl.core.Resources;
+import com.b2international.snowowl.core.events.util.Response;
+import com.b2international.snowowl.core.rate.ApiConfiguration;
+import com.b2international.snowowl.core.rate.ApiPlugin;
+import com.b2international.snowowl.core.rate.RateLimitConsumption;
+import com.b2international.snowowl.core.rate.RateLimiter;
+import com.b2international.snowowl.core.request.ResourceRequests;
+import com.b2international.snowowl.core.setup.Environment;
+import com.b2international.snowowl.test.commons.Services;
+import com.b2international.snowowl.test.commons.rest.ResourceApiAssert;
+import com.b2international.snowowl.test.commons.rest.RestExtensions;
+
+/**
+ * @since 8.10
+ */
+@FixMethodOrder(MethodSorters.NAME_ASCENDING)
+public class RateLimitTest {
+
+	private RateLimiter rateLimiter;
+
+	@Before
+	public void setup() throws Exception {
+		// inject rate limit feature into the system until we run the rate limit tests
+		Environment env = ApplicationContext.getServiceForClass(Environment.class);
+		ApiConfiguration apiConfig = new ApiConfiguration();
+		// this means that the user is able to perform 2 requests at a time, with one second refill rate (default value, but just in case fix it here as well)
+		apiConfig.setOverdraft(2L);
+		apiConfig.setRefillRate(1L);
+		new ApiPlugin().initRateLimiter(env, apiConfig);
+		this.rateLimiter = env.service(RateLimiter.class);
+
+	}
+
+	@After
+	public void after() throws Exception {
+		// revert back to noop rate limiter
+		ApplicationContext.getServiceForClass(Environment.class).services().registerService(RateLimiter.class, RateLimiter.NOOP);
+		rateLimiter = null;
+	}
+
+	@Test
+	public void rateLimitServiceTest() throws Exception {
+		// perform three consumptions at the same time
+		RateLimitConsumption firstRequest = rateLimiter.consume(RestExtensions.USER);
+		RateLimitConsumption secondRequest = rateLimiter.consume(RestExtensions.USER);
+		RateLimitConsumption thirdRequest = rateLimiter.consume(RestExtensions.USER);
+		// first two should succeed, third should fail to consume token
+		assertThat(firstRequest.isConsumed()).isTrue();
+		assertThat(secondRequest.isConsumed()).isTrue();
+		assertThat(thirdRequest.isConsumed()).isFalse();
+	}
+
+	@Test
+	public void rateLimitTest_JavaAPI() throws Exception {
+		Response<Resources> response = ResourceRequests.prepareSearch()
+			.setLimit(0)
+			.buildAsync()
+			.execute(Services.bus())
+			.getSyncResponse();
+
+		assertThat(response.getBody()).isNotNull();
+		assertThat(response.getHeaders()).containsEntry("X-Rate-Limit-Remaining", "1");
+	}
+
+	@Test
+	public void rateLimitTest_RestAPI() throws Exception {
+		ResourceApiAssert.assertResourceSearch(Map.of("limit", 0))
+			.statusCode(200)
+			.header("X-Rate-Limit-Remaining", "1");
+	}
+
+	@Test
+	public void rateLimitTest_JavaAPI_RateLimited() throws Exception {
+		// assume user consumed all available tokens from the bucket
+		rateLimiter.consume(RestExtensions.USER);
+		rateLimiter.consume(RestExtensions.USER);
+
+		try {
+			ResourceRequests.prepareSearch()
+				.setLimit(0)
+				.buildAsync()
+				.execute(Services.bus())
+				.getSyncResponse();
+			fail("Second request should throw a too many requests exception");
+		} catch (TooManyRequestsException e) {
+			// based on the fact that we have to 
+			assertThat(e.getSecondsToWait())
+				.isBetween(0L, 1L);
+			assertThat(e.getMessage())
+				.isEqualTo("Too many requests");
+		}
+	}
+
+	@Test
+	public void rateLimitTest_RestAPI_RateLimited() throws Exception {
+		// warm up request to initialize everything in the system before checking rate limits
+		ResourceApiAssert.assertResourceSearch(Map.of("limit", 0));
+
+		// assume user consumed all available tokens from the bucket
+		rateLimiter.consume(RestExtensions.USER);
+		rateLimiter.consume(RestExtensions.USER);
+
+		System.err.println("Sending resource request..." + LocalDateTime.now());
+		ResourceApiAssert.assertResourceSearch(Map.of("limit", 0))
+			.statusCode(429)
+			.header("X-Rate-Limit-Retry-After-Seconds", anyOf(equalTo("0"), equalTo("1")));
+	}
+
+}

core/com.b2international.snowowl.core.rest/src/com/b2international/snowowl/core/rest/util/PromiseMethodReturnValueHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 B2i Healthcare Pte Ltd, http://b2i.sg
+ * Copyright 2019-2023 B2i Healthcare Pte Ltd, http://b2i.sg
  * 
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,7 +18,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.core.MethodParameter;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.ResponseEntity;
+import org.springframework.http.ResponseEntity.BodyBuilder;
 import org.springframework.web.context.request.NativeWebRequest;
 import org.springframework.web.context.request.async.DeferredResult;
 import org.springframework.web.context.request.async.WebAsyncUtils;
@@ -27,6 +29,7 @@
 import org.springframework.web.method.support.ModelAndViewContainer;
 
 import com.b2international.snowowl.core.events.util.Promise;
+import com.b2international.snowowl.core.events.util.Response;
 
 /**
  * @since 7.2
@@ -51,20 +54,7 @@ public void handleReturnValue(Object returnValue, MethodParameter returnType, Mo
 			final Promise<?> promise = (Promise<?>) returnValue;
 			final DeferredResult<ResponseEntity<?>> result = new DeferredResult<>();
 			promise
-				.then(body -> {
-					if (result.isSetOrExpired()) {
-						LOG.warn("Deferred result is already set or expired, could not deliver result {}.", body);
-					} else { 
-						final ResponseEntity<?> response;
-						if (body instanceof ResponseEntity<?>) {
-							response = (ResponseEntity<?>) body;
-						} else {
-							response = ResponseEntity.ok().body(body);
-						}
-						result.setResult(response);
-					}
-					return null;
-				})
+				.thenRespond(promiseResponse -> setDeferredResult(result, promiseResponse))
 				.fail(err -> {
 					if (result.isSetOrExpired()) {
 						LOG.warn("Deferred result is already set or expired, could not deliver Throwable.", err);
@@ -80,6 +70,36 @@ public void handleReturnValue(Object returnValue, MethodParameter returnType, Mo
 		}
 	}
 
+	private Response<?> setDeferredResult(DeferredResult<ResponseEntity<?>> result, Response<?> promiseResponse) {
+		if (result.isSetOrExpired()) {
+			LOG.warn("Deferred result is already set or expired, could not deliver result {}.", promiseResponse);
+		} else { 
+			final Object body = promiseResponse.getBody();
+			final ResponseEntity<?> response;
+			if (body instanceof ResponseEntity<?> b) {
+				// return a custom ResponseEntity, copy it and apply headers returned from the system
+				HttpHeaders headers = b.getHeaders();
+				// append headers returned from system
+				promiseResponse.getHeaders().forEach((headerName, headerValue) -> {
+					headers.set(headerName, headerValue);
+				});
+				response = new ResponseEntity<>(b.getBody(), headers, b.getStatusCode());
+
+			} else {
+				// returning a standard object as reponse, use HTTP 200 OK
+				BodyBuilder responseBuilder = ResponseEntity.ok();
+				// append headers returned from system
+				promiseResponse.getHeaders().forEach((headerName, headerValue) -> {
+					responseBuilder.header(headerName, headerValue);
+				});
+				response = responseBuilder
+						.body(body);
+			}
+			result.setResult(response);
+		}
+		return null;
+	}
+
 	@Override
 	public boolean supportsReturnType(MethodParameter returnType) {
 		Class<?> type = returnType.getParameterType();