Move automatic token removal from xdg-activation to the token authority

src/include/server/mir/shell/token_authority.h

@@ -17,9 +17,17 @@
 #ifndef MIR_SHELL_TOKEN_AUTHORITY_H_
 #define MIR_SHELL_TOKEN_AUTHORITY_H_
 
+#include "mir/main_loop.h"
+#include "mir/time/alarm.h"
+
+#include <chrono>
+#include <cmath>
+#include <functional>
+#include <memory>
 #include <mutex>
-#include <set>
 #include <string>
+#include <unistd.h>
+#include <unordered_set>
 
 namespace mir
 {
@@ -28,32 +36,102 @@ namespace shell
 class TokenAuthority
 {
 public:
-    auto issue_token() -> std::string
+    class Token
+    {
+    public:
+        using RevocationListener = std::function<void(Token const& token)>;
+
+        Token(
+            std::string token,
+            std::unique_ptr<time::Alarm> alarm,
+            std::optional<RevocationListener> revocation_listener) :
+            token{token},
+            alarm{std::move(alarm)},
+            revocation_listener{revocation_listener}
+        {
+            alarm->reschedule_in(timeout_ms);
+        }
+
+        explicit Token(std::string token) :
+            token{token},
+            alarm{nullptr}
+        {
+        }
+
+        operator std::string() const
+        {
+            return token;
+        }
+
+        bool operator==(Token const& token) const
+        {
+            return this->token == token.token;
+        }
+
+        std::size_t hash() const
+        {
+            return std::hash<std::string>{}(token);
+        }
+
+    private:
+        friend TokenAuthority;
+        std::string token;
+        std::shared_ptr<mir::time::Alarm> const alarm;
+        std::optional<RevocationListener> revocation_listener;
+    };
+
+    TokenAuthority(std::shared_ptr<MainLoop> main_loop) :
+        main_loop{main_loop}
     {
-        std::scoped_lock lock{mutex};
-        auto const [iter, _] = issued_tokens.insert(generate_token());
-        return *iter;
     }
 
-    auto verify_token(std::string const& token) const -> bool
+    auto issue_token(std::optional<Token::RevocationListener> revocation_listener) -> Token
     {
         std::scoped_lock lock{mutex};
-        return issued_tokens.contains(token);
+
+        auto const generated = generate_token();
+        auto alarm = main_loop->create_alarm(
+            [this, generated]()
+            {
+                revoke_token(generated);
+            });
+
+        auto const [iter, _] = issued_tokens.insert(Token{generated ,std::move(alarm), revocation_listener});
+        return *iter;
     }
 
     void revoke_token(std::string to_remove)
     {
         std::scoped_lock lock{mutex};
-        issued_tokens.erase(to_remove);
+        auto const iter = issued_tokens.find(Token{to_remove});
+        if(iter != issued_tokens.end())
+        {
+            if(auto cb = iter->revocation_listener)
+                (*cb)(*iter);
+            issued_tokens.erase(iter);
+        }
     }
 
 private:
     static auto generate_token() -> std::string;
 
-    std::set<std::string> issued_tokens;
+    struct TokenHash
+    {
+        std::size_t operator()(mir::shell::TokenAuthority::Token const& s) const noexcept
+        {
+            return s.hash();
+        }
+    };
+    std::unordered_set<Token, TokenHash> issued_tokens;
+
+    std::shared_ptr<MainLoop> main_loop;
     std::mutex mutable mutex;
+
+    /// Time in milliseconds that the compositor will wait before invalidating a token
+    static auto constexpr timeout_ms = std::chrono::seconds(3000);
 };
 }
 }
 
+
 #endif

src/server/frontend_wayland/xdg_activation_v1.cpp

@@ -14,9 +14,9 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
+#include "mir/events/keyboard_event.h"
 #include "mir/shell/token_authority.h"
 #include "xdg_activation_v1.h"
-#include "mir/events/keyboard_event.h"
 #include "mir/input/keyboard_observer.h"
 #include "mir/log.h"
 #include "mir/main_loop.h"
@@ -43,11 +43,6 @@ namespace ms = mir::scene;
 namespace msh = mir::shell;
 namespace mi = mir::input;
 
