/*
 *
 * Copyright 2015 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include "src/core/lib/security/credentials/jwt/json_token.h"

#include <openssl/evp.h>
#include <string.h>

#include <grpc/grpc_security.h>
#include <grpc/slice.h>
#include <grpc/support/alloc.h>
#include <grpc/support/log.h>

#include "src/core/lib/json/json.h"
#include "src/core/lib/security/credentials/oauth2/oauth2_credentials.h"
#include "src/core/lib/slice/b64.h"
#include "src/core/lib/slice/slice_internal.h"
#include "test/core/util/test_config.h"

/* This JSON key was generated with the GCE console and revoked immediately.
   The identifiers have been changed as well.
   Maximum size for a string literal is 509 chars in C89, yay!  */
static const char test_json_key_str_part1[] =
    "{ \"private_key\": \"-----BEGIN PRIVATE KEY-----"
    "\\nMIICeAIBADANBgkqhkiG9w0BAQEFAASCAmIwggJeAgEAAoGBAOEvJsnoHnyHkXcp\\n7mJE"
    "qg"
    "WGjiw71NfXByguekSKho65FxaGbsnSM9SMQAqVk7Q2rG+I0OpsT0LrWQtZ\\nyjSeg/"
    "rWBQvS4hle4LfijkP3J5BG+"
    "IXDMP8RfziNRQsenAXDNPkY4kJCvKux2xdD\\nOnVF6N7dL3nTYZg+"
    "uQrNsMTz9UxVAgMBAAECgYEAzbLewe1xe9vy+2GoSsfib+28\\nDZgSE6Bu/"
    "zuFoPrRc6qL9p2SsnV7txrunTyJkkOnPLND9ABAXybRTlcVKP/sGgza\\n/"
    "8HpCqFYM9V8f34SBWfD4fRFT+n/"
    "73cfRUtGXdXpseva2lh8RilIQfPhNZAncenU\\ngqXjDvpkypEusgXAykECQQD+";
static const char test_json_key_str_part2[] =
    "53XxNVnxBHsYb+AYEfklR96yVi8HywjVHP34+OQZ\\nCslxoHQM8s+"
    "dBnjfScLu22JqkPv04xyxmt0QAKm9+vTdAkEA4ib7YvEAn2jXzcCI\\nEkoy2L/"
    "XydR1GCHoacdfdAwiL2npOdnbvi4ZmdYRPY1LSTO058tQHKVXV7NLeCa3\\nAARh2QJBAMKeDA"
    "G"
    "W303SQv2cZTdbeaLKJbB5drz3eo3j7dDKjrTD9JupixFbzcGw\\n8FZi5c8idxiwC36kbAL6Hz"
    "A"
    "ZoX+ofI0CQE6KCzPJTtYNqyShgKAZdJ8hwOcvCZtf\\n6z8RJm0+"
    "6YBd38lfh5j8mZd7aHFf6I17j5AQY7oPEc47TjJj/"
    "5nZ68ECQQDvYuI3\\nLyK5fS8g0SYbmPOL9TlcHDOqwG0mrX9qpg5DC2fniXNSrrZ64GTDKdzZ"
    "Y"
    "Ap6LI9W\\nIqv4vr6y38N79TTC\\n-----END PRIVATE KEY-----\\n\", ";
static const char test_json_key_str_part3[] =
    "\"private_key_id\": \"e6b5137873db8d2ef81e06a47289e6434ec8a165\", "
    "\"client_email\": "
    "\"777-abaslkan11hlb6nmim3bpspl31ud@developer.gserviceaccount."
    "com\", \"client_id\": "
    "\"777-abaslkan11hlb6nmim3bpspl31ud.apps.googleusercontent."
    "com\", \"type\": \"service_account\" }";

/* Test refresh token. */
static const char test_refresh_token_str[] =
    "{ \"client_id\": \"32555999999.apps.googleusercontent.com\","
    "  \"client_secret\": \"EmssLNjJy1332hD4KFsecret\","
    "  \"refresh_token\": \"1/Blahblasj424jladJDSGNf-u4Sua3HDA2ngjd42\","
    "  \"type\": \"authorized_user\"}";

