From 9a643c3a41feb67a473bc43d6ac23f4baa29b553 Mon Sep 17 00:00:00 2001
From: Mohit Yadav <105265192+mohityadav766@users.noreply.github.com>
Date: Thu, 13 Jul 2023 21:16:05 +0530
Subject: [PATCH] Fix Missing We Config on Error Page Response in #UI (#12428)

---
 .../service/OpenMetadataApplication.java      |  4 +-
 .../service/exception/OMErrorPageHandler.java | 79 +++++++++++++++++++
 2 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/exception/OMErrorPageHandler.java

diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java
index d804b201e2f..7bedba10191 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java
@@ -58,7 +58,6 @@ import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
 import org.eclipse.jetty.http.pathmap.ServletPathSpec;
-import org.eclipse.jetty.servlet.ErrorPageErrorHandler;
 import org.eclipse.jetty.servlet.FilterHolder;
 import org.eclipse.jetty.servlet.ServletHolder;
 import org.eclipse.jetty.websocket.server.NativeWebSocketServletContainerInitializer;
@@ -82,6 +81,7 @@ import org.openmetadata.service.events.scheduled.ReportsHandler;
 import org.openmetadata.service.exception.CatalogGenericExceptionMapper;
 import org.openmetadata.service.exception.ConstraintViolationExceptionMapper;
 import org.openmetadata.service.exception.JsonMappingExceptionMapper;
+import org.openmetadata.service.exception.OMErrorPageHandler;
 import org.openmetadata.service.extension.OpenMetadataExtension;
 import org.openmetadata.service.fernet.Fernet;
 import org.openmetadata.service.jdbi3.CollectionDAO;
@@ -437,7 +437,7 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
     CollectionRegistry.initialize(extensionResources);
     CollectionRegistry.getInstance().registerResources(jdbi, environment, config, authorizer, authenticatorHandler);
     environment.jersey().register(new JsonPatchProvider());
-    ErrorPageErrorHandler eph = new ErrorPageErrorHandler();
+    OMErrorPageHandler eph = new OMErrorPageHandler(config.getWebConfiguration());
     eph.addErrorPage(Response.Status.NOT_FOUND.getStatusCode(), "/");
     environment.getApplicationContext().setErrorHandler(eph);
   }
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/exception/OMErrorPageHandler.java b/openmetadata-service/src/main/java/org/openmetadata/service/exception/OMErrorPageHandler.java
new file mode 100644
index 00000000000..6e8707ddc7d
--- /dev/null
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/exception/OMErrorPageHandler.java
@@ -0,0 +1,79 @@
+package org.openmetadata.service.exception;
+
+import io.dropwizard.web.conf.WebConfiguration;
+import java.io.IOException;
+import javax.servlet.ServletException;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import lombok.extern.slf4j.Slf4j;
+import org.eclipse.jetty.http.HttpHeader;
+import org.eclipse.jetty.server.Dispatcher;
+import org.eclipse.jetty.server.Request;
+import org.eclipse.jetty.server.handler.ContextHandler;
+import org.eclipse.jetty.servlet.ErrorPageErrorHandler;
+
+@Slf4j
+public class OMErrorPageHandler extends ErrorPageErrorHandler {
+  private final WebConfiguration webConfiguration;
+
+  public OMErrorPageHandler(WebConfiguration webConfiguration) {
+    this.webConfiguration = webConfiguration;
+  }
+
+  @Override
+  public void handle(String target, Request baseRequest, HttpServletRequest request, HttpServletResponse response)
+      throws IOException, ServletException {
+    this.doError(target, baseRequest, request, response);
+  }
+
+  @Override
+  public void doError(String target, Request baseRequest, HttpServletRequest request, HttpServletResponse response)
+      throws IOException {
+    String cacheControl = this.getCacheControl();
+    if (cacheControl != null) {
+      response.setHeader(HttpHeader.CACHE_CONTROL.asString(), cacheControl);
+    }
+
+    // Attach Response Header from OM
+
+    // Hsts
+    webConfiguration.getHstsHeaderFactory().build().forEach(response::setHeader);
+
+    // Frame Options
+    webConfiguration.getFrameOptionsHeaderFactory().build().forEach(response::setHeader);
+
+    // Content Option
+    webConfiguration.getContentTypeOptionsHeaderFactory().build().forEach(response::setHeader);
+
+    // Xss Protections
+    webConfiguration.getXssProtectionHeaderFactory().build().forEach(response::setHeader);
+
+    String errorPage = ((ErrorPageMapper) this).getErrorPage(request);
+    ContextHandler.Context context = baseRequest.getErrorContext();
+    Dispatcher errorDispatcher =
+        errorPage != null && context != null ? (Dispatcher) context.getRequestDispatcher(errorPage) : null;
+
+    try {
+      if (errorDispatcher != null) {
+        try {
+          errorDispatcher.error(request, response);
+          return;
+        } catch (ServletException ex) {
+          LOG.debug("Error in OMErrorPageHandler", ex);
+          if (response.isCommitted()) {
+            return;
+          }
+        }
+      }
+
+      String message = (String) request.getAttribute("javax.servlet.error.message");
+      if (message == null) {
+        message = baseRequest.getResponse().getReason();
+      }
+
+      this.generateAcceptableResponse(baseRequest, request, response, response.getStatus(), message);
+    } finally {
+      baseRequest.setHandled(true);
+    }
+  }
+}