text: /* -*- Mode: C; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */
  
  /*  Fluent Bit
   *  ==========
   *  Copyright (C) 2019-2020 The Fluent Bit Authors
   *  Copyright (C) 2015-2018 Treasure Data Inc.
   *
   *  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 <fluent-bit/flb_info.h>
  #include <fluent-bit/flb_mem.h>
  #include <fluent-bit/flb_log.h>
  #include <fluent-bit/flb_utils.h>
  #include <fluent-bit/flb_oauth2.h>
  #include <fluent-bit/flb_upstream.h>
  #include <fluent-bit/flb_http_client.h>
  
  #include <jsmn/jsmn.h>
  
  #define free_temporary_buffers()                 \
      if (prot) {                                 \
          flb_free(prot);                         \
      }                                           \
      if (host) {                                 \
          flb_free(host);                         \
      }                                           \
      if (port) {                                 \
          flb_free(port);                         \
      }                                           \
      if (uri) {                                  \
          flb_free(uri);                          \
      }
  
  static inline int key_cmp(const char *str, int len, const char *cmp) {
  
      if (strlen(cmp) != len) {
          return -1;
      }
  
      return strncasecmp(str, cmp, len);
  }
  
  int flb_oauth2_parse_json_response(const char *json_data, size_t json_size,
                                     struct flb_oauth2 *ctx)
  {
      int i;
      int ret;
      int key_len;
      int val_len;
      int tokens_size = 32;
      const char *key;
      const char *val;
      jsmn_parser parser;
      jsmntok_t *t;
      jsmntok_t *tokens;
  
      jsmn_init(&parser);
      tokens = flb_calloc(1, sizeof(jsmntok_t) * tokens_size);
      if (!tokens) {
          flb_errno();
          return -1;
      }
  
      ret = jsmn_parse(&parser, json_data, json_size, tokens, tokens_size);
      if (ret <= 0) {
          flb_error("[oauth2] cannot parse payload:\n%s", json_data);
          flb_free(tokens);
          return -1;
      }
  
      t = &tokens[0];
      if (t->type != JSMN_OBJECT) {
          flb_error("[oauth2] invalid JSON response:\n%s", json_data);
          flb_free(tokens);
          return -1;
      }
  
      /* Parse JSON tokens */
      for (i = 1; i < ret; i++) {
          t = &tokens[i];
  
          if (t->type != JSMN_STRING) {
              continue;
          }
  
          if (t->start == -1 || t->end == -1 || (t->start == 0 && t->end == 0)){
              break;
          }
  
          /* Key */
          key = json_data + t->start;
          key_len = (t->end - t->start);
  
          /* Value */
          i++;
          t = &tokens[i];
          val = json_data + t->start;
          val_len = (t->end - t->start);
  
          if (key_cmp(key, key_len, "access_token") == 0) {
              ctx->access_token = flb_sds_create_len(val, val_len);
          }
          else if (key_cmp(key, key_len, "token_type") == 0) {
              ctx->token_type = flb_sds_create_len(val, val_len);
          }
          else if (key_cmp(key, key_len, "expires_in") == 0) {
              ctx->expires_in = atol(val);
          }
      }
  
      flb_free(tokens);
      if (!ctx->access_token || !ctx->token_type || ctx->expires_in < 60) {
          flb_sds_destroy(ctx->access_token);
          flb_sds_destroy(ctx->token_type);
          ctx->expires_in = 0;
          return -1;
      }
  
      return 0;
  }
  
  struct flb_oauth2 *flb_oauth2_create(struct flb_config *config,
                                       const char *auth_url, int expire_sec)
  {
      int ret;
      char *prot = NULL;
      char *host = NULL;
      char *port = NULL;
      char *uri = NULL;
      struct flb_oauth2 *ctx;
  
      /* allocate context */
      ctx = flb_calloc(1, sizeof(struct flb_oauth2));
      if (!ctx) {
          flb_errno();
          return NULL;
      }
  
      /* register token url */
      ctx->auth_url = flb_sds_create(auth_url);
      if (!ctx->auth_url) {
          flb_errno();
          flb_free(ctx);
          return NULL;
      }
  
      /* default payload size to 1kb */
      ctx->payload = flb_sds_create_size(1024);
      if (!ctx->payload) {
          flb_errno();
          flb_oauth2_destroy(ctx);
          return NULL;
      }
  
      ctx->issued = time(NULL);
      ctx->expires = ctx->issued + expire_sec;
  
      /* Parse and split URL */
      ret = flb_utils_url_split(auth_url, &prot, &host, &port, &uri);
      if (ret == -1) {
          flb_error("[oauth2] invalid URL: %s", auth_url);
          goto error;
      }
  
      if (!prot || strcmp(prot, "https") != 0) {
          flb_error("[oauth2] invalid endpoint protocol: %s", auth_url);
          goto error;
      }
  
      if (!host) {
          flb_error("[oauth2] invalid URL host: %s", auth_url);
          goto error;
      }
  
      /* Populate context */
      ctx->host = flb_sds_create(host);
      if (!ctx->host) {
          flb_errno();
          goto error;
      }
      if (port) {
          ctx->port = flb_sds_create(port);
      }
      else {
          ctx->port = flb_sds_create(FLB_OAUTH2_PORT);
      }
      if (!ctx->port) {
          flb_errno();
          goto error;
      }
      ctx->uri = flb_sds_create(uri);
      if (!ctx->uri) {
          flb_errno();
          goto error;
      }
  
      /* Create TLS context */
      ctx->tls = flb_tls_create(FLB_TRUE,  /* verify */
                                -1,        /* debug */
                                NULL,      /* vhost */
                                NULL,      /* ca_path */
                                NULL,      /* ca_file */
                                NULL,      /* crt_file */
                                NULL,      /* key_file */
                                NULL);     /* key_passwd */
      if (!ctx->tls) {
          flb_error("[oauth2] error initializing TLS context");
          goto error;
      }
  
      /* Create Upstream context */
      ctx->u = flb_upstream_create_url(config, auth_url,
                                       FLB_IO_TLS, ctx->tls);
      if (!ctx->u) {
          flb_error("[oauth2] error creating upstream context");
          goto error;
      }
  
      /* Remove Upstream Async flag */
      ctx->u->flags &= ~(FLB_IO_ASYNC);
  
      free_temporary_buffers();
      return ctx;
  
   error:
      free_temporary_buffers();
      flb_oauth2_destroy(ctx);
  
      return NULL;
  }
  
  /* Append a key/value to the request body */
  int flb_oauth2_payload_append(struct flb_oauth2 *ctx,
                                const char *key_str, int key_len,
                                const char *val_str, int val_len)
  {
      int size;
      flb_sds_t tmp;
  
      if (key_len == -1) {
          key_len = strlen(key_str);
      }
      if (val_len == -1) {
          val_len = strlen(val_str);
      }
  
      /*
       * Make sure we have enough space in the sds buffer, otherwise
       * add more capacity (so further flb_sds_cat calls do not
       * realloc().
       */
      size = key_len + val_len + 2;
      if (flb_sds_avail(ctx->payload) < size) {
          tmp = flb_sds_increase(ctx->payload, size);
          if (!tmp) {
              flb_errno();
              return -1;
          }
  
          if (tmp != ctx->payload) {
              ctx->payload = tmp;
          }
      }
  
      if (flb_sds_len(ctx->payload) > 0) {
          flb_sds_cat(ctx->payload, "&", 1);
      }
  
      /* Append key and value */
      flb_sds_cat(ctx->payload, key_str, key_len);
      flb_sds_cat(ctx->payload, "=", 1);
      flb_sds_cat(ctx->payload, val_str, val_len);
  
      return 0;
  }
  
  void flb_oauth2_destroy(struct flb_oauth2 *ctx)
  {
      flb_sds_destroy(ctx->auth_url);
      flb_sds_destroy(ctx->payload);
  
      flb_sds_destroy(ctx->host);
      flb_sds_destroy(ctx->port);
      flb_sds_destroy(ctx->uri);
  
      flb_sds_destroy(ctx->access_token);
      flb_sds_destroy(ctx->token_type);
  
      flb_upstream_destroy(ctx->u);
      flb_tls_destroy(ctx->tls);
  
      flb_free(ctx);
  }
  
  char *flb_oauth2_token_get(struct flb_oauth2 *ctx)
  {
      int ret;
      size_t b_sent;
      time_t now;
      struct flb_upstream_conn *u_conn;
      struct flb_http_client *c;
  
      now = time(NULL);
      if (ctx->access_token) {
          /* validate expired token */
          if (ctx->expires < now && flb_sds_len(ctx->access_token) > 0) {
              return ctx->access_token;
          }
      }
  
      /* Get Token and store it in the context */
      u_conn = flb_upstream_conn_get(ctx->u);
      if (!u_conn) {
          ctx->u->flags |= FLB_IO_IPV6;
          u_conn = flb_upstream_conn_get(ctx->u);
          if (!u_conn) {
              flb_error("[oauth2] could not get an upstream connection to %s:%i",
                        ctx->u->tcp_host, ctx->u->tcp_port);
              ctx->u->flags &= ~FLB_IO_IPV6;
              return NULL;
          }
      }
  
      /* Create HTTP client context */
      c = flb_http_client(u_conn, FLB_HTTP_POST, ctx->uri,
                          ctx->payload, flb_sds_len(ctx->payload),
                          ctx->host, atoi(ctx->port),
                          NULL, 0);
      if (!c) {
          flb_error("[oauth2] error creating HTTP client context");
          flb_upstream_conn_release(u_conn);
          return NULL;
      }
  
      /* Append HTTP Header */
      flb_http_add_header(c,
                          FLB_HTTP_HEADER_CONTENT_TYPE,
                          sizeof(FLB_HTTP_HEADER_CONTENT_TYPE) -1,
                          FLB_OAUTH2_HTTP_ENCODING,
                          sizeof(FLB_OAUTH2_HTTP_ENCODING) - 1);
  
      /* Issue request */
      ret = flb_http_do(c, &b_sent);
      if (ret != 0) {
          flb_warn("[oauth2] cannot issue request, http_do=%i", ret);
      }
      else {
          flb_info("[oauth2] HTTP Status=%i", c->resp.status);
          if (c->resp.payload_size > 0) {
              if (c->resp.status == 200) {
                  flb_debug("[oauth2] payload:\n%s", c->resp.payload);
              }
              else {
                  flb_info("[oauth2] payload:\n%s", c->resp.payload);
              }
          }
      }
  
      /* Extract token */
      if (c->resp.payload_size > 0 && c->resp.status == 200) {
          ret = flb_oauth2_parse_json_response(c->resp.payload,
                                               c->resp.payload_size, ctx);
          if (ret == 0) {
              flb_info("[oauth2] access token from '%s:%s' retrieved",
                       ctx->host, ctx->port);
              flb_http_client_destroy(c);
              return ctx->access_token;
          }
      }
  
      flb_http_client_destroy(c);
      return NULL;
  }
  
  int flb_oauth2_token_len(struct flb_oauth2 *ctx)
  {
      if (!ctx->access_token) {
          return -1;
      }
  
      return flb_sds_len(ctx->access_token);
  }
  
  int flb_oauth2_token_expired(struct flb_oauth2 *ctx)
  {
      time_t now;
  
      if (!ctx->access_token) {
          return FLB_TRUE;
      }
  
      now = time(NULL);
      if (ctx->expires <= now) {
          return FLB_TRUE;
      }
  
      return FLB_FALSE;
  }
  