Implement AuthSessionStorage

Bug: b:271249180
Change-Id: I37503ab8031271db23f58daa19622bd608fc4454
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4323390
Commit-Queue: Martin Bidlingmaier <mbid@google.com>
Auto-Submit: Denis Kuznetsov <antrim@chromium.org>
Reviewed-by: Martin Bidlingmaier <mbid@google.com>
Cr-Commit-Position: refs/heads/main@{#1118291}

chromeos/ash/components/login/auth/public/BUILD.gn

@@ -11,6 +11,7 @@ component("authpublic") {
   public_deps = [
     ":challenge_response_key",
     "//components/password_manager/core/browser:password_hash_data",
+    "//google_apis",  # due to AuthFailure.
   ]
   deps = [
     "//ash/constants",
@@ -21,7 +22,6 @@ component("authpublic") {
     "//components/prefs",
     "//components/user_manager",
     "//crypto",
-    "//google_apis",
     "//net",
     "//third_party/abseil-cpp:absl",
   ]

chromeos/ash/components/login/auth/public/auth_failure.cc

@@ -3,6 +3,7 @@
 // found in the LICENSE file.
 
 #include "chromeos/ash/components/login/auth/public/auth_failure.h"
+
 #include "base/check_op.h"
 
 namespace ash {

chromeos/ash/components/osauth/impl/BUILD.gn

@@ -12,9 +12,16 @@ component("impl") {
   defines = [ "IS_CHROMEOS_ASH_COMPONENTS_OSAUTH_IMPL" ]
   public_deps = [ "//chromeos/ash/components/osauth/public" ]
 
-  deps = [ "//base" ]
+  deps = [
+    "//base",
+    "//chromeos/ash/components/dbus/userdataauth",
+    "//chromeos/ash/components/login/auth",
+    "//chromeos/ash/components/login/auth/public:authpublic",
+  ]
   sources = [
     "auth_parts_impl.cc",
     "auth_parts_impl.h",
+    "auth_session_storage_impl.cc",
+    "auth_session_storage_impl.h",
   ]
 }

chromeos/ash/components/osauth/impl/DEPS

@@ -0,0 +1,5 @@
+include_rules = [
+  "+chromeos/ash/components/cryptohome",
+  "+chromeos/ash/components/dbus",
+  "+chromeos/ash/components/login",
+]

chromeos/ash/components/osauth/impl/auth_parts_impl.cc

@@ -2,11 +2,13 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include "chromeos/ash/components/osauth/impl/auth_parts_impl.h"
+
 #include <memory>
 
 #include "base/check.h"
 #include "base/check_op.h"
-#include "chromeos/ash/components/osauth/impl/auth_parts_impl.h"
+#include "chromeos/ash/components/osauth/impl/auth_session_storage_impl.h"
 #include "chromeos/ash/components/osauth/public/auth_parts.h"
 
 namespace ash {
@@ -26,7 +28,7 @@ std::unique_ptr<AuthPartsImpl> AuthPartsImpl::CreateTestInstance() {
 // static
 std::unique_ptr<AuthParts> AuthParts::Create() {
   std::unique_ptr<AuthPartsImpl> result = std::make_unique<AuthPartsImpl>();
-  // Create all parts here.
+  result->CreateDefaultComponents();
   return result;
 }
 
@@ -46,4 +48,13 @@ AuthPartsImpl::~AuthPartsImpl() {
   g_instance = nullptr;
 }
 
+void AuthPartsImpl::CreateDefaultComponents() {
+  session_storage_ = std::make_unique<AuthSessionStorageImpl>();
+}
+
+AuthSessionStorage* AuthPartsImpl::GetAuthSessionStorage() {
+  CHECK(session_storage_);
+  return session_storage_.get();
+}
+
 }  // namespace ash

chromeos/ash/components/osauth/impl/auth_parts_impl.h

@@ -20,6 +20,15 @@ class COMPONENT_EXPORT(CHROMEOS_ASH_COMPONENTS_OSAUTH) AuthPartsImpl
 
   AuthPartsImpl();
   ~AuthPartsImpl() override;
+
+  // AuthParts implementation:
+  AuthSessionStorage* GetAuthSessionStorage() override;
+
+ private:
+  friend class AuthParts;
+  void CreateDefaultComponents();
+
+  std::unique_ptr<AuthSessionStorage> session_storage_;
 };
 
 }  // namespace ash

chromeos/ash/components/osauth/impl/auth_session_storage_impl.cc

@@ -0,0 +1,125 @@
+// Copyright 2023 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chromeos/ash/components/osauth/impl/auth_session_storage_impl.h"
+
+#include <memory>
+#include <string>
+
+#include "base/check.h"
+#include "base/functional/callback.h"
+#include "base/location.h"
+#include "base/unguessable_token.h"
+#include "chromeos/ash/components/dbus/userdataauth/userdataauth_client.h"
+#include "chromeos/ash/components/login/auth/auth_performer.h"
+#include "chromeos/ash/components/login/auth/public/user_context.h"
+#include "chromeos/ash/components/osauth/public/common_types.h"
+#include "third_party/abseil-cpp/absl/types/optional.h"
+
+namespace ash {
+
+AuthSessionStorageImpl::AuthSessionStorageImpl() {
+  auth_performer_ = std::make_unique<AuthPerformer>(UserDataAuthClient::Get());
+}
+
+AuthSessionStorageImpl::~AuthSessionStorageImpl() = default;
+
+AuthSessionStorageImpl::TokenData::TokenData(
+    std::unique_ptr<UserContext> context)
+    : context(std::move(context)) {}
+AuthSessionStorageImpl::TokenData::~TokenData() = default;
+
+AuthProofToken AuthSessionStorageImpl::Store(
+    std::unique_ptr<UserContext> context) {
+  CHECK(context);
+  auto token = base::UnguessableToken::Create().ToString();
+  tokens_[token] =
+      std::make_unique<AuthSessionStorageImpl::TokenData>(std::move(context));
+  return token;
+}
+
+bool AuthSessionStorageImpl::IsValid(const AuthProofToken& token) {
+  auto data_it = tokens_.find(token);
+  if (data_it == std::end(tokens_)) {
+    return false;
+  }
+  switch (data_it->second->state) {
+    case TokenState::kBorrowed:
+      return !data_it->second->invalidate_on_return;
+    case TokenState::kOwned:
+      return true;
+    case TokenState::kInvalidating:
+      return false;
+  }
+}
+
+std::unique_ptr<UserContext> AuthSessionStorageImpl::Borrow(
+    const base::Location& borrow_location,
+    const AuthProofToken& token) {
+  auto data_it = tokens_.find(token);
+  CHECK(data_it != std::end(tokens_));
+  if (data_it->second->state == TokenState::kBorrowed) {
+    LOG(ERROR) << "Context was already borrowed from "
+               << data_it->second->borrow_location.ToString();
+  }
+  CHECK(data_it->second->state == TokenState::kOwned);
+  data_it->second->state = TokenState::kBorrowed;
+  data_it->second->borrow_location = borrow_location;
+
+  CHECK(data_it->second->context);
+  return std::move(data_it->second->context);
+}
+
+void AuthSessionStorageImpl::Return(const AuthProofToken& token,
+                                    std::unique_ptr<UserContext> context) {
+  CHECK(context);
+  auto data_it = tokens_.find(token);
+  CHECK(data_it != std::end(tokens_));
+  CHECK(data_it->second->state == TokenState::kBorrowed);
+  data_it->second->state = TokenState::kOwned;
+  CHECK(!data_it->second->context);
+  data_it->second->context = std::move(context);
+
+  if (data_it->second->invalidate_on_return) {
+    data_it->second->invalidate_on_return = false;
+    Invalidate(token, std::move(data_it->second->invalidation_closure));
+  }
+}
+
+void AuthSessionStorageImpl::Invalidate(const AuthProofToken& token,
+                                        base::OnceClosure on_invalidated) {
+  auto data_it = tokens_.find(token);
+  CHECK(data_it != std::end(tokens_));
+  if (data_it->second->state == TokenState::kBorrowed) {
+    data_it->second->invalidate_on_return = true;
+    data_it->second->invalidation_closure = std::move(on_invalidated);
+    return;
+  }
+  CHECK(data_it->second->state == TokenState::kOwned);
+  data_it->second->state = TokenState::kInvalidating;
+  auth_performer_->InvalidateAuthSession(
+      std::move(data_it->second->context),
+      base::BindOnce(&AuthSessionStorageImpl::OnSessionInvalidated,
+                     weak_factory_.GetWeakPtr(), token,
+                     std::move(on_invalidated)));
+}
+
+void AuthSessionStorageImpl::OnSessionInvalidated(
+    const AuthProofToken& token,
+    base::OnceClosure on_invalidated,
+    std::unique_ptr<UserContext> context,
+    absl::optional<AuthenticationError> error) {
+  if (error.has_value()) {
+    LOG(ERROR)
+        << "There was an error during attempt to invalidate auth session:"
+        << error.value().get_cryptohome_code();
+  };
+  auto data_it = tokens_.find(token);
+  CHECK(data_it != std::end(tokens_));
+  CHECK(data_it->second->state == TokenState::kInvalidating);
+  tokens_.erase(data_it);
+  std::move(on_invalidated).Run();
+}
+
+}  // namespace ash

chromeos/ash/components/osauth/impl/auth_session_storage_impl.h

@@ -0,0 +1,93 @@
+// Copyright 2023 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROMEOS_ASH_COMPONENTS_OSAUTH_IMPL_AUTH_SESSION_STORAGE_IMPL_H_
+#define CHROMEOS_ASH_COMPONENTS_OSAUTH_IMPL_AUTH_SESSION_STORAGE_IMPL_H_
+
+#include <memory>
+
+#include "base/component_export.h"
+#include "base/containers/flat_map.h"
+#include "base/functional/callback.h"
+#include "base/location.h"
+#include "chromeos/ash/components/login/auth/public/auth_callbacks.h"
+#include "chromeos/ash/components/osauth/public/auth_session_storage.h"
+#include "chromeos/ash/components/osauth/public/common_types.h"
+#include "third_party/abseil-cpp/absl/types/optional.h"
+
+namespace ash {
+
+class AuthenticationError;
+class AuthPerformer;
+class UserContext;
+
+// Helper class that stores and manages lifetime of authenticated UserContext.
+// Main use cases for this class are the situations where authenticated
+// operations do not happen immediately after authentication, but require some
+// user input, e.g. setting up additional factors during user onboarding on a
+// first run, or entering authentication-related section of
+// `chrome://os-settings`.
+//
+// When context is added to storage, storage would return a token as a
+// replacement, this token can be relatively safely be passed between components
+// as it does not contain any sensitive information.
+//
+// UserContext can be borrowed to perform authenticated operations and should be
+// returned to storage as soon as operation completes.
+class COMPONENT_EXPORT(CHROMEOS_ASH_COMPONENTS_OSAUTH) AuthSessionStorageImpl
+    : public AuthSessionStorage {
+ public:
+  AuthSessionStorageImpl();
+  ~AuthSessionStorageImpl() override;
+
+  // AuthSessionStorage implementation:
+  AuthProofToken Store(std::unique_ptr<UserContext> context) override;
+  bool IsValid(const AuthProofToken& token) override;
+  std::unique_ptr<UserContext> Borrow(const base::Location& location,
+                                      const AuthProofToken& token) override;
+  void Return(const AuthProofToken& token,
+              std::unique_ptr<UserContext> context) override;
+  void Invalidate(const AuthProofToken& token,
+                  base::OnceClosure on_invalidated) override;
+
+ private:
+  enum class TokenState {
+    kOwned,         // UserContext is owned by storage
+    kBorrowed,      // UserContext is currently borrowed
+    kInvalidating,  // token is being invalidated
+  };
+
+  struct TokenData {
+    explicit TokenData(std::unique_ptr<UserContext> context);
+    ~TokenData();
+
+    // Context associated with token
+    std::unique_ptr<UserContext> context;
+    TokenState state = TokenState::kOwned;
+
+    // Code location of the last borrow operation.
+    base::Location borrow_location;
+
+    // Data required to invalidate context upon return, if invalidation was
+    // requested while context is borrowed.
+    bool invalidate_on_return = false;
+    base::OnceClosure invalidation_closure;
+  };
+
+  void OnSessionInvalidated(const AuthProofToken& token,
+                            base::OnceClosure on_invalidated,
+                            std::unique_ptr<UserContext> context,
+                            absl::optional<AuthenticationError> error);
+
+  // Stored data for currently active tokens.
+  base::flat_map<AuthProofToken, std::unique_ptr<TokenData>> tokens_;
+
+  std::unique_ptr<AuthPerformer> auth_performer_;
+
+  base::WeakPtrFactory<AuthSessionStorageImpl> weak_factory_{this};
+};
+
+}  // namespace ash
+
+#endif  // CHROMEOS_ASH_COMPONENTS_OSAUTH_IMPL_AUTH_SESSION_STORAGE_IMPL_H_

chromeos/ash/components/osauth/public/BUILD.gn

@@ -10,5 +10,9 @@ assert(is_chromeos_ash,
 
 source_set("public") {
   deps = [ "//base" ]
-  sources = [ "auth_parts.h" ]
+  sources = [
+    "auth_parts.h",
+    "auth_session_storage.h",
+    "common_types.h",
+  ]
 }

chromeos/ash/components/osauth/public/auth_parts.h

@@ -11,6 +11,8 @@
 
 namespace ash {
 
+class AuthSessionStorage;
+
 // Central repository for accessing various OS authentication-related
 // objects.
 // When run normally or as a part of browser_tests it is created and
@@ -27,6 +29,8 @@ class COMPONENT_EXPORT(CHROMEOS_ASH_COMPONENTS_OSAUTH) AuthParts {
   static AuthParts* Get();
 
   virtual ~AuthParts() = default;
+
+  virtual AuthSessionStorage* GetAuthSessionStorage() = 0;
 };
 
 }  // namespace ash