From 152453790fbfce46a51bfe81a7d825957c9708d5 Mon Sep 17 00:00:00 2001
From: Stuart Douglas <stuart.w.douglas@gmail.com>
Date: Tue, 30 Apr 2024 17:27:35 +1000
Subject: [PATCH] Fix servlet error mapping

---
 .../undertow/deployment/UndertowBuildStep.java      | 11 +++++++++++
 .../undertow/test/ServletWebXmlTestCase.java        | 13 ++++++++++++-
 .../runtime/UndertowDeploymentRecorder.java         | 10 ++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/extensions/undertow/deployment/src/main/java/io/quarkus/undertow/deployment/UndertowBuildStep.java b/extensions/undertow/deployment/src/main/java/io/quarkus/undertow/deployment/UndertowBuildStep.java
index 415204b0d4fc2..98a1a8dbc3d2d 100644
--- a/extensions/undertow/deployment/src/main/java/io/quarkus/undertow/deployment/UndertowBuildStep.java
+++ b/extensions/undertow/deployment/src/main/java/io/quarkus/undertow/deployment/UndertowBuildStep.java
@@ -61,6 +61,7 @@
 import org.jboss.metadata.web.spec.CookieConfigMetaData;
 import org.jboss.metadata.web.spec.DispatcherType;
 import org.jboss.metadata.web.spec.EmptyRoleSemanticType;
+import org.jboss.metadata.web.spec.ErrorPageMetaData;
 import org.jboss.metadata.web.spec.FilterMappingMetaData;
 import org.jboss.metadata.web.spec.FilterMetaData;
 import org.jboss.metadata.web.spec.FiltersMetaData;
@@ -651,6 +652,16 @@ public ServletDeploymentManagerBuildItem build(List<ServletBuildItem> servlets,
             recorder.addServletContainerInitializer(deployment,
                     (Class<? extends ServletContainerInitializer>) context.classProxy(sci.sciClass), handlesTypes);
         }
+        if (webMetaData.getErrorPages() != null) {
+            for (ErrorPageMetaData errorPage : webMetaData.getErrorPages()) {
+                if (errorPage.getErrorCode() != null && !errorPage.getErrorCode().isBlank()) {
+                    recorder.addErrorPage(deployment, errorPage.getLocation(), Integer.parseInt(errorPage.getErrorCode()));
+                } else if (errorPage.getExceptionType() != null && !errorPage.getExceptionType().isBlank()) {
+                    recorder.addErrorPage(deployment, errorPage.getLocation(),
+                            (Class<? extends Throwable>) context.classProxy(errorPage.getExceptionType()));
+                }
+            }
+        }
         SessionConfigMetaData sessionConfig = webMetaData.getSessionConfig();
         if (sessionConfig != null) {
             if (sessionConfig.getSessionTimeoutSet()) {
diff --git a/extensions/undertow/deployment/src/test/java/io/quarkus/undertow/test/ServletWebXmlTestCase.java b/extensions/undertow/deployment/src/test/java/io/quarkus/undertow/test/ServletWebXmlTestCase.java
index 30b0e7c7122d6..64ca9d88894fc 100644
--- a/extensions/undertow/deployment/src/test/java/io/quarkus/undertow/test/ServletWebXmlTestCase.java
+++ b/extensions/undertow/deployment/src/test/java/io/quarkus/undertow/test/ServletWebXmlTestCase.java
@@ -33,7 +33,11 @@ public class ServletWebXmlTestCase {
             "  <mime-mapping>\n" +
             "    <extension>wasm</extension>\n" +
             "    <mime-type>application/wasm</mime-type>\n" +
-            "  </mime-mapping>" +
+            "  </mime-mapping>\n" +
+            "    <error-page>\n" +
+            "        <error-code>404</error-code>\n" +
+            "        <location>/mapped</location>\n" +
+            "    </error-page>\n" +
             "</web-app>";
 
     @RegisterExtension
@@ -56,4 +60,11 @@ public void testMimeMapping() {
                 .statusCode(200)
                 .contentType(is("application/wasm"));
     }
+
+    @Test
+    public void test404Mapping() {
+        RestAssured.when().get("/missing").then()
+                .statusCode(404)
+                .body(is("web xml servlet"));
+    }
 }
diff --git a/extensions/undertow/runtime/src/main/java/io/quarkus/undertow/runtime/UndertowDeploymentRecorder.java b/extensions/undertow/runtime/src/main/java/io/quarkus/undertow/runtime/UndertowDeploymentRecorder.java
index 082431efc6a9e..a0ccfe2a2d83d 100644
--- a/extensions/undertow/runtime/src/main/java/io/quarkus/undertow/runtime/UndertowDeploymentRecorder.java
+++ b/extensions/undertow/runtime/src/main/java/io/quarkus/undertow/runtime/UndertowDeploymentRecorder.java
@@ -747,6 +747,16 @@ public void setSessionCookieConfig(ServletSessionConfig config, String name, Str
         }
     }
 
+    public void addErrorPage(RuntimeValue<DeploymentInfo> deployment, String location, int errorCode) {
+        deployment.getValue().addErrorPage(new ErrorPage(location, errorCode));
+
+    }
+
+    public void addErrorPage(RuntimeValue<DeploymentInfo> deployment, String location,
+            Class<? extends Throwable> exceptionType) {
+        deployment.getValue().addErrorPage(new ErrorPage(location, exceptionType));
+    }
+
     /**
      * we can't have SecureRandom in the native image heap, so we need to lazy init
      */