static const char test_scope[] = "myperm1 myperm2";

static const char test_service_url[] = "https://foo.com/foo.v1";

static char* test_json_key_str(const char* bad_part3) {
  const char* part3 =
      bad_part3 != nullptr ? bad_part3 : test_json_key_str_part3;
  size_t result_len = strlen(test_json_key_str_part1) +
                      strlen(test_json_key_str_part2) + strlen(part3);
  char* result = static_cast<char*>(gpr_malloc(result_len + 1));
  char* current = result;
  strcpy(result, test_json_key_str_part1);
  current += strlen(test_json_key_str_part1);
  strcpy(current, test_json_key_str_part2);
  current += strlen(test_json_key_str_part2);
  strcpy(current, part3);
  return result;
}

static void test_parse_json_key_success(void) {
  char* json_string = test_json_key_str(nullptr);
  grpc_auth_json_key json_key =
      grpc_auth_json_key_create_from_string(json_string);
  GPR_ASSERT(grpc_auth_json_key_is_valid(&json_key));
  GPR_ASSERT(json_key.type != nullptr &&
             strcmp(json_key.type, "service_account") == 0);
  GPR_ASSERT(json_key.private_key_id != nullptr &&
             strcmp(json_key.private_key_id,
                    "e6b5137873db8d2ef81e06a47289e6434ec8a165") == 0);
  GPR_ASSERT(json_key.client_id != nullptr &&
             strcmp(json_key.client_id,
                    "777-abaslkan11hlb6nmim3bpspl31ud.apps."
                    "googleusercontent.com") == 0);
  GPR_ASSERT(json_key.client_email != nullptr &&
             strcmp(json_key.client_email,
                    "777-abaslkan11hlb6nmim3bpspl31ud@developer."
                    "gserviceaccount.com") == 0);
  GPR_ASSERT(json_key.private_key != nullptr);
  gpr_free(json_string);
  grpc_auth_json_key_destruct(&json_key);
}

static void test_parse_json_key_failure_bad_json(void) {
  const char non_closing_part3[] =
      "\"private_key_id\": \"e6b5137873db8d2ef81e06a47289e6434ec8a165\", "
      "\"client_email\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud@developer.gserviceaccount."
      "com\", \"client_id\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud.apps.googleusercontent."
      "com\", \"type\": \"service_account\" ";
  char* json_string = test_json_key_str(non_closing_part3);
  grpc_auth_json_key json_key =
      grpc_auth_json_key_create_from_string(json_string);
  GPR_ASSERT(!grpc_auth_json_key_is_valid(&json_key));
  gpr_free(json_string);
  grpc_auth_json_key_destruct(&json_key);
}

static void test_parse_json_key_failure_no_type(void) {
  const char no_type_part3[] =
      "\"private_key_id\": \"e6b5137873db8d2ef81e06a47289e6434ec8a165\", "
      "\"client_email\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud@developer.gserviceaccount."
      "com\", \"client_id\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud.apps.googleusercontent."
      "com\" }";
  char* json_string = test_json_key_str(no_type_part3);
  grpc_auth_json_key json_key =
      grpc_auth_json_key_create_from_string(json_string);
  GPR_ASSERT(!grpc_auth_json_key_is_valid(&json_key));
  gpr_free(json_string);
  grpc_auth_json_key_destruct(&json_key);
}

static void test_parse_json_key_failure_no_client_id(void) {
  const char no_client_id_part3[] =
      "\"private_key_id\": \"e6b5137873db8d2ef81e06a47289e6434ec8a165\", "
      "\"client_email\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud@developer.gserviceaccount."
      "com\", "
      "\"type\": \"service_account\" }";
  char* json_string = test_json_key_str(no_client_id_part3);
  grpc_auth_json_key json_key =
      grpc_auth_json_key_create_from_string(json_string);
  GPR_ASSERT(!grpc_auth_json_key_is_valid(&json_key));
  gpr_free(json_string);
  grpc_auth_json_key_destruct(&json_key);
}

