Add an authenticated direct-query PEM mode for vizier deployments

src/carnot/blocking_agg_benchmark.cc

@@ -89,9 +89,9 @@ void BM_Query(benchmark::State& state, std::vector<types::DataType> types,
               int64_t num_batches, const datagen::DistributionParams* dist_vars,
               const datagen::DistributionParams* len_vars) {
   auto table_store = std::make_shared<table_store::TableStore>();
-  auto server = LocalGRPCResultSinkServer();
+  auto server = LocalGRPCResultSinkServer::Create().ConsumeValueOrDie();
 
-  auto carnot = SetUpCarnot(table_store, &server);
+  auto carnot = SetUpCarnot(table_store, server.get());
   auto table = table_store::CreateTable(types, distribution_types, state.range(0), num_batches,
                                         dist_vars, len_vars)
                    .ConsumeValueOrDie();
@@ -105,8 +105,8 @@ void BM_Query(benchmark::State& state, std::vector<types::DataType> types,
     if (!res.ok()) {
       LOG(FATAL) << "Aggregate benchmark query did not execute successfully.";
     }
-    bytes_processed += server.exec_stats().ConsumeValueOrDie().execution_stats().bytes_processed();
-    server.ResetQueryResults();
+    bytes_processed += server->exec_stats().ConsumeValueOrDie().execution_stats().bytes_processed();
+    server->ResetQueryResults();
     ++i;
   }
 

src/carnot/carnot_executable.cc

@@ -240,13 +240,18 @@ int main(int argc, char* argv[]) {
 
   // Execute query.
   auto table_store = std::make_shared<px::table_store::TableStore>();
-  auto result_server = px::carnot::exec::LocalGRPCResultSinkServer();
+  auto result_server_or = px::carnot::exec::LocalGRPCResultSinkServer::Create();
+  if (!result_server_or.ok()) {
+    LOG(FATAL) << absl::Substitute("Failed to start result server: $0",
+                                   result_server_or.status().msg());
+  }
+  auto result_server = result_server_or.ConsumeValueOrDie();
   auto func_registry = std::make_unique<px::carnot::udf::Registry>("default_registry");
   px::carnot::funcs::RegisterFuncsOrDie(func_registry.get());
   auto clients_config =
       std::make_unique<px::carnot::Carnot::ClientsConfig>(px::carnot::Carnot::ClientsConfig{
           [&result_server](const std::string& address, const std::string&) {
-            return result_server.StubGenerator(address);
+            return result_server->StubGenerator(address);
           },
           [](grpc::ClientContext*) {},
       });
@@ -263,13 +268,13 @@ int main(int argc, char* argv[]) {
     LOG(FATAL) << absl::Substitute("Query failed to execute: $0", exec_status.msg());
   }
 
-  auto output_names = result_server.output_tables();
+  auto output_names = result_server->output_tables();
   if (!output_names.size()) {
     LOG(FATAL) << "Query produced no output tables.";
   }
-  std::string output_name = *(result_server.output_tables().begin());
+  std::string output_name = *(result_server->output_tables().begin());
   LOG(INFO) << absl::Substitute("Writing results for output table: $0", output_name);
   // Write output table to CSV.
-  TableToCsv(output_filename, result_server.query_results(output_name));
+  TableToCsv(output_filename, result_server->query_results(output_name));
   return 0;
 }

src/carnot/carnot_test.cc

@@ -50,7 +50,7 @@ class CarnotTest : public ::testing::Test {
   void SetUp() override {
     Test::SetUp();
     table_store_ = std::make_shared<table_store::TableStore>();
-    result_server_ = std::make_unique<exec::LocalGRPCResultSinkServer>();
+    ASSERT_OK_AND_ASSIGN(result_server_, exec::LocalGRPCResultSinkServer::Create());
 
     auto func_registry = std::make_unique<px::carnot::udf::Registry>("default_registry");
     funcs::RegisterFuncsOrDie(func_registry.get());

src/carnot/end_to_end_join_test.cc

@@ -44,7 +44,7 @@ class JoinTest : public ::testing::Test {
   void SetUp() override {
     Test::SetUp();
     table_store_ = std::make_shared<table_store::TableStore>();
-    result_server_ = std::make_unique<exec::LocalGRPCResultSinkServer>();
+    ASSERT_OK_AND_ASSIGN(result_server_, exec::LocalGRPCResultSinkServer::Create());
     auto func_registry = std::make_unique<px::carnot::udf::Registry>("default_registry");
     px::carnot::funcs::RegisterFuncsOrDie(func_registry.get());
     auto clients_config = std::make_unique<Carnot::ClientsConfig>(Carnot::ClientsConfig{

src/carnot/engine_state.h

@@ -128,14 +128,18 @@ class EngineState : public NotCopyable {
     return std::make_unique<plan::PlanState>(func_registry_.get());
   }
 
+  // max_output_rows_per_table of 0 means unlimited. That is the right default
+  // for the single-node executable and for tests, where the caller controls the
+  // input, but a server accepting queries from a peer should pass a real cap:
+  // results are buffered in memory before they are streamed out.
   std::unique_ptr<planner::CompilerState> CreateLocalExecutionCompilerState(
-      types::Time64NSValue time_now) {
+      types::Time64NSValue time_now, int64_t max_output_rows_per_table = 0) {
     auto rel_map = table_store_->GetRelationMap();
     // Use an empty string for query result address, because the local execution mode should use
     // the Local GRPC result server to send results to.
     return std::make_unique<planner::CompilerState>(
         std::move(rel_map), planner::SensitiveColumnMap{}, registry_info_.get(), time_now,
-        /* max_output_rows_per_table */ 0,
+        max_output_rows_per_table,
         /* result address */ "",
         /* ssl target name override*/ "", planner::RedactionOptions{}, nullptr, nullptr,
         planner::DebugInfo{});

src/carnot/exec/local_grpc_result_server.h

@@ -64,7 +64,12 @@ class LocalResultSinkServer final : public carnotpb::ResultSinkService::Service
     return ::grpc::Status::OK;
   }
 
-  void ResetQueryResults() { query_results_.clear(); }
+  void ResetQueryResults() {
+    // Takes the same lock TransferResultChunk does: a caller resetting between
+    // queries races an in-flight chunk write otherwise.
+    const std::lock_guard<std::mutex> lock(result_mutex_);
+    query_results_.clear();
+  }
 
  private:
   // List of the query results received.
@@ -75,13 +80,14 @@ class LocalResultSinkServer final : public carnotpb::ResultSinkService::Service
 
 class LocalGRPCResultSinkServer {
  public:
-  LocalGRPCResultSinkServer() {
-    grpc::ServerBuilder builder;
-
-    builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials());
-    builder.RegisterService(&result_sink_server_);
-    grpc_server_ = builder.BuildAndStart();
-    CHECK(grpc_server_ != nullptr);
+  // Starts the in-process result sink server. Returns an error if it fails to bind.
+  static StatusOr<std::unique_ptr<LocalGRPCResultSinkServer>> Create() {
+    auto server = std::unique_ptr<LocalGRPCResultSinkServer>(new LocalGRPCResultSinkServer());
+    if (server->grpc_server_ == nullptr) {
+      return error::ResourceUnavailable(
+          "failed to start the local result sink server on 127.0.0.1:0");
+    }
+    return server;
   }
 
   ~LocalGRPCResultSinkServer() {
@@ -151,6 +157,14 @@ class LocalGRPCResultSinkServer {
   }
 
  private:
+  LocalGRPCResultSinkServer() {
+    grpc::ServerBuilder builder;
+
+    builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials());
+    builder.RegisterService(&result_sink_server_);
+    grpc_server_ = builder.BuildAndStart();
+  }
+
   std::unique_ptr<grpc::Server> grpc_server_;
   LocalResultSinkServer result_sink_server_;
 };

src/vizier/services/agent/pem/BUILD.bazel

@@ -31,12 +31,22 @@ pl_cc_library(
     ),
     hdrs = glob(["*.h"]),
     deps = [
+        "//src/api/proto/vizierpb:vizier_pl_cc_proto",
+        "//src/carnot",
+        "//src/carnot/carnotpb:carnot_pl_cc_proto",
+        "//src/carnot/funcs:cc_library",
+        "//src/carnot/planner/compiler:cc_library",
         "//src/carnot/planner/dynamic_tracing/ir/logicalpb:logical_pl_cc_proto",
+        "//src/carnot/planpb:plan_pl_cc_proto",
         "//src/integrations/grpc_clocksync:cc_library",
+        "//src/shared/services/jwt:cc_library",
         "//src/shared/tracepoint_translation:cc_library",
         "//src/stirling:cc_library",
         "//src/stirling/source_connectors/dynamic_tracer/dynamic_tracing/ir/logicalpb:logical_pl_cc_proto",
         "//src/vizier/services/agent/shared/manager:cc_library",
+        "//src/vizier/services/agent/shared/vizier_results:cc_library",
+        "@com_github_grpc_grpc//:grpc++",
+        "@com_github_rlyeh_sole//:sole",
     ],
 )
 
@@ -51,6 +61,22 @@ pl_cc_test(
     ],
 )
 
+pl_cc_test(
+    name = "direct_query_server_test",
+    srcs = ["direct_query_server_test.cc"],
+    deps = [
+        ":cc_library",
+        "//src/api/proto/vizierpb:vizier_pl_cc_proto",
+        "//src/carnot",
+        "//src/carnot/funcs:cc_library",
+        "//src/shared/services/jwt:cc_library",
+        "//src/shared/types:cc_library",
+        "//src/table_store:cc_library",
+        "@com_github_grpc_grpc//:grpc++",
+        "@com_github_rlyeh_sole//:sole",
+    ],
+)
+
 pl_cc_binary(
     name = "pem",
     srcs = ["pem_main.cc"],

src/vizier/services/agent/pem/direct_query_server.cc

@@ -0,0 +1,165 @@
+/*
+ * Copyright 2018- The Pixie 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "src/vizier/services/agent/pem/direct_query_server.h"
+
+#include <chrono>
+#include <cstdint>
+#include <cstring>
+#include <iterator>
+#include <string>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+#include <absl/strings/str_split.h>
+#include <absl/strings/string_view.h>
+#include <absl/strings/substitute.h>
+#include <jwt/base64.hpp>
+#include <sole.hpp>
+
+#include "src/carnot/carnot.h"
+#include "src/carnot/carnotpb/carnot.pb.h"
+#include "src/carnot/engine_state.h"
+#include "src/carnot/exec/local_grpc_result_server.h"
+#include "src/carnot/planner/compiler/compiler.h"
+#include "src/carnot/planpb/plan.pb.h"
+#include "src/common/base/base.h"
+#include "src/shared/services/jwt/service_token.h"
+#include "src/shared/types/typespb/wrapper/types_pb_wrapper.h"
+#include "src/vizier/services/agent/shared/vizier_results/result_conversion.h"
+
+namespace px {
+namespace vizier {
+namespace agent {
+
+namespace {
+
+constexpr char kAuthorizationMetadataKey[] = "authorization";
+
+}  // namespace
+
+::grpc::Status AuthenticateRequest(::grpc::ServerContext* ctx, const std::string& jwt_signing_key) {
+  if (jwt_signing_key.empty()) {
+    return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED,
+                          "direct-query: signing key not configured");
+  }
+  const auto& md = ctx->client_metadata();
+  const auto range = md.equal_range(kAuthorizationMetadataKey);
+  if (range.first == range.second) {
+    return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED,
+                          "direct-query: missing authorization metadata");
+  }
+  if (std::next(range.first) != range.second) {
+    return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED,
+                          "direct-query: duplicate authorization metadata");
+  }
+  const std::string_view raw(range.first->second.data(), range.first->second.size());
+  const std::string_view token = ::px::services::StripBearerPrefix(raw);
+  if (token.empty()) {
+    return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED,
+                          "direct-query: authorization is not a Bearer token");
+  }
+  const auto s = ::px::services::VerifyServiceJWT(token, jwt_signing_key);
+  if (!s.ok()) {
+    VLOG(1) << "direct-query: rejecting bearer token: " << s.msg();
+    return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED,
+                          "direct-query: invalid bearer token");
+  }
+  return ::grpc::Status::OK;
+}
+
+namespace {
+
+constexpr int64_t kMaxOutputRowsPerTable = 10000;
+
+std::vector<::px::api::vizierpb::ExecuteScriptResponse> collectResponses(
+    ::px::carnot::exec::LocalGRPCResultSinkServer* result_server, const std::string& query_id) {
+  std::vector<::px::api::vizierpb::ExecuteScriptResponse> out;
+  for (const auto& chunk : result_server->raw_query_results()) {
+    auto resp = ExecuteScriptResponseFromChunk(chunk, query_id);
+    if (!resp.has_value()) {
+      continue;
+    }
+    out.push_back(std::move(*resp));
+  }
+  return out;
+}
+
+}  // namespace
+
+::grpc::Status DirectQueryServer::ExecuteScript(
+    ::grpc::ServerContext* context, const ::px::api::vizierpb::ExecuteScriptRequest* request,
+    ::grpc::ServerWriter<::px::api::vizierpb::ExecuteScriptResponse>* writer) {
+  if (auto s = AuthenticateRequest(context, jwt_signing_key_); !s.ok()) {
+    return s;
+  }
+  if (request->mutation()) {
+    return ::grpc::Status(::grpc::StatusCode::UNIMPLEMENTED,
+                          "direct-query: mutations are out of scope (read-only endpoint)");
+  }
+  if (carnot_ == nullptr || engine_state_ == nullptr || result_server_ == nullptr) {
+    return ::grpc::Status(::grpc::StatusCode::FAILED_PRECONDITION,
+                          "direct-query: server not wired with a live Carnot");
+  }
+  const auto query_id = sole::uuid4();
+  const std::string query_id_str = query_id.str();
+
+  const auto time_now = ::px::CurrentTimeNS();
+
+  auto compiler_state =
+      engine_state_->CreateLocalExecutionCompilerState(time_now, kMaxOutputRowsPerTable);
+  auto plan_or = ::px::carnot::planner::compiler::Compiler().Compile(request->query_str(),
+                                                                     compiler_state.get());
+  if (!plan_or.ok()) {
+    auto msg = absl::Substitute("direct-query: PxL compile failed ($0)", plan_or.msg());
+    VLOG(1) << msg;
+    return ::grpc::Status(::grpc::StatusCode::INVALID_ARGUMENT, msg);
+  }
+  const auto plan = plan_or.ConsumeValueOrDie();
+  EmitSchemaResponses(plan, query_id_str, writer);
+
+  std::vector<::px::api::vizierpb::ExecuteScriptResponse> responses;
+  {
+    absl::MutexLock lk(&exec_mu_);
+    result_server_->ResetQueryResults();
+    auto exec_s = carnot_->ExecuteQuery(request->query_str(), query_id, time_now);
+    if (!exec_s.ok()) {
+      auto msg = absl::Substitute("direct-query: PxL execute failed ($0)", exec_s.msg());
+      VLOG(1) << msg;
+      return ::grpc::Status(::grpc::StatusCode::INTERNAL, msg);
+    }
+    responses = collectResponses(result_server_, query_id_str);
+  }
+
+  for (const auto& resp : responses) {
+    if (context->IsCancelled()) {
+      return ::grpc::Status(::grpc::StatusCode::CANCELLED, "direct-query: client cancelled");
+    }
+    if (!writer->Write(resp)) {
+      VLOG(1) << "direct-query: client stream closed mid-response, abandoning query "
+              << query_id_str;
+      break;
+    }
+  }
+  return ::grpc::Status::OK;
+}
+
+}  // namespace agent
+}  // namespace vizier
+}  // namespace px