tlbmc/hft/core/hft_service.cc

#include "tlbmc/hft/core/hft_service.h"

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <queue>
#include <string>
#include <utility>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/functional/any_invocable.h"
#include "absl/functional/function_ref.h"
#include "absl/log/log.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/substitute.h"
#include "absl/synchronization/mutex.h"
#include "peer_identity.h"
#include "peer_identity_grpc.h"
#include "redfish_authorizer_interface.h"
#include "hft_capabilities.pb.h"
#include "hft_service.pb.h"
#include "payload.pb.h"
#include "subscription_params.pb.h"
#include "tlbmc/hft/core/edge_filter.h"
#include "tlbmc/hft/core/manager.h"
#include "g3/grpc_headers.h"

namespace milotic_hft {

namespace {

// One-shot reactor that finishes the stream immediately with the given
// status. Used by the service to reject calls (cert / authz / empty
// request / sample-rate) without bringing up a full ServerReactorImpl.
class RejectReactor : public grpc::ServerWriteReactor<HftResponse> {
 public:
  explicit RejectReactor(const grpc::Status& status) { Finish(status); }
  void OnDone() override { delete this; }
};

}  // namespace

namespace internal {

ServerReactorImpl::ServerReactorImpl(
    SubscriptionManager* manager,
    absl::AnyInvocable<void(ServerReactorImpl*)> on_done,
    std::size_t maximum_event_queue_size)
    : manager_(manager),
      on_done_(std::move(on_done)),
      maximum_event_queue_size_(maximum_event_queue_size) {}

std::shared_ptr<ServerReactorImpl> ServerReactorImpl::Create(
    SubscriptionManager* manager,
    absl::AnyInvocable<void(ServerReactorImpl*)> on_done,
    std::size_t maximum_event_queue_size) {
  auto reactor = std::shared_ptr<ServerReactorImpl>(new ServerReactorImpl(
      manager, std::move(on_done), maximum_event_queue_size));
  absl::MutexLock lock(reactor->mutex_);
  reactor->self_ = reactor;
  return reactor;
}

void ServerReactorImpl::Begin(const HftRequest& request,
                              const grpc::Status& setup_status) {
  if (!setup_status.ok()) {
    SafeFinish(setup_status);
    return;
  }
  if (request.subscriptions_size() == 0) {
    SafeFinish(grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                            "No subscription is provided."));
    return;
  }

  std::shared_ptr<ServerReactorImpl> self;
  {
    absl::MutexLock lock(mutex_);
    self = self_;
  }

