From f6f3cbcb91decf9a9794ce7cc46d5aef5526396e Mon Sep 17 00:00:00 2001
From: Georgios Andrianakis <geoand@gmail.com>
Date: Mon, 5 Aug 2024 09:58:33 +0300
Subject: [PATCH] Add more supported types to @ClientExceptionMapper

Closes: #42293
---
 .../ClientExceptionMapperHandler.java         | 57 +++++++++++++++++--
 .../client/reactive/deployment/DotNames.java  |  8 +++
 .../RegisteredClientExceptionMapperTest.java  | 13 ++++-
 .../reactive/ClientExceptionMapper.java       | 17 +++++-
 .../handlers/ClientSendRequestHandler.java    | 13 ++---
 .../client/impl/RestClientRequestContext.java |  4 ++
 6 files changed, 95 insertions(+), 17 deletions(-)

diff --git a/extensions/resteasy-reactive/rest-client/deployment/src/main/java/io/quarkus/rest/client/reactive/deployment/ClientExceptionMapperHandler.java b/extensions/resteasy-reactive/rest-client/deployment/src/main/java/io/quarkus/rest/client/reactive/deployment/ClientExceptionMapperHandler.java
index ae97af6d7c3b3..1730c5a93e0bd 100644
--- a/extensions/resteasy-reactive/rest-client/deployment/src/main/java/io/quarkus/rest/client/reactive/deployment/ClientExceptionMapperHandler.java
+++ b/extensions/resteasy-reactive/rest-client/deployment/src/main/java/io/quarkus/rest/client/reactive/deployment/ClientExceptionMapperHandler.java
@@ -2,9 +2,13 @@
 
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
+import java.net.URI;
 import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
 
 import jakarta.ws.rs.Priorities;
+import jakarta.ws.rs.core.MultivaluedMap;
 import jakarta.ws.rs.core.Response;
 
 import org.eclipse.microprofile.rest.client.ext.ResponseExceptionMapper;
@@ -12,7 +16,9 @@
 import org.jboss.jandex.AnnotationTarget;
 import org.jboss.jandex.AnnotationValue;
 import org.jboss.jandex.ClassInfo;
+import org.jboss.jandex.DotName;
 import org.jboss.jandex.MethodInfo;
+import org.jboss.jandex.ParameterizedType;
 import org.jboss.jandex.Type;
 import org.jboss.resteasy.reactive.client.impl.RestClientRequestContext;
 import org.jboss.resteasy.reactive.common.processor.ResteasyReactiveDotNames;