static void test_parse_json_key_failure_no_client_email(void) {
  const char no_client_email_part3[] =
      "\"private_key_id\": \"e6b5137873db8d2ef81e06a47289e6434ec8a165\", "
      "\"client_id\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud.apps.googleusercontent."
      "com\", \"type\": \"service_account\" }";
  char* json_string = test_json_key_str(no_client_email_part3);
  grpc_auth_json_key json_key =
      grpc_auth_json_key_create_from_string(json_string);
  GPR_ASSERT(!grpc_auth_json_key_is_valid(&json_key));
  gpr_free(json_string);
  grpc_auth_json_key_destruct(&json_key);
}

static void test_parse_json_key_failure_no_private_key_id(void) {
  const char no_private_key_id_part3[] =
      "\"client_email\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud@developer.gserviceaccount."
      "com\", \"client_id\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud.apps.googleusercontent."
      "com\", \"type\": \"service_account\" }";
  char* json_string = test_json_key_str(no_private_key_id_part3);
  grpc_auth_json_key json_key =
      grpc_auth_json_key_create_from_string(json_string);
  GPR_ASSERT(!grpc_auth_json_key_is_valid(&json_key));
  gpr_free(json_string);
  grpc_auth_json_key_destruct(&json_key);
}

static void test_parse_json_key_failure_no_private_key(void) {
  const char no_private_key_json_string[] =
      "{ \"private_key_id\": \"e6b5137873db8d2ef81e06a47289e6434ec8a165\", "
      "\"client_email\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud@developer.gserviceaccount."
      "com\", \"client_id\": "
      "\"777-abaslkan11hlb6nmim3bpspl31ud.apps.googleusercontent."
      "com\", \"type\": \"service_account\" }";
  grpc_auth_json_key json_key =
      grpc_auth_json_key_create_from_string(no_private_key_json_string);
  GPR_ASSERT(!grpc_auth_json_key_is_valid(&json_key));
  grpc_auth_json_key_destruct(&json_key);
}

static grpc_json* parse_json_part_from_jwt(const char* str, size_t len,
                                           char** scratchpad) {
  grpc_core::ExecCtx exec_ctx;
  char* b64;
  char* decoded;
  grpc_json* json;
  grpc_slice slice;
  b64 = static_cast<char*>(gpr_malloc(len + 1));
  strncpy(b64, str, len);
  b64[len] = '\0';
  slice = grpc_base64_decode(b64, 1);
  GPR_ASSERT(!GRPC_SLICE_IS_EMPTY(slice));
  decoded = static_cast<char*>(gpr_malloc(GRPC_SLICE_LENGTH(slice) + 1));
  strncpy(decoded, reinterpret_cast<const char*> GRPC_SLICE_START_PTR(slice),
          GRPC_SLICE_LENGTH(slice));
  decoded[GRPC_SLICE_LENGTH(slice)] = '\0';
  json = grpc_json_parse_string(decoded);
  gpr_free(b64);
  *scratchpad = decoded;
  grpc_slice_unref(slice);

  return json;
}

static void check_jwt_header(grpc_json* header) {
  grpc_json* ptr;
  grpc_json* alg = nullptr;
  grpc_json* typ = nullptr;
  grpc_json* kid = nullptr;

  for (ptr = header->child; ptr; ptr = ptr->next) {
    if (strcmp(ptr->key, "alg") == 0) {
      alg = ptr;
    } else if (strcmp(ptr->key, "typ") == 0) {
      typ = ptr;
    } else if (strcmp(ptr->key, "kid") == 0) {
      kid = ptr;
    }
  }
  GPR_ASSERT(alg != nullptr);
  GPR_ASSERT(alg->type == GRPC_JSON_STRING);
  GPR_ASSERT(strcmp(alg->value, "RS256") == 0);

  GPR_ASSERT(typ != nullptr);
  GPR_ASSERT(typ->type == GRPC_JSON_STRING);
  GPR_ASSERT(strcmp(typ->value, "JWT") == 0);

  GPR_ASSERT(kid != nullptr);
  GPR_ASSERT(kid->type == GRPC_JSON_STRING);
  GPR_ASSERT(strcmp(kid->value, "e6b5137873db8d2ef81e06a47289e6434ec8a165") ==
             0);
}

