1 files changed, 76 insertions, 29 deletions

diff --git a/src/irc/irc_client.cpp b/src/irc/irc_client.cpp
index d7fa2cd..de38d42 100644
--- a/src/irc/irc_client.cpp
+++ b/src/irc/irc_client.cpp
@@ -135,7 +135,7 @@ IrcClient::IrcClient(std::shared_ptr<Poller>& poller, std::string hostname,
                      std::string realname, std::string user_hostname,
                      Bridge& bridge):
   TCPClientSocketHandler(poller),
-  hostname(std::move(hostname)),
+  hostname(hostname),
   user_hostname(std::move(user_hostname)),
   username(std::move(username)),
   realname(std::move(realname)),
@@ -143,7 +143,14 @@ IrcClient::IrcClient(std::shared_ptr<Poller>& poller, std::string hostname,
   bridge(bridge),
   welcomed(false),
   chanmodes({"", "", "", ""}),
-  chantypes({'#', '&'})
+  chantypes({'#', '&'}),
+  tokens_bucket(this->get_throttle_limit(), 1s, [this]() {
+    if (message_queue.empty())
+      return true;
+    this->actual_send(std::move(this->message_queue.front()));
+    this->message_queue.pop_front();
+    return false;
+  }, "TokensBucket" + this->hostname + this->bridge.get_jid())
 {
 #ifdef USE_DATABASE
   auto options = Database::get_irc_server_options(this->bridge.get_bare_jid(),
@@ -171,6 +178,7 @@ IrcClient::~IrcClient()
   // This event may or may not exist (if we never got connected, it
   // doesn't), but it's ok
   TimedEventsManager::instance().cancel("PING" + this->hostname + this->bridge.get_jid());
+  TimedEventsManager::instance().cancel("TokensBucket" + this->hostname + this->bridge.get_jid());
 }
 
 void IrcClient::start()
@@ -317,9 +325,21 @@ IrcChannel* IrcClient::get_channel(const std::string& n)
     }
   catch (const std::out_of_range& exception)
     {
-      this->channels.emplace(name, std::make_unique<IrcChannel>());
+      return this->channels.emplace(name, std::make_unique<IrcChannel>()).first->second.get();
+    }
+}
+
+const IrcChannel* IrcClient::find_channel(const std::string& n) const
+{
+  const std::string name = utils::tolower(n);
+  try
+    {
       return this->channels.at(name).get();
     }
+  catch (const std::out_of_range& exception)
+    {
+      return nullptr;
+    }
 }
 
 bool IrcClient::is_channel_joined(const std::string& name)
@@ -378,25 +398,39 @@ void IrcClient::parse_in_buffer(const size_t)
     }
 }
 
-void IrcClient::send_message(IrcMessage&& message)
+void IrcClient::actual_send(std::pair<IrcMessage, MessageCallback>&& message_pair)
 {
-  log_debug("IRC SENDING: (", this->get_hostname(), ") ", message);
-  std::string res;
-  if (!message.prefix.empty())
-    res += ":" + std::move(message.prefix) + " ";
-  res += message.command;
-  for (const std::string& arg: message.arguments)
-    {
-      if (arg.find(' ') != std::string::npos ||
-          (!arg.empty() && arg[0] == ':'))
-        {
-          res += " :" + arg;
-          break;
-        }
-      res += " " + arg;
-    }
-  res += "\r\n";
-  this->send_data(std::move(res));
+  const IrcMessage& message = message_pair.first;
+  const MessageCallback& callback = message_pair.second;
+   log_debug("IRC SENDING: (", this->get_hostname(), ") ", message);
+    std::string res;
+    if (!message.prefix.empty())
+      res += ":" + message.prefix + " ";
+    res += message.command;
+    for (const std::string& arg: message.arguments)
+      {
+        if (arg.find(' ') != std::string::npos
+            || (!arg.empty() && arg[0] == ':'))
+          {
+            res += " :" + arg;
+            break;
+          }
+        res += " " + arg;
+      }
+    res += "\r\n";
+    this->send_data(std::move(res));
+
+    if (callback)
+      callback(this, message);
+ }
+
+void IrcClient::send_message(IrcMessage message, MessageCallback callback, bool throttle)
+{
+  auto message_pair = std::make_pair(std::move(message), std::move(callback));
+  if (this->tokens_bucket.use_token() || !throttle)
+    this->actual_send(std::move(message_pair));
+  else
+    message_queue.push_back(std::move(message_pair));
 }
 
 void IrcClient::send_raw(const std::string& txt)