  for (const SubscriptionParams& subscription : request.subscriptions()) {
    absl::AnyInvocable<void(Payload&&)> on_data_callback =
        [self](Payload&& payload) {
          HftResponse response;
          *response.add_payloads() = std::move(payload);
          self->AddResponse(std::move(response));
        };
    absl::StatusOr<std::shared_ptr<SubscriptionManager::Subscription>>
        subscription_handle = manager_->AddSubscription(
            subscription, std::move(on_data_callback));
    if (!subscription_handle.ok()) {
      SafeFinish(grpc::Status(
          grpc::StatusCode::INTERNAL,
          absl::Substitute("Failed to add subscription: $0",
                           subscription_handle.status().message())));
      break;
    }
    absl::MutexLock lock(mutex_);
    subscriptions_.push_back(std::move(*subscription_handle));
  }
}

ServerReactorImpl::~ServerReactorImpl() {
  LOG(INFO) << "~Destroying reactor()";
}

void ServerReactorImpl::SafeFinish(const grpc::Status& status)
    ABSL_LOCKS_EXCLUDED(mutex_) {
  absl::MutexLock lock(mutex_);
  if (status_ != Status::kOpen) {
    return;
  }
  status_ = Status::kFinishCalled;
  Finish(status);
}

bool ServerReactorImpl::AddResponse(HftResponse&& response)
    ABSL_LOCKS_EXCLUDED(mutex_) {
  absl::MutexLock lock(mutex_);
  if (status_ != Status::kOpen) {
    DLOG(WARNING) << "Stream is finished or finishing";
    return false;
  }
  if (!current_write_.has_value()) {
    current_write_ = std::move(response);
    StartWrite(&*current_write_);
    return true;
  }
  if (pending_.size() >= maximum_event_queue_size_) {
    DLOG(WARNING) << "Response queue is full, dropping response";
    return false;
  }
  pending_.push(std::move(response));
  return true;
}

ServerReactorImpl::Status ServerReactorImpl::GetStatus() const
    ABSL_LOCKS_EXCLUDED(mutex_) {
  absl::MutexLock lock(mutex_);
  return status_;
}

std::size_t ServerReactorImpl::NumSubscriptions() const
    ABSL_LOCKS_EXCLUDED(mutex_) {
  absl::MutexLock lock(mutex_);
  return subscriptions_.size();
}

std::vector<std::shared_ptr<SubscriptionManager::Subscription>>
ServerReactorImpl::GetSubscriptionsForDebug() const
    ABSL_LOCKS_EXCLUDED(mutex_) {
  absl::MutexLock lock(mutex_);
  return subscriptions_;
}

void ServerReactorImpl::OnWriteDone(bool ok) ABSL_LOCKS_EXCLUDED(mutex_) {
  absl::MutexLock lock(mutex_);
  if (!ok) {
    LOG(ERROR) << "Failed to write response";
    if (status_ == Status::kOpen) {
      status_ = Status::kFinishCalled;
      Finish(
          grpc::Status(grpc::StatusCode::INTERNAL, "Failed to write response"));
    }
    return;
  }
  if (status_ != Status::kOpen || pending_.empty()) {
    current_write_.reset();
    return;
  }
  current_write_ = std::move(pending_.front());
  pending_.pop();
  StartWrite(&*current_write_);
}

void ServerReactorImpl::OnCancel() {
  DLOG(INFO) << "OnCancel";
  absl::MutexLock lock(mutex_);
  if (status_ != Status::kOpen) {
    return;
  }
  status_ = Status::kFinishCalled;
  Finish(
      grpc::Status(grpc::StatusCode::CANCELLED, "Client cancelled requests."));
}

void ServerReactorImpl::OnDone() {
  DLOG(INFO) << "OnDone";
  std::shared_ptr<ServerReactorImpl> self;
  std::vector<std::shared_ptr<SubscriptionManager::Subscription>> subscriptions;
  {
    absl::MutexLock lock(mutex_);
    if (status_ == Status::kFinished) {
      return;
    }
    status_ = Status::kFinished;
    self = std::move(self_);
    // Move the handles out under the mutex; this also breaks the
    // reactor <-> subscription cycle (subscriptions hold the on_data_callback
    // closure, which captures shared_ptr<this>).
    subscriptions = std::move(subscriptions_);
  }
  for (const std::shared_ptr<SubscriptionManager::Subscription>& sub :
       subscriptions) {
    if (absl::Status status = manager_->Unsubscribe(sub); !status.ok()) {
      LOG(WARNING) << "Failed to unsubscribe subscription: " << status;
    }
  }
  on_done_(this);
  // `self` goes out of scope here. If no other holder kept a reference to
  // the reactor, this destroys it.
}

}  // namespace internal

std::vector<std::pair<std::string, uint64_t>>
HftServiceImpl::GetRoleSampleRatesForDebug() const {
  absl::MutexLock lock(mutex_);
  return {role_to_total_sample_rate_.begin(), role_to_total_sample_rate_.end()};
}

void HftServiceImpl::ForEachReactorForDebug(
    absl::FunctionRef<void(const internal::ServerReactorImpl&)> fn) const {
  absl::MutexLock lock(mutex_);
  for (const internal::ServerReactorImpl* reactor_ptr : reactors_) {
    fn(*reactor_ptr);
  }
}

uint64_t HftServiceImpl::GetAccumulativeSampleRate(
    const std::string& role) const {
  absl::MutexLock lock(mutex_);
  auto it = role_to_total_sample_rate_.find(role);
  if (it != role_to_total_sample_rate_.end()) {
    return it->second;
  }
  return 0;
}

