From 5fb158635c0f282d672590a1c4e6dccfe273c3f2 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Fri, 2 Dec 2022 12:14:11 -0600
Subject: [PATCH] Use existing `WebApplicationException` entity, if available

---
 .../mappers/CompletionExceptionMapper.java    | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/mappers/CompletionExceptionMapper.java b/service/src/main/java/org/whispersystems/textsecuregcm/mappers/CompletionExceptionMapper.java
index 829fcc7c9..9cb6dbc44 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/mappers/CompletionExceptionMapper.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/mappers/CompletionExceptionMapper.java
@@ -5,7 +5,9 @@
 
 package org.whispersystems.textsecuregcm.mappers;
 
+import java.util.Optional;
 import java.util.concurrent.CompletionException;
+import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.ext.ExceptionMapper;
@@ -23,8 +25,25 @@ public class CompletionExceptionMapper implements ExceptionMapper<CompletionExce
     final Throwable cause = exception.getCause();
 
     if (cause != null) {
+
       final Class type = cause.getClass();
-      return providers.getExceptionMapper(type).toResponse(cause);
+      final ExceptionMapper exceptionMapper = providers.getExceptionMapper(type);
+
+      // some exception mappers, like LoggingExceptionMapper, have side effects (e.g., logging)
+      // so we always build their response…
+      final Response exceptionMapperResponse = exceptionMapper.toResponse(cause);
+
+      final Optional<Response> webApplicationExceptionResponse;
+      if (cause instanceof WebApplicationException webApplicationException) {
+        webApplicationExceptionResponse = Optional.of(webApplicationException.getResponse());
+      } else {
+        webApplicationExceptionResponse = Optional.empty();
+      }
+
+      // …but if the exception was a WebApplicationException, and provides an entity, we want to keep it
+      return webApplicationExceptionResponse
+          .filter(Response::hasEntity)
+          .orElse(exceptionMapperResponse);
     }
 
     return Response.serverError().build();