@@ -447,12 +481,12 @@ void IrcClient::send_topic_command(const std::string& chan_name, const std::stri
 
 void IrcClient::send_quit_command(const std::string& reason)
 {
-  this->send_message(IrcMessage("QUIT", {reason}));
+  this->send_message(IrcMessage("QUIT", {reason}), {}, false);
 }
 
 void IrcClient::send_join_command(const std::string& chan_name, const std::string& password)
 {
-  if (this->welcomed == false)
+  if (!this->welcomed)
     {
       const auto it = std::find_if(begin(this->channels_to_join), end(this->channels_to_join),
                                    [&chan_name](const auto& pair) { return std::get<0>(pair) == chan_name; });
@@ -466,10 +500,11 @@ void IrcClient::send_join_command(const std::string& chan_name, const std::strin
   this->start();
 }
 
-bool IrcClient::send_channel_message(const std::string& chan_name, const std::string& body)
+bool IrcClient::send_channel_message(const std::string& chan_name, const std::string& body,
+                                     MessageCallback callback)
 {
   IrcChannel* channel = this->get_channel(chan_name);
-  if (channel->joined == false)
+  if (!channel->joined)
     {
       log_warning("Cannot send message to channel ", chan_name, ", it is not joined");
       return false;
@@ -489,7 +524,7 @@ bool IrcClient::send_channel_message(const std::string& chan_name, const std::st
                          ::strlen(":!@ PRIVMSG ") - chan_name.length() - ::strlen(" :\r\n");
   const auto lines = cut(body, line_size);
   for (const auto& line: lines)
-    this->send_message(IrcMessage("PRIVMSG", {chan_name, line}));
+    this->send_message(IrcMessage("PRIVMSG", {chan_name, line}), callback);
   return true;
 }
 
@@ -1123,8 +1158,6 @@ void IrcClient::on_channel_bad_key(const IrcMessage& message)
 
 void IrcClient::on_channel_mode(const IrcMessage& message)
 {
-  // For now, just transmit the modes so the user can know what happens
-  // TODO, actually interprete the mode.
   Iid iid;
   iid.set_local(message.arguments[0]);
   iid.set_server(this->hostname);
@@ -1142,7 +1175,7 @@ void IrcClient::on_channel_mode(const IrcMessage& message)
     }
   this->bridge.send_message(iid, "", "Mode " + iid.get_local() +
                                       " [" + mode_arguments + "] by " + user.nick,
-                             true);
+                             true, this->is_channel_joined(iid.get_local()));
   const IrcChannel* channel = this->get_channel(iid.get_local());
   if (!channel)
     return;
@@ -1215,6 +1248,11 @@ void IrcClient::on_channel_mode(const IrcMessage& message)
     }
 }
 
+void IrcClient::set_throttle_limit(long int limit)
+{
+  this->tokens_bucket.set_limit(limit);
+}
+
 void IrcClient::on_user_mode(const IrcMessage& message)
 {
   this->bridge.send_xmpp_message(this->hostname, "",
@@ -1252,3 +1290,12 @@ bool IrcClient::abort_on_invalid_cert() const
   return true;
 }
 #endif
+
+long int IrcClient::get_throttle_limit() const
+{
+#ifdef USE_DATABASE
+  return Database::get_irc_server_options(this->bridge.get_bare_jid(), this->hostname).col<Database::ThrottleLimit>();
+#else
+  return 10;
+#endif
+}