uint64_t HftServiceImpl::GetReactorsCount() const {
  absl::MutexLock lock(mutex_);
  return reactors_.size();
}

uint64_t HftServiceImpl::GetRoleToTotalSampleRateCount() const {
  absl::MutexLock lock(mutex_);
  return role_to_total_sample_rate_.size();
}

uint64_t HftServiceImpl::GetReactorToSubscriptionIdsCount() const {
  absl::MutexLock lock(mutex_);
  uint64_t count = 0;
  for (const auto* reactor_ptr : reactors_) {
    if (reactor_ptr->NumSubscriptions() > 0) {
      ++count;
    }
  }
  return count;
}

void HftServiceImpl::WaitUntilNoReactors() {
  absl::MutexLock lock(mutex_);
  mutex_.Await(absl::Condition(this, &HftServiceImpl::NoReactorsLeft));
}

bool HftServiceImpl::NoReactorsLeft() const { return reactors_.empty(); }

HftServiceImpl::HftServiceImpl(
    HftServiceOptions options,
    std::unique_ptr<SubscriptionManager> subscription_manager,
    const milotic::authz::RedfishAuthorizerInterface* authorizer)
    : options_(std::move(options)),
      subscription_manager_(std::move(subscription_manager)),
      authorizer_(authorizer) {}

ServiceFeatures HftServiceImpl::GetSupportedServiceFeatures() {
  ServiceFeatures features;
  features.set_edge_filtering_supported(TRISTATE_TRUE);
  *features.mutable_edge_filtering_detail() = GetSupportedEdgeFilteringDetail();
  return features;
}

HftServiceImpl::AuthorizeResult HftServiceImpl::AuthorizeRequest(
    grpc::CallbackServerContext* context) const {
  if (!options_.enable_authorization) {
    return {.status = grpc::Status::OK, .role = ""};
  }
  // Only allow requests when the server is ready to serve (for tlbmc: it has
  // root certs and a prod-signed cert). A missing authorizer or serving hook
  // fails closed.
  if (authorizer_ == nullptr || options_.is_serving_allowed == nullptr ||
      !options_.is_serving_allowed()) {
    return {.status = grpc::Status(
                grpc::StatusCode::PERMISSION_DENIED,
                "The server does not have root certs and prod signed cert."),
            .role = ""};
  }
  milotic::authz::PeerSpiffeIdentity peer_identity;
  grpc::Status authz_status =
      milotic::authz::ExtractPeerIdentityFromAuthContext(
          *context->auth_context(), peer_identity);
  if (!authz_status.ok()) {
    return {.status = authz_status, .role = ""};
  }
  std::string role = authorizer_->GetPeerRedfishRole(peer_identity);
  if (role.empty()) {
    return {.status = grpc::Status(
                grpc::StatusCode::PERMISSION_DENIED,
                absl::StrCat("Peer role is not specified in the auth config: "
                             "peer_identity=",
                             peer_identity.spiffe_id)),
            .role = ""};
  }
  return {.status = grpc::Status::OK, .role = std::move(role)};
}

grpc::ServerUnaryReactor* HftServiceImpl::GetCapabilities(
    grpc::CallbackServerContext* context,
    const GetCapabilitiesRequest* /*request*/, CapabilitiesResponse* response) {
  grpc::ServerUnaryReactor* reactor = context->DefaultReactor();
  if (AuthorizeResult authz = AuthorizeRequest(context); !authz.status.ok()) {
    reactor->Finish(authz.status);
    return reactor;
  }
  *response->mutable_features() = GetSupportedServiceFeatures();
  reactor->Finish(grpc::Status::OK);
  return reactor;
}

grpc::ServerWriteReactor<HftResponse>* HftServiceImpl::Subscribe(
    grpc::CallbackServerContext* context, const HftRequest* request) {
  AuthorizeResult authz = AuthorizeRequest(context);
  return SubscribeWithStatus(request, authz.role, authz.status);
}