-namespace
-{
-/// Time in milliseconds that the compositor will wait before invalidating a token
-auto constexpr timeout_ms = std::chrono::seconds(3000);
-}
 
 namespace mir
 {
@@ -57,17 +52,13 @@ struct XdgActivationTokenData
 {
     XdgActivationTokenData(
         std::string const& token,
-        std::unique_ptr<time::Alarm> alarm_,
         std::shared_ptr<ms::Session> const& session)
         : token{token},
-        alarm{std::move(alarm_)},
         session{session}
     {
-        alarm->reschedule_in(timeout_ms);
     }
 
     std::string const token;
-    std::unique_ptr<time::Alarm> const alarm;
     std::weak_ptr<ms::Session> session;
 
     std::optional<uint32_t> serial;
@@ -139,11 +130,13 @@ mf::XdgActivationV1::~XdgActivationV1()
 
 std::shared_ptr<mf::XdgActivationTokenData> const& mf::XdgActivationV1::create_token(std::shared_ptr<ms::Session> const& session)
 {
-    auto generated = token_authority->issue_token();
-    auto token = std::make_shared<XdgActivationTokenData>(generated, main_loop->create_alarm([this, generated]()
-    {
-        try_consume_token(generated);
-    }), session);
+    auto generated = token_authority->issue_token(
+        [this](auto const& token)
+        {
+            this->try_consume_token(token);
+        });
+
+    auto token = std::make_shared<XdgActivationTokenData>(generated, session);
 
     {
         std::lock_guard guard(pending_tokens_mutex);
@@ -168,7 +161,6 @@ std::shared_ptr<mf::XdgActivationTokenData> mf::XdgActivationV1::try_consume_tok
         if (iter != pending_tokens.end())
         {
             auto result = *iter;
-            token_authority->revoke_token(iter->get()->token);
             pending_tokens.erase(iter);
             return result;
         }
@@ -182,7 +174,7 @@ void mf::XdgActivationV1::invalidate_all()
     std::lock_guard guard(pending_tokens_mutex);
     for (auto const& it : pending_tokens)
         token_authority->revoke_token(it->token);
-    pending_tokens.clear();
+    // No need to clear pending_tokens, since revoke_token() erases each token
 }
 
 void mf::XdgActivationV1::invalidate_if_not_from_session(std::shared_ptr<ms::Session> const& session)
@@ -196,8 +188,6 @@ void mf::XdgActivationV1::invalidate_if_not_from_session(std::shared_ptr<ms::Ses
     std::lock_guard guard(pending_tokens_mutex);
     for (auto const& expired_token : pending_tokens | srv::filter(is_invalid))
         token_authority->revoke_token(expired_token->token);
-
-    std::erase_if(pending_tokens, is_invalid);
 }
 
 void mf::XdgActivationV1::bind(struct wl_resource* resource)
@@ -270,13 +260,18 @@ void mf::XdgActivationV1::Instance::activate(std::string const& token, struct wl
     // 2. The surface failed to use the token in the alotted period of time
     // 3. A key was pressed down between the issuing of the token and the activation
     //    of the surface with that token
+    //
     auto xdg_token = xdg_activation_v1->try_consume_token(token);
     if (!xdg_token)
     {
         mir::log_error("XdgActivationV1::activate invalid token: %s", token.c_str());
         return;
     }
 
+    // try_consume_token removes the token from our list of tokens, need to
+    // inform the token authority as well.
+    xdg_activation_v1->token_authority->revoke_token(token);
+
     main_loop->enqueue(this, [
         surface=surface,
         shell=shell,

src/server/server.cpp

@@ -483,8 +483,7 @@ auto mir::Server::get_activation_token() const -> std::optional<std::string>
 {
     if (auto const config = self->server_config)
     {
-        // TODO We should revoke the token at some time
-        return config->the_token_authority()->issue_token();
+        return config->the_token_authority()->issue_token({});
     }
 
     BOOST_THROW_EXCEPTION(std::logic_error("Cannot open connection when not running"));

src/server/shell/default_configuration.cpp

@@ -150,9 +150,9 @@ auto mir::DefaultServerConfiguration::the_idle_handler() -> std::shared_ptr<msh:
 
 auto mir::DefaultServerConfiguration::the_token_authority() -> std::shared_ptr<msh::TokenAuthority>
 {
-    return token_authority([]()
+    return token_authority([this]()
         {
-            return std::make_shared<shell::TokenAuthority>();
+            return std::make_shared<shell::TokenAuthority>(the_main_loop());
         });
 }