Add plugins (#27)

plugins/wasm-cpp/extensions/jwt_auth/plugin.cc

@@ -0,0 +1,397 @@
+// Copyright (c) 2022 Alibaba Group Holding Ltd.
+//
+// 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 "extensions/jwt_auth/plugin.h"
+
+#include <algorithm>
+#include <array>
+#include <cstdint>
+#include <string>
+#include <unordered_set>
+#include <utility>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/str_split.h"
+#include "common/common_util.h"
+#include "common/http_util.h"
+#include "common/json_util.h"
+
+using ::nlohmann::json;
+using ::Wasm::Common::JsonArrayIterate;
+using ::Wasm::Common::JsonGetField;
+using ::Wasm::Common::JsonObjectIterate;
+using ::Wasm::Common::JsonValueAs;
+
+#ifdef NULL_PLUGIN
+
+namespace proxy_wasm {
+namespace null_plugin {
+namespace jwt_auth {
+
+PROXY_WASM_NULL_PLUGIN_REGISTRY
+
+#endif
+namespace {
+constexpr absl::string_view InvalidTokenErrorString =
+    ", error=\"invalid_token\"";
+constexpr uint32_t MaximumUriLength = 256;
+constexpr std::string_view kRcDetailJwtAuthnPrefix = "jwt_authn_access_denied";
+std::string generateRcDetails(std::string_view error_msg) {
+  // Replace space with underscore since RCDetails may be written to access log.
+  // Some log processors assume each log segment is separated by whitespace.
+  return absl::StrCat(kRcDetailJwtAuthnPrefix, "{",
+                      absl::StrJoin(absl::StrSplit(error_msg, ' '), "_"), "}");
+}
+
+}  // namespace
+static RegisterContextFactory register_JwtAuth(CONTEXT_FACTORY(PluginContext),
+                                               ROOT_FACTORY(PluginRootContext));
+
+#define JSON_FIND_FIELD(dict, field)               \
+  auto dict##_##field##_json = dict.find(#field);  \
+  if (dict##_##field##_json == dict.end()) {       \
+    LOG_WARN("can't find '" #field "' in " #dict); \
+    return false;                                  \
+  }
+
+#define JSON_VALUE_AS(type, src, dst, err_msg)                      \
+  auto dst##_v = JsonValueAs<type>(src);                            \
+  if (dst##_v.second != Wasm::Common::JsonParserResultDetail::OK || \
+      !dst##_v.first) {                                             \
+    LOG_WARN(#err_msg);                                             \
+    return false;                                                   \
+  }                                                                 \
+  auto& dst = dst##_v.first.value();
+
+#define JSON_FIELD_VALUE_AS(type, dict, field)                       \
+  JSON_VALUE_AS(type, dict##_##field##_json.value(), dict##_##field, \
+                "'" #field "' field in " #dict "convert to " #type " failed")
+
+bool PluginRootContext::parsePluginConfig(const json& configuration,
+                                          JwtAuthConfigRule& rule) {
+  std::unordered_set<std::string> name_set;
+  if (!JsonArrayIterate(
+          configuration, "consumers", [&](const json& consumer) -> bool {
+            Consumer c;
+            JSON_FIND_FIELD(consumer, name);
+            JSON_FIELD_VALUE_AS(std::string, consumer, name);
+            if (name_set.count(consumer_name) != 0) {
+              LOG_WARN("consumer already exists: " + consumer_name);
+              return false;
+            }
+            c.name = consumer_name;
+            JSON_FIND_FIELD(consumer, jwks);
+            JSON_FIELD_VALUE_AS(std::string, consumer, jwks);
+            c.jwks = google::jwt_verify::Jwks::createFrom(
+                consumer_jwks, google::jwt_verify::Jwks::JWKS);
+            if (c.jwks->getStatus() != Status::Ok) {
+              LOG_WARN(absl::StrFormat(
+                  "jwks is invalid, consumer:%s, status:%s, jwks:%s",
+                  consumer_name,
+                  google::jwt_verify::getStatusString(c.jwks->getStatus()),
+                  consumer_jwks));
+              return false;
+            }
+            std::unordered_map<std::string, std::string> claims;
+            auto consumer_claims_json = consumer.find("claims");
+            if (consumer_claims_json != consumer.end()) {
+              JSON_FIELD_VALUE_AS(Wasm::Common::JsonObject, consumer, claims);
+              if (!JsonObjectIterate(
+                      consumer_claims, [&](std::string key) -> bool {
+                        auto claims_claim_json = consumer_claims.find(key);
+                        JSON_FIELD_VALUE_AS(std::string, claims, claim);
+                        claims.emplace(std::make_pair(
+                            key, Wasm::Common::trim(claims_claim)));
+                        return true;
+                      })) {
+                LOG_WARN("failed to parse 'claims' in consumer: " +
+                         consumer_name);
+                return false;
+              }
+            }
+            auto consumer_issuer_json = consumer.find("issuer");
+            if (consumer_issuer_json != consumer.end()) {
+              JSON_FIELD_VALUE_AS(std::string, consumer, issuer);
+              claims.emplace(
+                  std::make_pair("iss", Wasm::Common::trim(consumer_issuer)));
+            }
+            c.allowd_claims = std::move(claims);
+            std::vector<FromHeader> from_headers;
+            if (!JsonArrayIterate(
+                    consumer, "from_headers",
+                    [&](const json& from_header) -> bool {
+                      JSON_FIND_FIELD(from_header, name);
+                      JSON_FIELD_VALUE_AS(std::string, from_header, name);
+                      JSON_FIND_FIELD(from_header, value_prefix);
+                      JSON_FIELD_VALUE_AS(std::string, from_header,
+                                          value_prefix);
+                      from_headers.push_back(FromHeader{
+                          from_header_name, from_header_value_prefix});
+                      return true;
+                    })) {
+              LOG_WARN("failed to parse 'from_headers' in consumer: " +
+                       consumer_name);
+              return false;
+            }
+            std::vector<std::string> from_params;
+            if (!JsonArrayIterate(consumer, "from_params",
+                                  [&](const json& from_param_json) -> bool {
+                                    JSON_VALUE_AS(std::string, from_param_json,
+                                                  from_param, "invalid item");
+                                    from_params.push_back(from_param);
+                                    return true;
+                                  })) {
+              LOG_WARN("failed to parse 'from_params' in consumer: " +
+                       consumer_name);
+              return false;
+            }
+            std::vector<std::string> from_cookies;
+            if (!JsonArrayIterate(consumer, "from_cookies",
+                                  [&](const json& from_cookie_json) -> bool {
+                                    JSON_VALUE_AS(std::string, from_cookie_json,
+                                                  from_cookie, "invalid item");
+                                    from_cookies.push_back(from_cookie);
+                                    return true;
+                                  })) {
+              LOG_WARN("failed to parse 'from_cookies' in consumer: " +
+                       consumer_name);
+              return false;
+            }
+            if (!from_headers.empty() || !from_params.empty() ||
+                !from_cookies.empty()) {
+              c.from_headers = std::move(from_headers);
+              c.from_params = std::move(from_params);
+              c.from_cookies = std::move(from_cookies);
+            }
+            std::unordered_map<std::string, ClaimToHeader> claims_to_headers;
+            if (!JsonArrayIterate(
+                    consumer, "claims_to_headers",
+                    [&](const json& item_json) -> bool {
+                      JSON_VALUE_AS(Wasm::Common::JsonObject, item_json, item,
+                                    "invalid item");
+                      JSON_FIND_FIELD(item, claim);
+                      JSON_FIELD_VALUE_AS(std::string, item, claim);
+                      auto c2h_it = claims_to_headers.find(item_claim);
+                      if (c2h_it != claims_to_headers.end()) {
+                        LOG_WARN("claim to header already exists: " +
+                                 item_claim);
+                        return false;
+                      }
+                      auto& c2h = claims_to_headers[item_claim];
+                      JSON_FIND_FIELD(item, header);
+                      JSON_FIELD_VALUE_AS(std::string, item, header);
+                      c2h.header = std::move(item_header);
+                      auto item_override_json = item.find("override");
+                      if (item_override_json != item.end()) {
+                        JSON_FIELD_VALUE_AS(bool, item, override);
+                        c2h.override = item_override;
+                      }
+                      return true;
+                    })) {
+              LOG_WARN("failed to parse 'claims_to_headers' in consumer: " +
+                       consumer_name);
+              return false;
+            }
+            c.claims_to_headers = std::move(claims_to_headers);
+            auto consumer_clock_skew_seconds_json =
+                consumer.find("clock_skew_seconds");
+            if (consumer_clock_skew_seconds_json != consumer.end()) {
+              JSON_FIELD_VALUE_AS(uint64_t, consumer, clock_skew_seconds);
+              c.clock_skew = consumer_clock_skew_seconds;
+            }
+            auto consumer_keep_token_json = consumer.find("keep_token");
+            if (consumer_keep_token_json != consumer.end()) {
+              JSON_FIELD_VALUE_AS(bool, consumer, keep_token);
+              c.keep_token = consumer_keep_token;
+            }
+            c.extractor = Extractor::create(c);
+            rule.consumers.push_back(std::move(c));
+            name_set.insert(consumer_name);
+            return true;
+          })) {
+    LOG_WARN("failed to parse configuration for consumers.");
+    return false;
+  }
+  if (rule.consumers.empty()) {
+    LOG_INFO("at least one consumer has to be configured for a rule.");
+    return false;
+  }
+  return true;
+}
+
+Status PluginRootContext::consumerVerify(
+    const Consumer& consumer, uint64_t now,
+    std::vector<JwtLocationConstPtr>& jwt_tokens) {
+  auto tokens = consumer.extractor->extract();
+  if (tokens.empty()) {
+    return Status::JwtMissed;
+  }
+  for (auto& token : tokens) {
+    google::jwt_verify::Jwt jwt;
+    Status status = jwt.parseFromString(token->token());
+    if (status != Status::Ok) {
+      LOG_INFO(absl::StrFormat(
+          "jwt parse failed, consumer:%s, token:%s, status:%s", consumer.name,
+          token->token(), google::jwt_verify::getStatusString(status)));
+      return status;
+    }
+    StructUtils payload_getter(jwt.payload_pb_);
+    if (!consumer.allowd_claims.empty()) {
+      for (const auto& claim : consumer.allowd_claims) {
+        std::string value;
+        if (payload_getter.GetString(claim.first, &value) ==
+            StructUtils::WRONG_TYPE) {
+          LOG_INFO(absl::StrFormat(
+              "jwt payload invalid, consumer:%s, token:%s, claim:%s",
+              consumer.name, jwt.payload_str_, claim.first));
+          return Status::JwtVerificationFail;
+        }
+        if (value != claim.second) {
+          LOG_INFO(absl::StrFormat(
+              "jwt payload invalid, consumer:%s, claim:%s, value:%s, expect:%s",
+              consumer.name, claim.first, value, claim.second));
+          return Status::JwtVerificationFail;
+        }
+      }
+    }
+    status = jwt.verifyTimeConstraint(now, consumer.clock_skew);
+    if (status != Status::Ok) {
+      LOG_DEBUG(absl::StrFormat(
+          "jwt verify time failed, consumer:%s,  token:%s, status:%s",
+          consumer.name, token->token(),
+          google::jwt_verify::getStatusString(status)));
+      return status;
+    }
+    status =
+        google::jwt_verify::verifyJwtWithoutTimeChecking(jwt, *consumer.jwks);
+    if (status != Status::Ok) {
+      LOG_DEBUG(absl::StrFormat(
+          "jwt verify failed, consumer:%s, token:%s, status:%s", consumer.name,
+          token->token(), google::jwt_verify::getStatusString(status)));
+      return status;
+    }
+    for (const auto& claim_to_header : consumer.claims_to_headers) {
+      std::string value;
+      if (payload_getter.GetString(claim_to_header.first, &value) !=
+          StructUtils::WRONG_TYPE) {
+        token->addClaimToHeader(claim_to_header.second.header, value,
+                                claim_to_header.second.override);
+      } else {
+        uint64_t num_value;
+        if (payload_getter.GetUInt64(claim_to_header.first, &num_value) !=
+            StructUtils::WRONG_TYPE) {
+          token->addClaimToHeader(claim_to_header.second.header,
+                                  std::to_string((unsigned long long)num_value),
+                                  claim_to_header.second.override);
+        }
+      }
+    }
+  }
+  jwt_tokens = std::move(tokens);
+  return Status::Ok;
+}
+
+bool PluginRootContext::checkPlugin(
+    const JwtAuthConfigRule& rule,
+    const std::optional<std::unordered_set<std::string>>& allow_set) {
+  std::optional<Status> err_status;
+  bool verified = false;
+  uint64_t now = getCurrentTimeNanoseconds() / 1e9;
+  for (const auto& consumer : rule.consumers) {
+    std::vector<JwtLocationConstPtr> tokens;
+    auto status = consumerVerify(consumer, now, tokens);
+    if (status == Status::Ok) {
+      verified = true;
+      // global config without allow_set field allows any consumers
+      if (!allow_set ||
+          allow_set.value().find(consumer.name) != allow_set.value().end()) {
+        addRequestHeader("X-Mse-Consumer", consumer.name);
+        for (auto& token : tokens) {
+          if (!consumer.keep_token) {
+            token->removeJwt();
+          }
+          token->claimsToHeaders();
+        }
+        return true;
+      }
+    }
+    // use the first status
+    if (!err_status) {
+      err_status = status;
+    }
+  }
+  if (!verified) {
+    auto status = err_status ? err_status.value() : Status::JwtMissed;
+    auto err_str = google::jwt_verify::getStatusString(status);
+    auto authn_value = absl::StrCat(
+        "Bearer realm=\"",
+        Wasm::Common::Http::buildOriginalUri(MaximumUriLength), "\"");
+    if (status != Status::JwtMissed) {
+      absl::StrAppend(&authn_value, InvalidTokenErrorString);
+    }
+    sendLocalResponse(401, generateRcDetails(err_str), err_str,
+                      {{"WWW-Authenticate", authn_value}});
+  } else {
+    sendLocalResponse(403, kRcDetailJwtAuthnPrefix, "Access Denied", {});
+  }
+  return false;
+}
+
+bool PluginRootContext::onConfigure(size_t size) {
+  // Parse configuration JSON string.
+  if (size > 0 && !configure(size)) {
+    LOG_WARN("configuration has errors initialization will not continue.");
+    setInvalidConfig();
+    return true;
+  }
+  return true;
+}
+
+bool PluginRootContext::configure(size_t configuration_size) {
+  auto configuration_data = getBufferBytes(WasmBufferType::PluginConfiguration,
+                                           0, configuration_size);
+  // Parse configuration JSON string.
+  auto result = ::Wasm::Common::JsonParse(configuration_data->view());
+  if (!result) {
+    LOG_WARN(absl::StrCat("cannot parse plugin configuration JSON string: ",
+                          configuration_data->view()));
+    return false;
+  }
+  if (!parseAuthRuleConfig(result.value())) {
+    LOG_WARN(absl::StrCat("cannot parse plugin configuration JSON string: ",
+                          configuration_data->view()));
+    return false;
+  }
+  return true;
+}
+
+FilterHeadersStatus PluginContext::onRequestHeaders(uint32_t, bool) {
+  auto* rootCtx = rootContext();
+  return rootCtx->checkAuthRule(
+             [rootCtx](const auto& config, const auto& allow_set) {
+               return rootCtx->checkPlugin(config, allow_set);
+             })
+             ? FilterHeadersStatus::Continue
+             : FilterHeadersStatus::StopIteration;
+}
+
+#ifdef NULL_PLUGIN
+
+}  // namespace jwt_auth
+}  // namespace null_plugin
+}  // namespace proxy_wasm
+
+#endif