fix(health): fix health check url authentication (#9117)

metadata-auth/auth-api/src/main/java/com/datahub/authentication/AuthenticationRequest.java

@@ -1,6 +1,8 @@
 package com.datahub.authentication;
 
 import com.datahub.plugins.auth.authentication.Authenticator;
+import lombok.Getter;
+
 import java.util.Map;
 import java.util.Objects;
 import java.util.TreeMap;
@@ -13,14 +15,24 @@ import javax.annotation.Nonnull;
  * Currently, this class only hold the inbound request's headers, but could certainly be extended
  * to contain additional information like the request parameters, body, ip, etc as needed.
  */
+@Getter
 public class AuthenticationRequest {
 
   private final Map<String, String> caseInsensitiveHeaders;
 
+  private final String servletInfo;
+  private final String pathInfo;
+
   public AuthenticationRequest(@Nonnull final Map<String, String> requestHeaders) {
+    this("", "", requestHeaders);
+  }
+
+  public AuthenticationRequest(@Nonnull String servletInfo, @Nonnull String pathInfo, @Nonnull final Map<String, String> requestHeaders) {
     Objects.requireNonNull(requestHeaders);
     caseInsensitiveHeaders = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
     caseInsensitiveHeaders.putAll(requestHeaders);
+    this.servletInfo = servletInfo;
+    this.pathInfo = pathInfo;
   }
 
   /**

metadata-service/auth-filter/src/main/java/com/datahub/auth/authentication/filter/AuthenticationFilter.java

@@ -2,6 +2,7 @@ package com.datahub.auth.authentication.filter;
 
 import com.datahub.authentication.authenticator.AuthenticatorChain;
 import com.datahub.authentication.authenticator.DataHubSystemAuthenticator;
+import com.datahub.authentication.authenticator.HealthStatusAuthenticator;
 import com.datahub.authentication.authenticator.NoOpAuthenticator;
 import com.datahub.authentication.token.StatefulTokenService;
 import com.datahub.plugins.PluginConstant;
@@ -29,6 +30,7 @@ import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -148,7 +150,7 @@ public class AuthenticationFilter implements Filter {
   }
 
   private AuthenticationRequest buildAuthContext(HttpServletRequest request) {
-    return new AuthenticationRequest(Collections.list(request.getHeaderNames())
+    return new AuthenticationRequest(request.getServletPath(), request.getPathInfo(), Collections.list(request.getHeaderNames())
         .stream()
         .collect(Collectors.toMap(headerName -> headerName, request::getHeader)));
   }
@@ -242,7 +244,14 @@ public class AuthenticationFilter implements Filter {
         final Authenticator authenticator = clazz.newInstance();
         // Successfully created authenticator. Now init and register it.
         log.debug(String.format("Initializing Authenticator with name %s", type));
+        if (authenticator instanceof HealthStatusAuthenticator) {
+          Map<String, Object> authenticatorConfig = new HashMap<>(Map.of(SYSTEM_CLIENT_ID_CONFIG,
+                  this.configurationProvider.getAuthentication().getSystemClientId()));
+          authenticatorConfig.putAll(Optional.ofNullable(internalAuthenticatorConfig.getConfigs()).orElse(Collections.emptyMap()));
+          authenticator.init(authenticatorConfig, authenticatorContext);
+        } else {
           authenticator.init(configs, authenticatorContext);
+        }
         log.info(String.format("Registering Authenticator with name %s", type));
         authenticatorChain.register(authenticator);
       } catch (Exception e) {

metadata-service/auth-impl/src/main/java/com/datahub/authentication/authenticator/HealthStatusAuthenticator.java

@@ -0,0 +1,55 @@
+package com.datahub.authentication.authenticator;
+
+import com.datahub.authentication.Actor;
+import com.datahub.authentication.ActorType;
+import com.datahub.authentication.Authentication;
+import com.datahub.authentication.AuthenticationException;
+import com.datahub.authentication.AuthenticationRequest;
+import com.datahub.authentication.AuthenticatorContext;
+import com.datahub.plugins.auth.authentication.Authenticator;
+import lombok.extern.slf4j.Slf4j;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+import static com.datahub.authentication.AuthenticationConstants.SYSTEM_CLIENT_ID_CONFIG;
+
+
+/**
+ * This Authenticator is used for allowing access for unauthenticated health check endpoints
+ *
+ * It exists to support load balancers, liveness/readiness checks
+ *
+ */
+@Slf4j
+public class HealthStatusAuthenticator implements Authenticator {
+  private static final Set<String> HEALTH_ENDPOINTS = Set.of(
+          "/openapi/check/",
+          "/openapi/up/"
+  );
+  private String systemClientId;
+
+  @Override
+  public void init(@Nonnull final Map<String, Object> config, @Nullable final AuthenticatorContext context) {
+    Objects.requireNonNull(config, "Config parameter cannot be null");
+    this.systemClientId = Objects.requireNonNull((String) config.get(SYSTEM_CLIENT_ID_CONFIG),
+            String.format("Missing required config %s", SYSTEM_CLIENT_ID_CONFIG));
+  }
+
+  @Override
+  public Authentication authenticate(@Nonnull AuthenticationRequest context) throws AuthenticationException {
+    Objects.requireNonNull(context);
+    if (HEALTH_ENDPOINTS.stream().anyMatch(prefix -> String.join("", context.getServletInfo(), context.getPathInfo()).startsWith(prefix))) {
+      return new Authentication(
+              new Actor(ActorType.USER, systemClientId),
+              "",
+              Collections.emptyMap()
+      );
+    }
+    throw new AuthenticationException("Authorization not allowed. Non-health check endpoint.");
+  }
+}

metadata-service/configuration/src/main/resources/application.yml

@@ -11,6 +11,8 @@ authentication:
         # Key used to validate incoming tokens. Should typically be the same as authentication.tokenService.signingKey
         signingKey: ${DATAHUB_TOKEN_SERVICE_SIGNING_KEY:WnEdIeTG/VVCLQqGwC/BAkqyY0k+H8NEAtWGejrBI94=}
         salt: ${DATAHUB_TOKEN_SERVICE_SALT:ohDVbJBvHHVJh9S/UA4BYF9COuNnqqVhr9MLKEGXk1O=}
+    # Required for unauthenticated health check endpoints - best not to remove.
+    - type: com.datahub.authentication.authenticator.HealthStatusAuthenticator
 
   # Normally failures are only warnings, enable this to throw them.
   logAuthenticatorExceptions: ${METADATA_SERVICE_AUTHENTICATOR_EXCEPTIONS_ENABLED:false}

metadata-service/health-servlet/build.gradle

@@ -1,22 +0,0 @@
-apply plugin: 'java'
-
-dependencies {
-
-  implementation project(':metadata-service:factories')
-
-  implementation externalDependency.guava
-  implementation externalDependency.reflections
-  implementation externalDependency.springBoot
-  implementation externalDependency.springCore
-  implementation externalDependency.springDocUI
-  implementation externalDependency.springWeb
-  implementation externalDependency.springWebMVC
-  implementation externalDependency.springBeans
-  implementation externalDependency.springContext
-  implementation externalDependency.slf4jApi
-  compileOnly externalDependency.lombok
-  implementation externalDependency.antlr4Runtime
-  implementation externalDependency.antlr4
-
-  annotationProcessor externalDependency.lombok
-}

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java

@@ -44,7 +44,6 @@ public class SpringWebConfig implements WebMvcConfigurer {
             .group("default")
             .packagesToExclude(
                     "io.datahubproject.openapi.operations",
-                    "com.datahub.health",
                     "io.datahubproject.openapi.health"
             ).build();
   }
@@ -55,7 +54,6 @@ public class SpringWebConfig implements WebMvcConfigurer {
             .group("operations")
             .packagesToScan(
                     "io.datahubproject.openapi.operations",
-                    "com.datahub.health",
                     "io.datahubproject.openapi.health"
             ).build();
   }

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/health/HealthCheckController.java

@@ -1,5 +1,6 @@
-package com.datahub.health.controller;
+package io.datahubproject.openapi.health;
 
+import com.google.common.base.Supplier;
 import com.google.common.base.Suppliers;
 import com.linkedin.gms.factory.config.ConfigurationProvider;
 import io.swagger.v3.oas.annotations.tags.Tag;
@@ -9,7 +10,6 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Supplier;
 
 import org.opensearch.action.admin.cluster.health.ClusterHealthRequest;
 import org.opensearch.action.admin.cluster.health.ClusterHealthResponse;
@@ -27,7 +27,7 @@ import org.springframework.web.bind.annotation.RestController;
 
 
 @RestController
-@RequestMapping("/check")
+@RequestMapping("/")
 @Tag(name = "HealthCheck", description = "An API for checking health of GMS and its clients.")
 public class HealthCheckController {
   @Autowired
@@ -41,6 +41,12 @@ public class HealthCheckController {
         this::getElasticHealth, config.getHealthCheck().getCacheDurationSeconds(), TimeUnit.SECONDS);
   }
 
+  @GetMapping(path = "/check/ready", produces = MediaType.APPLICATION_JSON_VALUE)
+  public ResponseEntity<Boolean> getCombinedHealthCheck(String... checks) {
+    return ResponseEntity.status(getCombinedDebug(checks).getStatusCode())
+            .body(getCombinedDebug(checks).getStatusCode().is2xxSuccessful());
+  }
+
   /**
    * Combined health check endpoint for checking GMS clients.
    * For now, just checks the health of the ElasticSearch client
@@ -48,11 +54,10 @@ public class HealthCheckController {
    * that component). The status code will be 200 if all components are okay, and 500 if one or more components are not
    * healthy.
    */
-  @GetMapping(path = "/ready", produces = MediaType.APPLICATION_JSON_VALUE)
-  public ResponseEntity<Map<String, ResponseEntity<String>>> getCombinedHealthCheck(String... checks) {
-
+  @GetMapping(path = "/debug/ready", produces = MediaType.APPLICATION_JSON_VALUE)
+  public ResponseEntity<Map<String, ResponseEntity<String>>> getCombinedDebug(String... checks) {
     Map<String, Supplier<ResponseEntity<String>>> healthChecks = new HashMap<>();
-    healthChecks.put("elasticsearch", this::getElasticHealthWithCache);
+    healthChecks.put("elasticsearch", this::getElasticDebugWithCache);
     // Add new components here
 
     List<String> componentsToCheck = checks != null && checks.length > 0
@@ -67,7 +72,6 @@ public class HealthCheckController {
               .get());
     }
 
-
     boolean isHealthy = componentHealth.values().stream().allMatch(resp -> resp.getStatusCode() == HttpStatus.OK);
     if (isHealthy) {
       return ResponseEntity.ok(componentHealth);
@@ -75,12 +79,18 @@ public class HealthCheckController {
     return ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE).body(componentHealth);
   }
 
+  @GetMapping(path = "/check/elastic", produces = MediaType.APPLICATION_JSON_VALUE)
+  public ResponseEntity<Boolean> getElasticHealthWithCache() {
+    return ResponseEntity.status(getElasticDebugWithCache().getStatusCode())
+            .body(getElasticDebugWithCache().getStatusCode().is2xxSuccessful());
+  }
+
   /**
    * Checks the memoized cache for the latest elastic health check result
    * @return The ResponseEntity containing the health check result
    */
-  @GetMapping(path = "/elastic", produces = MediaType.APPLICATION_JSON_VALUE)
-  public ResponseEntity<String> getElasticHealthWithCache() {
+  @GetMapping(path = "/debug/elastic", produces = MediaType.APPLICATION_JSON_VALUE)
+  public ResponseEntity<String> getElasticDebugWithCache() {
     return this.memoizedSupplier.get();
   }
 

metadata-service/war/build.gradle

@@ -17,7 +17,6 @@ dependencies {
   runtimeOnly project(':metadata-service:servlet')
   runtimeOnly project(':metadata-service:auth-servlet-impl')
   runtimeOnly project(':metadata-service:graphql-servlet-impl')
-  runtimeOnly project(':metadata-service:health-servlet')
   runtimeOnly project(':metadata-service:openapi-servlet')
   runtimeOnly project(':metadata-service:openapi-entity-servlet')
   runtimeOnly project(':metadata-service:openapi-analytics-servlet')

metadata-service/war/src/main/webapp/WEB-INF/openapiServlet-servlet.xml

@@ -3,7 +3,7 @@
        xmlns:context="http://www.springframework.org/schema/context"
        xsi:schemaLocation="http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans-3.0.xsd http://www.springframework.org/schema/context http://www.springframework.org/schema/context/spring-context-3.0.xsd">
 
-  <context:component-scan base-package="io.datahubproject.openapi,com.datahub.health"/>
+  <context:component-scan base-package="io.datahubproject.openapi"/>
   <context:component-scan base-package="org.springdoc.webmvc.ui,org.springdoc.core,org.springdoc.webmvc.core,org.springframework.boot.autoconfigure.jackson"/>
 
 <bean id="yamlProperties" class="org.springframework.beans.factory.config.YamlPropertiesFactoryBean">

settings.gradle

@@ -8,7 +8,6 @@ include 'metadata-service:auth-config'
 include 'metadata-service:auth-impl'
 include 'metadata-service:auth-filter'
 include 'metadata-service:auth-servlet-impl'
-include 'metadata-service:health-servlet'
 include 'metadata-service:restli-api'
 include 'metadata-service:restli-client'
 include 'metadata-service:restli-servlet-impl'