// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef REMOTING_PROTOCOL_SESSION_AUTHZ_AUTHENTICATOR_H_
#define REMOTING_PROTOCOL_SESSION_AUTHZ_AUTHENTICATOR_H_

#include <memory>
#include <optional>
#include <string>
#include <string_view>

#include "base/functional/callback.h"
#include "base/memory/weak_ptr.h"
#include "base/time/time.h"
#include "remoting/base/constants.h"
#include "remoting/base/http_status.h"
#include "remoting/base/session_authz_service_client.h"
#include "remoting/base/session_policies.h"
#include "remoting/proto/session_authz_service.h"
#include "remoting/protocol/authenticator.h"
#include "remoting/protocol/credentials_type.h"
#include "remoting/protocol/session_authz_reauthorizer.h"

namespace remoting::protocol {

// Class that implements the host leg of the SessionAuthz mechanism. Note that
// this class only handles the initial authorization process. Inter-session
// reauthorization will be handled by a different class.
// See go/crd-sessionauthz-integration for internal details.
class SessionAuthzAuthenticator : public Authenticator {
 public:
  SessionAuthzAuthenticator(
      CredentialsType credentials_type,
      std::unique_ptr<SessionAuthzServiceClient> service_client,
      const CreateBaseAuthenticatorCallback&
          create_base_authenticator_callback);
  ~SessionAuthzAuthenticator() override;

  SessionAuthzAuthenticator(const SessionAuthzAuthenticator&) = delete;
  SessionAuthzAuthenticator& operator=(const SessionAuthzAuthenticator&) =
      delete;

  // Start the authenticator. The authenticator won't be in the `MESSAGE_READY`
  // state until |resume_callback| is called.
  void Start(base::OnceClosure resume_callback);

  const SessionAuthzReauthorizer* reauthorizer() const {
    return reauthorizer_.get();
  }

  // The session ID generated by the SessionAuthz server. This is only set if
  // the GenerateHostToken API call has completed and succeeded. Otherwise, this
  // will be an empty string.
  const std::string& session_id() const { return session_id_; }

  // The host token returned by the SessionAuthz server. This is only set if
  // the GenerateHostToken API call has completed and succeeded. Otherwise, this
  // will be an empty string.
  const std::string& host_token() const { return host_token_; }

  // Authenticator implementation.
  CredentialsType credentials_type() const override;
  const Authenticator& implementing_authenticator() const override;
  State state() const override;
  bool started() const override;
  RejectionReason rejection_reason() const override;
  RejectionDetails rejection_details() const override;
  void ProcessMessage(const JingleAuthentication& message,
                      base::OnceClosure resume_callback) override;
  JingleAuthentication GetNextMessage() override;
  const std::string& GetAuthKey() const override;
  const SessionPolicies* GetSessionPolicies() const override;

  void SetReauthorizerForTesting(
      std::unique_ptr<SessionAuthzReauthorizer> reauthorizer);
  void SetSessionIdForTesting(std::string_view session_id);
  void SetHostTokenForTesting(std::string_view host_token);

 private:
  enum class SessionAuthzState {
    NOT_STARTED,

    // A request was made to call GenerateHostToken on SessionAuthz, and the
    // response is pending.
    GENERATING_HOST_TOKEN,

    // The response for GenerateHostToken has been received, and is ready to be
    // sent to the client.
    READY_TO_SEND_HOST_TOKEN,

    // Waiting for the client to send the session token.
    WAITING_FOR_SESSION_TOKEN,

    // A request was made to call VerifySessionToken on SessionAuthz, with the
    // session token sent by the client, and the response is pending.
    VERIFYING_SESSION_TOKEN,

    // The shared secret has been returned by the VerifySessionToken call. The
    // rest of the authentication process is now handed over to the underlying
    // authenticator.
    SHARED_SECRET_FETCHED,

    // An error has occurred when trying to fetch the shared secret. Note that
    // an error might occur in the underlying authenticator, in which case the
    // SessionAuthz state will be `SHARED_SECRET_FETCHED` and
    // `underlying_->state()` will be `REJECTED`.
    FAILED,
  };

  void GenerateHostToken(base::OnceClosure resume_callback);
  void OnHostTokenGenerated(
      base::OnceClosure resume_callback,
      const HttpStatus& status,
      std::unique_ptr<internal::GenerateHostTokenResponseStruct> response);
  void VerifySessionToken(const JingleAuthentication& message,
                          base::OnceClosure resume_callback);
  void OnVerifiedSessionToken(
      const JingleAuthentication& message,
      base::OnceClosure resume_callback,
      const HttpStatus& status,
      std::unique_ptr<internal::VerifySessionTokenResponseStruct> response);
  void HandleSessionAuthzError(const std::string_view& action_name,
                               const HttpStatus& status);
  void StartReauthorizerIfNecessary();
  void OnReauthorizationFailed(HttpStatus::Code error_code,
                               const Authenticator::RejectionDetails& details);

  CredentialsType credentials_type_;
  std::unique_ptr<SessionAuthzServiceClient> service_client_;
  CreateBaseAuthenticatorCallback create_base_authenticator_callback_;
  std::unique_ptr<Authenticator> underlying_;
  std::unique_ptr<internal::VerifySessionTokenResponseStruct>
      verify_token_response_;
  std::optional<SessionPolicies> session_policies_;
  std::unique_ptr<SessionAuthzReauthorizer> reauthorizer_;

  SessionAuthzState session_authz_state_ = SessionAuthzState::NOT_STARTED;

  // The rejection reason specific to fetching the shared secret from
  // SessionAuthz. If |session_authz_state_| is NOT `ERROR`, the actual
  // rejection reason is delegated to `underlying_->rejection_reason()`.
  RejectionReason session_authz_rejection_reason_;

  RejectionDetails rejection_details_;

  std::string session_id_;
  std::string host_token_;

  base::WeakPtrFactory<SessionAuthzAuthenticator> weak_factory_{this};
};

}  // namespace remoting::protocol

#endif  // REMOTING_PROTOCOL_SESSION_AUTHZ_AUTHENTICATOR_H_