datahub/datahub-frontend/app/client/AuthServiceClient.java

package client;

import static com.linkedin.metadata.Constants.DATAHUB_LOGIN_SOURCE_HEADER_NAME;

import com.datahub.authentication.Authentication;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.inject.Inject;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
import javax.annotation.Nonnull;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.HttpEntity;
import org.apache.http.HttpStatus;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.util.EntityUtils;
import play.mvc.Http;

/** This class is responsible for coordinating authentication with the backend Metadata Service. */
@Slf4j
public class AuthServiceClient {

  private static final String GENERATE_SESSION_TOKEN_ENDPOINT = "auth/generateSessionTokenForUser";
  private static final String SIGN_UP_ENDPOINT = "auth/signUp";
  private static final String RESET_NATIVE_USER_CREDENTIALS_ENDPOINT =
      "auth/resetNativeUserCredentials";
  private static final String VERIFY_NATIVE_USER_CREDENTIALS_ENDPOINT =
      "auth/verifyNativeUserCredentials";
  private static final String TRACK_ENDPOINT = "auth/track";
  private static final String ACCESS_TOKEN_FIELD = "accessToken";
  private static final String USER_ID_FIELD = "userId";
  private static final String USER_URN_FIELD = "userUrn";
  private static final String FULL_NAME_FIELD = "fullName";
  private static final String EMAIL_FIELD = "email";
  private static final String TITLE_FIELD = "title";
  private static final String PASSWORD_FIELD = "password";
  private static final String INVITE_TOKEN_FIELD = "inviteToken";
  private static final String RESET_TOKEN_FIELD = "resetToken";
  private static final String IS_NATIVE_USER_CREATED_FIELD = "isNativeUserCreated";
  private static final String ARE_NATIVE_USER_CREDENTIALS_RESET_FIELD =
      "areNativeUserCredentialsReset";
  private static final String DOES_PASSWORD_MATCH_FIELD = "doesPasswordMatch";

  private final String metadataServiceHost;
  private final Integer metadataServicePort;
  private final Boolean metadataServiceUseSsl;
  private final Authentication systemAuthentication;
  private final CloseableHttpClient httpClient;

  @Inject
  public AuthServiceClient(
      @Nonnull final String metadataServiceHost,
      @Nonnull final Integer metadataServicePort,
      @Nonnull final Boolean useSsl,
      @Nonnull final Authentication systemAuthentication,
      @Nonnull final CloseableHttpClient httpClient) {
    this.metadataServiceHost = Objects.requireNonNull(metadataServiceHost);
    this.metadataServicePort = Objects.requireNonNull(metadataServicePort);
    this.metadataServiceUseSsl = Objects.requireNonNull(useSsl);
    this.systemAuthentication = Objects.requireNonNull(systemAuthentication);
    this.httpClient = Objects.requireNonNull(httpClient);
  }