static void check_jwt_claim(grpc_json* claim, const char* expected_audience,
                            const char* expected_scope) {
  gpr_timespec expiration = gpr_time_0(GPR_CLOCK_REALTIME);
  gpr_timespec issue_time = gpr_time_0(GPR_CLOCK_REALTIME);
  gpr_timespec parsed_lifetime;
  grpc_json* iss = nullptr;
  grpc_json* scope = nullptr;
  grpc_json* aud = nullptr;
  grpc_json* exp = nullptr;
  grpc_json* iat = nullptr;
  grpc_json* sub = nullptr;
  grpc_json* ptr;

  for (ptr = claim->child; ptr; ptr = ptr->next) {
    if (strcmp(ptr->key, "iss") == 0) {
      iss = ptr;
    } else if (strcmp(ptr->key, "sub") == 0) {
      sub = ptr;
    } else if (strcmp(ptr->key, "scope") == 0) {
      scope = ptr;
    } else if (strcmp(ptr->key, "aud") == 0) {
      aud = ptr;
    } else if (strcmp(ptr->key, "exp") == 0) {
      exp = ptr;
    } else if (strcmp(ptr->key, "iat") == 0) {
      iat = ptr;
    }
  }

  GPR_ASSERT(iss != nullptr);
  GPR_ASSERT(iss->type == GRPC_JSON_STRING);
  GPR_ASSERT(
      strcmp(
          iss->value,
          "777-abaslkan11hlb6nmim3bpspl31ud@developer.gserviceaccount.com") ==
      0);

  if (expected_scope != nullptr) {
    GPR_ASSERT(scope != nullptr);
    GPR_ASSERT(sub == nullptr);
    GPR_ASSERT(scope->type == GRPC_JSON_STRING);
    GPR_ASSERT(strcmp(scope->value, expected_scope) == 0);
  } else {
    /* Claims without scope must have a sub. */
    GPR_ASSERT(scope == nullptr);
    GPR_ASSERT(sub != nullptr);
    GPR_ASSERT(sub->type == GRPC_JSON_STRING);
    GPR_ASSERT(strcmp(iss->value, sub->value) == 0);
  }

  GPR_ASSERT(aud != nullptr);
  GPR_ASSERT(aud->type == GRPC_JSON_STRING);
  GPR_ASSERT(strcmp(aud->value, expected_audience) == 0);

  GPR_ASSERT(exp != nullptr);
  GPR_ASSERT(exp->type == GRPC_JSON_NUMBER);
  expiration.tv_sec = strtol(exp->value, nullptr, 10);

  GPR_ASSERT(iat != nullptr);
  GPR_ASSERT(iat->type == GRPC_JSON_NUMBER);
  issue_time.tv_sec = strtol(iat->value, nullptr, 10);

  parsed_lifetime = gpr_time_sub(expiration, issue_time);
  GPR_ASSERT(parsed_lifetime.tv_sec == grpc_max_auth_token_lifetime().tv_sec);
}

static void check_jwt_signature(const char* b64_signature, RSA* rsa_key,
                                const char* signed_data,
                                size_t signed_data_size) {
  grpc_core::ExecCtx exec_ctx;

  EVP_MD_CTX* md_ctx = EVP_MD_CTX_create();
  EVP_PKEY* key = EVP_PKEY_new();

  grpc_slice sig = grpc_base64_decode(b64_signature, 1);
  GPR_ASSERT(!GRPC_SLICE_IS_EMPTY(sig));
  GPR_ASSERT(GRPC_SLICE_LENGTH(sig) == 128);

  GPR_ASSERT(md_ctx != nullptr);
  GPR_ASSERT(key != nullptr);
  EVP_PKEY_set1_RSA(key, rsa_key);

  GPR_ASSERT(
      EVP_DigestVerifyInit(md_ctx, nullptr, EVP_sha256(), nullptr, key) == 1);
  GPR_ASSERT(EVP_DigestVerifyUpdate(md_ctx, signed_data, signed_data_size) ==
             1);
  GPR_ASSERT(EVP_DigestVerifyFinal(md_ctx, GRPC_SLICE_START_PTR(sig),
                                   GRPC_SLICE_LENGTH(sig)) == 1);

  grpc_slice_unref_internal(sig);
  if (key != nullptr) EVP_PKEY_free(key);
  if (md_ctx != nullptr) EVP_MD_CTX_destroy(md_ctx);
}