grpc::ServerWriteReactor<HftResponse>* HftServiceImpl::SubscribeWithStatus(
    const HftRequest* request, const std::string& role,
    grpc::Status authz_status) {
  using internal::ServerReactorImpl;
  // Only allow subscriptions when the server is ready to serve (for tlbmc: it
  // has root certs and a prod-signed cert). A missing authorizer or serving
  // hook fails closed.
  if (options_.enable_authorization &&
      (authorizer_ == nullptr || options_.is_serving_allowed == nullptr ||
       !options_.is_serving_allowed())) {
    return new RejectReactor(grpc::Status(
        grpc::StatusCode::PERMISSION_DENIED,
        "The server does not have root certs and prod signed cert."));
  }
  if (options_.enable_authorization && !authz_status.ok()) {
    return new RejectReactor(authz_status);
  }
  if (request->subscriptions_size() == 0) {
    return new RejectReactor(grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                                          "No subscription is provided."));
  }

  uint64_t total_sample_rate = 0;
  // TODO(nanzhou): limit subscriptions for "all" resources as well. For now,
  // assume that subscriptions all come with explicit identifiers.
  for (const SubscriptionParams& subscription : request->subscriptions()) {
    // Bill the requested rate. Sensor sampling intervals are validated to a
    // power of two in AddSubscription, so this is the rate the subscription
    // will actually run at; FRU subscriptions are billed at their requested
    // interval. (An invalid sensor interval is rejected later by
    // AddSubscription; guard the division so a malformed request cannot divide
    // by zero here.)
    const int32_t effective_interval_ms =
        std::max(1, subscription.sampling_interval_ms());
    total_sample_rate += 1000 / effective_interval_ms;
  }

  // Sample-rate gate. Reads and writes role_to_total_sample_rate_, so the
  // check and the bookkeeping update are taken under the service mutex
  // atomically. On failure we exit before touching anything.
  bool should_decrease_sample_rate = false;
  if (options_.enable_authorization) {
    absl::MutexLock lock(mutex_);
    uint64_t new_total_sample_rate = total_sample_rate;
    auto it = role_to_total_sample_rate_.find(role);
    if (it != role_to_total_sample_rate_.end()) {
      new_total_sample_rate += it->second;
    }
    uint64_t sample_rate_limit = authorizer_->GetSampleRateLimit(role);
    if (new_total_sample_rate > sample_rate_limit) {
      return new RejectReactor(grpc::Status(
          grpc::StatusCode::RESOURCE_EXHAUSTED,
          absl::Substitute(
              "Role $0 has reached the maximum allowed sample "
              "rate of $1, currently $2, new total sample rate "
              "after this RPC: $3",
              role, sample_rate_limit,
              it == role_to_total_sample_rate_.end() ? 0 : it->second,
              new_total_sample_rate)));
    }
    role_to_total_sample_rate_[role] = new_total_sample_rate;
    should_decrease_sample_rate = true;
  }

  absl::AnyInvocable<void(ServerReactorImpl*)> on_reactor_done =
      [this, total_sample_rate, role,
       should_decrease =
           should_decrease_sample_rate](ServerReactorImpl* reactor) {
        absl::MutexLock lock(mutex_);
        if (auto it = role_to_total_sample_rate_.find(role);
            should_decrease && it != role_to_total_sample_rate_.end()) {
          it->second = it->second - total_sample_rate;
          DLOG(INFO) << "role: " << role
                     << " total_sample_rate: " << total_sample_rate
                     << " role_to_total_sample_rate_: "
                     << role_to_total_sample_rate_[role];
          if (it->second == 0) {
            role_to_total_sample_rate_.erase(it);
          }
        }
        reactors_.erase(reactor);
      };

  std::shared_ptr<ServerReactorImpl> reactor = ServerReactorImpl::Create(
      subscription_manager_.get(), std::move(on_reactor_done),
      options_.maximum_event_queue_size);
  {
    absl::MutexLock lock(mutex_);
    reactors_.insert(reactor.get());
  }

  reactor->Begin(*request, grpc::Status::OK);
  return reactor.get();
}

}  // namespace milotic_hft