  /**
   * Call the Auth Service to generate a session token for a particular user with a unique actor id,
   * or throws an exception if generation fails.
   *
   * <p>Notice that the "userId" parameter should NOT be of type "urn", but rather the unique id of
   * an Actor of type USER.
   */
  @Nonnull
  public String generateSessionTokenForUser(@Nonnull final String userId, String loginSource) {
    Objects.requireNonNull(userId, "userId must not be null");
    CloseableHttpResponse response = null;

    try {
      final String protocol = this.metadataServiceUseSsl ? "https" : "http";
      final HttpPost request =
          new HttpPost(
              String.format(
                  "%s://%s:%s/%s",
                  protocol,
                  this.metadataServiceHost,
                  this.metadataServicePort,
                  GENERATE_SESSION_TOKEN_ENDPOINT));

      log.info("Requesting session token for user: {}", userId);

      // Build JSON request to generate a token on behalf of a user.
      final ObjectMapper objectMapper = new ObjectMapper();
      final ObjectNode objectNode = objectMapper.createObjectNode();
      objectNode.put(USER_ID_FIELD, userId);
      final String json =
          objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(objectNode);
      request.setEntity(new StringEntity(json, StandardCharsets.UTF_8));

      // Add authorization header with DataHub frontend system id and secret.
      request.addHeader(Http.HeaderNames.AUTHORIZATION, this.systemAuthentication.getCredentials());

      request.addHeader(DATAHUB_LOGIN_SOURCE_HEADER_NAME, loginSource);

      response = httpClient.execute(request);
      final HttpEntity entity = response.getEntity();
      if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK && entity != null) {
        log.info("Successfully received session token for user: {}", userId);
        final String jsonStr = EntityUtils.toString(entity);
        return getAccessTokenFromJson(jsonStr);
      } else {
        throw new RuntimeException(
            String.format(
                "Bad response from the Metadata Service: %s %s",
                response.getStatusLine().toString(), response.getEntity().toString()));
      }
    } catch (Exception e) {
      log.error("Failed to generate session token for user: {}", userId, e);
      throw new RuntimeException("Failed to generate session token for user", e);
    } finally {
      try {
        if (response != null) {
          response.close();
        }
      } catch (Exception e) {
        log.error("Failed to close http response", e);
      }
    }
  }

  /** Call the Auth Service to create a native Datahub user. */
  public boolean signUp(
      @Nonnull final String userUrn,
      @Nonnull final String fullName,
      @Nonnull final String email,
      @Nonnull final String title,
      @Nonnull final String password,
      @Nonnull final String inviteToken) {
    Objects.requireNonNull(userUrn, "userUrn must not be null");
    Objects.requireNonNull(fullName, "fullName must not be null");
    Objects.requireNonNull(email, "email must not be null");
    Objects.requireNonNull(title, "title must not be null");
    Objects.requireNonNull(password, "password must not be null");
    Objects.requireNonNull(inviteToken, "inviteToken must not be null");
    CloseableHttpResponse response = null;

    try {

      final String protocol = this.metadataServiceUseSsl ? "https" : "http";
      final HttpPost request =
          new HttpPost(
              String.format(
                  "%s://%s:%s/%s",
                  protocol, this.metadataServiceHost, this.metadataServicePort, SIGN_UP_ENDPOINT));

      // Build JSON request to sign up a native user.
      final ObjectMapper objectMapper = new ObjectMapper();
      final ObjectNode objectNode = objectMapper.createObjectNode();
      objectNode.put(USER_URN_FIELD, userUrn);
      objectNode.put(FULL_NAME_FIELD, fullName);
      objectNode.put(EMAIL_FIELD, email);
      objectNode.put(TITLE_FIELD, title);
      objectNode.put(PASSWORD_FIELD, password);
      objectNode.put(INVITE_TOKEN_FIELD, inviteToken);
      final String json =
          objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(objectNode);
      request.setEntity(new StringEntity(json, StandardCharsets.UTF_8));

      // Add authorization header with DataHub frontend system id and secret.
      request.addHeader(Http.HeaderNames.AUTHORIZATION, this.systemAuthentication.getCredentials());

      response = httpClient.execute(request);
      final HttpEntity entity = response.getEntity();
      if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK && entity != null) {
        // Successfully generated a token for the User
        final String jsonStr = EntityUtils.toString(entity);
        return getIsNativeUserCreatedFromJson(jsonStr);
      } else {
        String content =
            response.getEntity().getContent() == null
                ? ""
                : new String(
                    response.getEntity().getContent().readAllBytes(), StandardCharsets.UTF_8);
        throw new RuntimeException(
            String.format(
                "Bad response from the Metadata Service: %s %s Body: %s",
                response.getStatusLine().toString(), response.getEntity().toString(), content));
      }
    } catch (Exception e) {
      throw new RuntimeException(String.format("Failed to create user %s", userUrn), e);
    } finally {
      try {
        if (response != null) {
          response.close();
        }
      } catch (Exception e) {
        log.error("Failed to close http response", e);
      }
    }
  }

  /** Call the Auth Service to reset credentials for a native DataHub user. */
  public boolean resetNativeUserCredentials(
      @Nonnull final String userUrn,
      @Nonnull final String password,
      @Nonnull final String resetToken) {
    Objects.requireNonNull(userUrn, "userUrn must not be null");
    Objects.requireNonNull(password, "password must not be null");
    Objects.requireNonNull(resetToken, "reset token must not be null");
    CloseableHttpResponse response = null;

    try {

      final String protocol = this.metadataServiceUseSsl ? "https" : "http";
      final HttpPost request =
          new HttpPost(
              String.format(
                  "%s://%s:%s/%s",
                  protocol,
                  this.metadataServiceHost,
                  this.metadataServicePort,
                  RESET_NATIVE_USER_CREDENTIALS_ENDPOINT));

      // Build JSON request to verify credentials for a native user.
      final ObjectMapper objectMapper = new ObjectMapper();
      final ObjectNode objectNode = objectMapper.createObjectNode();
      objectNode.put(USER_URN_FIELD, userUrn);
      objectNode.put(PASSWORD_FIELD, password);
      objectNode.put(RESET_TOKEN_FIELD, resetToken);
      final String json =
          objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(objectNode);
      request.setEntity(new StringEntity(json, StandardCharsets.UTF_8));

      // Add authorization header with DataHub frontend system id and secret.
      request.addHeader(Http.HeaderNames.AUTHORIZATION, this.systemAuthentication.getCredentials());

      response = httpClient.execute(request);
      final HttpEntity entity = response.getEntity();
      if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK && entity != null) {
        // Successfully generated a token for the User
        final String jsonStr = EntityUtils.toString(entity);
        return getAreNativeUserCredentialsResetFromJson(jsonStr);
      } else {
        throw new RuntimeException(
            String.format(
                "Bad response from the Metadata Service: %s %s",
                response.getStatusLine().toString(), response.getEntity().toString()));
      }
    } catch (Exception e) {
      throw new RuntimeException("Failed to reset credentials for user", e);
    } finally {
      try {
        if (response != null) {
          response.close();
        }
      } catch (Exception e) {
        log.error("Failed to close http response", e);
      }
    }
  }

  /** Call the Auth Service to verify the credentials for a native Datahub user. */
  public boolean verifyNativeUserCredentials(
      @Nonnull final String userUrn, @Nonnull final String password) {
    Objects.requireNonNull(userUrn, "userUrn must not be null");
    Objects.requireNonNull(password, "password must not be null");
    CloseableHttpResponse response = null;

    try {

      final String protocol = this.metadataServiceUseSsl ? "https" : "http";
      final HttpPost request =
          new HttpPost(
              String.format(
                  "%s://%s:%s/%s",
                  protocol,
                  this.metadataServiceHost,
                  this.metadataServicePort,
                  VERIFY_NATIVE_USER_CREDENTIALS_ENDPOINT));

      // Build JSON request to verify credentials for a native user.
      final ObjectMapper objectMapper = new ObjectMapper();
      final ObjectNode objectNode = objectMapper.createObjectNode();
      objectNode.put(USER_URN_FIELD, userUrn);
      objectNode.put(PASSWORD_FIELD, password);
      final String json =
          objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(objectNode);
      request.setEntity(new StringEntity(json, StandardCharsets.UTF_8));

      // Add authorization header with DataHub frontend system id and secret.
      request.addHeader(Http.HeaderNames.AUTHORIZATION, this.systemAuthentication.getCredentials());

      response = httpClient.execute(request);
      final HttpEntity entity = response.getEntity();
      if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK && entity != null) {
        // Successfully generated a token for the User
        final String jsonStr = EntityUtils.toString(entity);
        return getDoesPasswordMatchFromJson(jsonStr);
      } else {
        throw new RuntimeException(
            String.format(
                "Bad response from the Metadata Service: %s %s",
                response.getStatusLine().toString(), response.getEntity().toString()));
      }
    } catch (Exception e) {
      throw new RuntimeException("Failed to verify credentials for user", e);
    } finally {
      try {
        if (response != null) {
          response.close();
        }
      } catch (Exception e) {
        log.error("Failed to close http response", e);
      }
    }
  }

  /** Call the Auth Service to track an analytics event */
  public void track(@Nonnull final String event) {
    Objects.requireNonNull(event, "event must not be null");
    CloseableHttpResponse response = null;

    try {
      final String protocol = this.metadataServiceUseSsl ? "https" : "http";
      final HttpPost request =
          new HttpPost(
              String.format(
                  "%s://%s:%s/%s",
                  protocol, this.metadataServiceHost, this.metadataServicePort, TRACK_ENDPOINT));

      // Build JSON request to track event.
      request.setEntity(new StringEntity(event, StandardCharsets.UTF_8));

      // Add authorization header with DataHub frontend system id and secret.
      request.addHeader(Http.HeaderNames.AUTHORIZATION, this.systemAuthentication.getCredentials());

      response = httpClient.execute(request);
      final HttpEntity entity = response.getEntity();

      if (response.getStatusLine().getStatusCode() != HttpStatus.SC_OK || entity == null) {
        throw new RuntimeException(
            String.format(
                "Bad response from the Metadata Service: %s %s",
                response.getStatusLine().toString(), response.getEntity().toString()));
      }
    } catch (Exception e) {
      throw new RuntimeException("Failed to track event", e);
    } finally {
      try {
        if (response != null) {
          response.close();
        }
      } catch (Exception e) {
        log.error("Failed to close http response", e);
      }
    }
  }

  private String getAccessTokenFromJson(final String jsonStr) {
    ObjectMapper mapper = new ObjectMapper();
    try {
      return mapper.readTree(jsonStr).get(ACCESS_TOKEN_FIELD).asText();
    } catch (Exception e) {
      // Do not log the raw json in case it contains access token.
      throw new IllegalArgumentException("Failed to parse JSON received from the MetadataService!");
    }
  }

  private boolean getIsNativeUserCreatedFromJson(final String jsonStr) {
    ObjectMapper mapper = new ObjectMapper();
    try {
      return mapper.readTree(jsonStr).get(IS_NATIVE_USER_CREATED_FIELD).asBoolean();
    } catch (Exception e) {
      throw new IllegalArgumentException("Failed to parse JSON received from the MetadataService!");
    }
  }

  private boolean getAreNativeUserCredentialsResetFromJson(final String jsonStr) {
    ObjectMapper mapper = new ObjectMapper();
    try {
      return mapper.readTree(jsonStr).get(ARE_NATIVE_USER_CREDENTIALS_RESET_FIELD).asBoolean();
    } catch (Exception e) {
      throw new IllegalArgumentException("Failed to parse JSON received from the MetadataService!");
    }
  }

  private boolean getDoesPasswordMatchFromJson(final String jsonStr) {
    ObjectMapper mapper = new ObjectMapper();
    try {
      return mapper.readTree(jsonStr).get(DOES_PASSWORD_MATCH_FIELD).asBoolean();
    } catch (Exception e) {
      throw new IllegalArgumentException("Failed to parse JSON received from the MetadataService!");
    }
  }
}