static char* service_account_creds_jwt_encode_and_sign(
    const grpc_auth_json_key* key) {
  return grpc_jwt_encode_and_sign(key, GRPC_JWT_OAUTH2_AUDIENCE,
                                  grpc_max_auth_token_lifetime(), test_scope);
}

static char* jwt_creds_jwt_encode_and_sign(const grpc_auth_json_key* key) {
  return grpc_jwt_encode_and_sign(key, test_service_url,
                                  grpc_max_auth_token_lifetime(), nullptr);
}

static void service_account_creds_check_jwt_claim(grpc_json* claim) {
  check_jwt_claim(claim, GRPC_JWT_OAUTH2_AUDIENCE, test_scope);
}

static void jwt_creds_check_jwt_claim(grpc_json* claim) {
  check_jwt_claim(claim, test_service_url, nullptr);
}

static void test_jwt_encode_and_sign(
    char* (*jwt_encode_and_sign_func)(const grpc_auth_json_key*),
    void (*check_jwt_claim_func)(grpc_json*)) {
  char* json_string = test_json_key_str(nullptr);
  grpc_json* parsed_header = nullptr;
  grpc_json* parsed_claim = nullptr;
  char* scratchpad;
  grpc_auth_json_key json_key =
      grpc_auth_json_key_create_from_string(json_string);
  const char* b64_signature;
  size_t offset = 0;
  char* jwt = jwt_encode_and_sign_func(&json_key);
  const char* dot = strchr(jwt, '.');
  GPR_ASSERT(dot != nullptr);
  parsed_header = parse_json_part_from_jwt(jwt, static_cast<size_t>(dot - jwt),
                                           &scratchpad);
  GPR_ASSERT(parsed_header != nullptr);
  check_jwt_header(parsed_header);
  offset = static_cast<size_t>(dot - jwt) + 1;
  grpc_json_destroy(parsed_header);
  gpr_free(scratchpad);

  dot = strchr(jwt + offset, '.');
  GPR_ASSERT(dot != nullptr);
  parsed_claim = parse_json_part_from_jwt(
      jwt + offset, static_cast<size_t>(dot - (jwt + offset)), &scratchpad);
  GPR_ASSERT(parsed_claim != nullptr);
  check_jwt_claim_func(parsed_claim);
  offset = static_cast<size_t>(dot - jwt) + 1;
  grpc_json_destroy(parsed_claim);
  gpr_free(scratchpad);

  dot = strchr(jwt + offset, '.');
  GPR_ASSERT(dot == nullptr); /* no more part. */
  b64_signature = jwt + offset;
  check_jwt_signature(b64_signature, json_key.private_key, jwt, offset - 1);

  gpr_free(json_string);
  grpc_auth_json_key_destruct(&json_key);
  gpr_free(jwt);
}

static void test_service_account_creds_jwt_encode_and_sign(void) {
  test_jwt_encode_and_sign(service_account_creds_jwt_encode_and_sign,
                           service_account_creds_check_jwt_claim);
}

static void test_jwt_creds_jwt_encode_and_sign(void) {
  test_jwt_encode_and_sign(jwt_creds_jwt_encode_and_sign,
                           jwt_creds_check_jwt_claim);
}