@@ -35,6 +41,12 @@ class ClientExceptionMapperHandler {
     private static final ResultHandle[] EMPTY_RESULT_HANDLES_ARRAY = new ResultHandle[0];
     private static final MethodDescriptor GET_INVOKED_METHOD =
             MethodDescriptor.ofMethod(RestClientRequestContext.class, "getInvokedMethod", Method.class);
+    private static final MethodDescriptor GET_URI =
+            MethodDescriptor.ofMethod(RestClientRequestContext.class, "getUri", URI.class);
+    private static final MethodDescriptor GET_PROPERTIES =
+            MethodDescriptor.ofMethod(RestClientRequestContext.class, "getProperties", Map.class);
+    private static final MethodDescriptor GET_REQUEST_HEADERS_AS_MAP =
+            MethodDescriptor.ofMethod(RestClientRequestContext.class, "getRequestHeadersAsMap", MultivaluedMap.class);
     private final ClassOutput classOutput;
 
     ClientExceptionMapperHandler(ClassOutput classOutput) {
@@ -105,17 +117,24 @@ GeneratedClassResult generateResponseExceptionMapper(AnnotationInstance instance
             LinkedHashMap<String, ResultHandle> targetMethodParams = new LinkedHashMap<>();
             for (Type paramType : targetMethod.parameterTypes()) {
                 ResultHandle targetMethodParamHandle;
-                if (paramType.name().equals(ResteasyReactiveDotNames.RESPONSE)) {
+                DotName paramTypeName = paramType.name();
+                if (paramTypeName.equals(ResteasyReactiveDotNames.RESPONSE)) {
                     targetMethodParamHandle = toThrowable.getMethodParam(0);
-                } else if (paramType.name().equals(DotNames.METHOD)) {
+                } else if (paramTypeName.equals(DotNames.METHOD)) {
                     targetMethodParamHandle = toThrowable.invokeVirtualMethod(GET_INVOKED_METHOD, toThrowable.getMethodParam(1));
+                } else if (paramTypeName.equals(DotNames.URI)) {
+                    targetMethodParamHandle = toThrowable.invokeVirtualMethod(GET_URI, toThrowable.getMethodParam(1));
+                } else if (isMapStringToObject(paramType)) {
+                    targetMethodParamHandle = toThrowable.invokeVirtualMethod(GET_PROPERTIES, toThrowable.getMethodParam(1));
+                } else if (isMultivaluedMapStringToString(paramType)) {
+                    targetMethodParamHandle = toThrowable.invokeVirtualMethod(GET_REQUEST_HEADERS_AS_MAP, toThrowable.getMethodParam(1));
                 } else {
-                    String message = DotNames.CLIENT_EXCEPTION_MAPPER + " can only take parameters of type '" + ResteasyReactiveDotNames.RESPONSE + "' or '" + DotNames.METHOD + "'"
+                    String message = "Unsupported parameter type used in " + DotNames.CLIENT_EXCEPTION_MAPPER + ". See the Javadoc of the annotation for the supported types."
                     + " Offending instance is '" + targetMethod.declaringClass().name().toString() + "#"
                             + targetMethod.name() + "'";
                     throw new IllegalStateException(message);
                 }
-                targetMethodParams.put(paramType.name().toString(), targetMethodParamHandle);
+                targetMethodParams.put(paramTypeName.toString(), targetMethodParamHandle);
             }
 
             ResultHandle resultHandle = toThrowable.invokeStaticInterfaceMethod(
@@ -136,6 +155,36 @@ GeneratedClassResult generateResponseExceptionMapper(AnnotationInstance instance
         return new GeneratedClassResult(restClientInterfaceClassInfo.name().toString(), generatedClassName, priority);
     }
 
+    private boolean isMapStringToObject(Type paramType) {
+        if (paramType.kind() != Type.Kind.PARAMETERIZED_TYPE) {
+            return false;
+        }
+        ParameterizedType parameterizedType = paramType.asParameterizedType();
+        if (!parameterizedType.name().equals(DotNames.MAP)) {
+            return false;
+        }
+        List<Type> arguments = parameterizedType.arguments();
+        if (arguments.size() != 2) {
+            return false;
+        }
+        return arguments.get(0).name().equals(DotNames.STRING) && arguments.get(1).name().equals(DotNames.OBJECT);
+    }
+
+    private boolean isMultivaluedMapStringToString(Type paramType) {
+        if (paramType.kind() != Type.Kind.PARAMETERIZED_TYPE) {
+            return false;
+        }
+        ParameterizedType parameterizedType = paramType.asParameterizedType();
+        if (!parameterizedType.name().equals(DotNames.MULTIVALUED_MAP)) {
+            return false;
+        }
+        List<Type> arguments = parameterizedType.arguments();
+        if (arguments.size() != 2) {
+            return false;
+        }
+        return arguments.get(0).name().equals(DotNames.STRING) && arguments.get(1).name().equals(DotNames.STRING);
+    }
+
     public static String getGeneratedClassName(MethodInfo methodInfo) {
         StringBuilder sigBuilder = new StringBuilder();
         sigBuilder.append(methodInfo.name()).append("_").append(methodInfo.returnType().name().toString());
diff --git a/extensions/resteasy-reactive/rest-client/deployment/src/main/java/io/quarkus/rest/client/reactive/deployment/DotNames.java b/extensions/resteasy-reactive/rest-client/deployment/src/main/java/io/quarkus/rest/client/reactive/deployment/DotNames.java
index 897409ee3d7d4..dba2de1a3b40e 100644
--- a/extensions/resteasy-reactive/rest-client/deployment/src/main/java/io/quarkus/rest/client/reactive/deployment/DotNames.java
+++ b/extensions/resteasy-reactive/rest-client/deployment/src/main/java/io/quarkus/rest/client/reactive/deployment/DotNames.java
@@ -1,9 +1,12 @@
 package io.quarkus.rest.client.reactive.deployment;
 
 import java.lang.reflect.Method;
+import java.net.URI;
+import java.util.Map;
 
 import jakarta.ws.rs.client.ClientRequestFilter;
 import jakarta.ws.rs.client.ClientResponseFilter;
+import jakarta.ws.rs.core.MultivaluedMap;
 
 import org.eclipse.microprofile.rest.client.annotation.ClientHeaderParam;
 import org.eclipse.microprofile.rest.client.annotation.ClientHeaderParams;
@@ -44,6 +47,11 @@ public class DotNames {
     public static final DotName RESPONSE_EXCEPTION_MAPPER = DotName.createSimple(ResponseExceptionMapper.class.getName());
 
     static final DotName METHOD = DotName.createSimple(Method.class.getName());
+    static final DotName URI = DotName.createSimple(URI.class.getName());
+    static final DotName MAP = DotName.createSimple(Map.class.getName());
+    static final DotName MULTIVALUED_MAP = DotName.createSimple(MultivaluedMap.class.getName());
+    static final DotName STRING = DotName.createSimple(String.class.getName());
+    static final DotName OBJECT = DotName.createSimple(Object.class.getName());
 
     public static final DotName SSE_EVENT_FILTER = DotName.createSimple(SseEventFilter.class);
 
diff --git a/extensions/resteasy-reactive/rest-client/deployment/src/test/java/io/quarkus/rest/client/reactive/error/clientexceptionmapper/RegisteredClientExceptionMapperTest.java b/extensions/resteasy-reactive/rest-client/deployment/src/test/java/io/quarkus/rest/client/reactive/error/clientexceptionmapper/RegisteredClientExceptionMapperTest.java
index 815e8c1b21c2c..85315052f150d 100644
--- a/extensions/resteasy-reactive/rest-client/deployment/src/test/java/io/quarkus/rest/client/reactive/error/clientexceptionmapper/RegisteredClientExceptionMapperTest.java
+++ b/extensions/resteasy-reactive/rest-client/deployment/src/test/java/io/quarkus/rest/client/reactive/error/clientexceptionmapper/RegisteredClientExceptionMapperTest.java
@@ -5,10 +5,13 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.lang.reflect.Method;
+import java.net.URI;
+import java.util.Map;
 
 import jakarta.ws.rs.GET;
 import jakarta.ws.rs.Path;
 import jakarta.ws.rs.Priorities;
+import jakarta.ws.rs.core.MultivaluedMap;
 import jakarta.ws.rs.core.Response;
 
 import org.eclipse.microprofile.rest.client.annotation.RegisterProvider;
@@ -120,8 +123,14 @@ public interface ClientWithRegisteredLowPriorityMapper {
         Dto get400();
 
         @ClientExceptionMapper
-        static DummyException map(Method method, Response response) {
-            if ((response.getStatus() == 404) && method.getName().equals("get404")) {
+        static DummyException map(Method method, Response response, URI uri,
+                Map<String, Object> properties, MultivaluedMap<String, String> requestHeaders) {
+            // the conditions here make sure that the mapper is passed all the data we expect it to be passed
+            if ((response.getStatus() == 404)
+                    && method.getName().equals("get404")
+                    && uri.getPath().equals("/error/404")
+                    && properties.containsKey("org.eclipse.microprofile.rest.client.invokedMethod")
+                    && requestHeaders.containsKey("User-Agent")) {
                 return new DummyException();
             }
             return null;
diff --git a/extensions/resteasy-reactive/rest-client/runtime/src/main/java/io/quarkus/rest/client/reactive/ClientExceptionMapper.java b/extensions/resteasy-reactive/rest-client/runtime/src/main/java/io/quarkus/rest/client/reactive/ClientExceptionMapper.java
index 6541390621c05..67a5d37b9b677 100644
--- a/extensions/resteasy-reactive/rest-client/runtime/src/main/java/io/quarkus/rest/client/reactive/ClientExceptionMapper.java
+++ b/extensions/resteasy-reactive/rest-client/runtime/src/main/java/io/quarkus/rest/client/reactive/ClientExceptionMapper.java
@@ -10,19 +10,30 @@
 /**
  * Used to easily define an exception mapper for the specific REST Client on which it's used.
  * This method is called when the HTTP response from the invoked service has a status code of 400 or higher.
- *
+ * <p>
  * The annotation MUST be placed on a method of the REST Client interface that meets the following criteria:
+ *
  * <ul>
  * <li>Is a {@code static} method</li>
  * <li>Returns any subclass of {@link RuntimeException}</li>
- * <li>Takes a single parameter of type {@link jakarta.ws.rs.core.Response}</li>
+ * </ul>
+ *
+ * The method can utilize any combination of the following parameters:
+ *
+ * <ul>
+ * <li>{@code jakarta.ws.rs.core.Response} which represents the HTTP response</li>
+ * <li>{@code Method} which represents the invoked method of the client</li>
+ * <li>{@code URI} which represents the the request URI</li>
+ * <li>{@code Map<String, Object>} which gives access to the properties that are available to (and potentially changed by)
+ * {@link jakarta.ws.rs.client.ClientRequestContext}</li>
+ * <li>{@code jakarta.ws.rs.core.MultivaluedMap} containing the request headers</li>
  * </ul>
  *
  * An example method could look like the following:
  *
  * <pre>
  * {@code
- * &#64;ClientExceptionMapper
+ * @ClientExceptionMapper
  * static DummyException map(Response response, Method method) {
  *     if (response.getStatus() == 404) {
  *         return new DummyException();
diff --git a/independent-projects/resteasy-reactive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/handlers/ClientSendRequestHandler.java b/independent-projects/resteasy-reactive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/handlers/ClientSendRequestHandler.java
index 93a15882c887f..2d34bc05690ea 100644
--- a/independent-projects/resteasy-reactive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/handlers/ClientSendRequestHandler.java
+++ b/independent-projects/resteasy-reactive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/handlers/ClientSendRequestHandler.java
@@ -158,8 +158,7 @@ public void handle(AsyncResult<AsyncFile> openedAsyncFile) {
                                                 return;
                                             }
 
-                                            MultivaluedMap<String, String> headerMap = requestContext.getRequestHeaders()
-                                                    .asMap();
+                                            MultivaluedMap<String, String> headerMap = requestContext.getRequestHeadersAsMap();
                                             updateRequestHeadersFromConfig(requestContext, headerMap);
 
                                             // set the Vertx headers after we've run the interceptors because they can modify them
@@ -170,8 +169,7 @@ public void handle(AsyncResult<AsyncFile> openedAsyncFile) {
                                         }
                                     });
                 } else if (requestContext.isInputStreamUpload() && !hasWriterInterceptors(requestContext)) {
-                    MultivaluedMap<String, String> headerMap = requestContext.getRequestHeaders()
-                            .asMap();
+                    MultivaluedMap<String, String> headerMap = requestContext.getRequestHeadersAsMap();
                     updateRequestHeadersFromConfig(requestContext, headerMap);
                     setVertxHeaders(httpClientRequest, headerMap);
                     Future<HttpClientResponse> sent = httpClientRequest.send(
@@ -180,8 +178,7 @@ public void handle(AsyncResult<AsyncFile> openedAsyncFile) {
                                     httpClientRequest));
                     attachSentHandlers(sent, httpClientRequest, requestContext);
                 } else if (requestContext.isMultiBufferUpload()) {
-                    MultivaluedMap<String, String> headerMap = requestContext.getRequestHeaders()
-                            .asMap();
+                    MultivaluedMap<String, String> headerMap = requestContext.getRequestHeadersAsMap();
                     updateRequestHeadersFromConfig(requestContext, headerMap);
                     setVertxHeaders(httpClientRequest, headerMap);
                     Future<HttpClientResponse> sent = httpClientRequest.send(ReadStreamSubscriber.asReadStream(
@@ -483,7 +480,7 @@ private QuarkusMultipartFormUpload setMultipartHeadersAndPrepareBody(HttpClientR
                     "Multipart form upload expects an entity of type MultipartForm, got: " + state.getEntity().getEntity());
         }
 
-        MultivaluedMap<String, String> headerMap = state.getRequestHeaders().asMap();
+        MultivaluedMap<String, String> headerMap = state.getRequestHeadersAsMap();
         updateRequestHeadersFromConfig(state, headerMap);
         QuarkusMultipartForm multipartForm = (QuarkusMultipartForm) state.getEntity().getEntity();
         multipartForm.preparePojos(state);
@@ -524,7 +521,7 @@ private QuarkusMultipartFormUpload setMultipartHeadersAndPrepareBody(HttpClientR
     private Buffer setRequestHeadersAndPrepareBody(HttpClientRequest httpClientRequest,
             RestClientRequestContext state)
             throws IOException {
-        MultivaluedMap<String, String> headerMap = state.getRequestHeaders().asMap();
+        MultivaluedMap<String, String> headerMap = state.getRequestHeadersAsMap();
         updateRequestHeadersFromConfig(state, headerMap);
 
         Buffer actualEntity = AsyncInvokerImpl.EMPTY_BUFFER;
diff --git a/independent-projects/resteasy-reactive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/impl/RestClientRequestContext.java b/independent-projects/resteasy-reactive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/impl/RestClientRequestContext.java
index 134cf0d594313..9ce0103471997 100644
--- a/independent-projects/resteasy-reactive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/impl/RestClientRequestContext.java
+++ b/independent-projects/resteasy-reactive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/impl/RestClientRequestContext.java
@@ -429,6 +429,10 @@ public ClientRequestHeaders getRequestHeaders() {
         return requestHeaders;
     }
 
+    public MultivaluedMap<String, String> getRequestHeadersAsMap() {
+        return requestHeaders.asMap();
+    }
+
     public String getHttpMethod() {
         return httpMethod;
     }