static void test_parse_refresh_token_success(void) {
  grpc_auth_refresh_token refresh_token =
      grpc_auth_refresh_token_create_from_string(test_refresh_token_str);
  GPR_ASSERT(grpc_auth_refresh_token_is_valid(&refresh_token));
  GPR_ASSERT(refresh_token.type != nullptr &&
             (strcmp(refresh_token.type, "authorized_user") == 0));
  GPR_ASSERT(refresh_token.client_id != nullptr &&
             (strcmp(refresh_token.client_id,
                     "32555999999.apps.googleusercontent.com") == 0));
  GPR_ASSERT(
      refresh_token.client_secret != nullptr &&
      (strcmp(refresh_token.client_secret, "EmssLNjJy1332hD4KFsecret") == 0));
  GPR_ASSERT(refresh_token.refresh_token != nullptr &&
             (strcmp(refresh_token.refresh_token,
                     "1/Blahblasj424jladJDSGNf-u4Sua3HDA2ngjd42") == 0));
  grpc_auth_refresh_token_destruct(&refresh_token);
}

static void test_parse_refresh_token_failure_no_type(void) {
  const char refresh_token_str[] =
      "{ \"client_id\": \"32555999999.apps.googleusercontent.com\","
      "  \"client_secret\": \"EmssLNjJy1332hD4KFsecret\","
      "  \"refresh_token\": \"1/Blahblasj424jladJDSGNf-u4Sua3HDA2ngjd42\"}";
  grpc_auth_refresh_token refresh_token =
      grpc_auth_refresh_token_create_from_string(refresh_token_str);
  GPR_ASSERT(!grpc_auth_refresh_token_is_valid(&refresh_token));
}

static void test_parse_refresh_token_failure_no_client_id(void) {
  const char refresh_token_str[] =
      "{ \"client_secret\": \"EmssLNjJy1332hD4KFsecret\","
      "  \"refresh_token\": \"1/Blahblasj424jladJDSGNf-u4Sua3HDA2ngjd42\","
      "  \"type\": \"authorized_user\"}";
  grpc_auth_refresh_token refresh_token =
      grpc_auth_refresh_token_create_from_string(refresh_token_str);
  GPR_ASSERT(!grpc_auth_refresh_token_is_valid(&refresh_token));
}

static void test_parse_refresh_token_failure_no_client_secret(void) {
  const char refresh_token_str[] =
      "{ \"client_id\": \"32555999999.apps.googleusercontent.com\","
      "  \"refresh_token\": \"1/Blahblasj424jladJDSGNf-u4Sua3HDA2ngjd42\","
      "  \"type\": \"authorized_user\"}";
  grpc_auth_refresh_token refresh_token =
      grpc_auth_refresh_token_create_from_string(refresh_token_str);
  GPR_ASSERT(!grpc_auth_refresh_token_is_valid(&refresh_token));
}

static void test_parse_refresh_token_failure_no_refresh_token(void) {
  const char refresh_token_str[] =
      "{ \"client_id\": \"32555999999.apps.googleusercontent.com\","
      "  \"client_secret\": \"EmssLNjJy1332hD4KFsecret\","
      "  \"type\": \"authorized_user\"}";
  grpc_auth_refresh_token refresh_token =
      grpc_auth_refresh_token_create_from_string(refresh_token_str);
  GPR_ASSERT(!grpc_auth_refresh_token_is_valid(&refresh_token));
}

int main(int argc, char** argv) {
  grpc::testing::TestEnvironment env(argc, argv);
  grpc_init();
  test_parse_json_key_success();
  test_parse_json_key_failure_bad_json();
  test_parse_json_key_failure_no_type();
  test_parse_json_key_failure_no_client_id();
  test_parse_json_key_failure_no_client_email();
  test_parse_json_key_failure_no_private_key_id();
  test_parse_json_key_failure_no_private_key();
  test_service_account_creds_jwt_encode_and_sign();
  test_jwt_creds_jwt_encode_and_sign();
  test_parse_refresh_token_success();
  test_parse_refresh_token_failure_no_type();
  test_parse_refresh_token_failure_no_client_id();
  test_parse_refresh_token_failure_no_client_secret();
  test_parse_refresh_token_failure_no_refresh_token();
  grpc_shutdown();
  return 0;
}