From bf8e025b2c1a969dfe7440e12fb4b5644db59359 Mon Sep 17 00:00:00 2001
From: josh <josh@elixxir.io>
Date: Wed, 15 Dec 2021 15:55:06 -0800
Subject: [PATCH] Update bucket for every CMIX message sent

---
 api/send.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/api/send.go b/api/send.go
index 3c6cc0178..08531b696 100644
--- a/api/send.go
+++ b/api/send.go
@@ -50,6 +50,7 @@ func (c *Client) SendUnsafe(m message.Send, param params.Unsafe) ([]id.Round,
 func (c *Client) SendCMIX(msg format.Message, recipientID *id.ID,
 	param params.CMIX) (id.Round, ephemeral.Id, error) {
 	jww.INFO.Printf("SendCMIX(%s)", string(msg.GetContents()))
+	c.OnSend(1)
 	return c.network.SendCMIX(msg, recipientID, param)
 }
 
@@ -58,6 +59,7 @@ func (c *Client) SendCMIX(msg format.Message, recipientID *id.ID,
 // round ID of the round the payload was sent or an error if it fails.
 func (c *Client) SendManyCMIX(messages []message.TargetedCmixMessage,
 	params params.CMIX) (id.Round, []ephemeral.Id, error) {
+	c.OnSend(uint32(len(messages)))
 	return c.network.SendManyCMIX(messages, params)
 }
 
@@ -73,3 +75,13 @@ func (c *Client) NewCMIXMessage(contents []byte) (format.Message, error) {
 	msg.SetContents(contents)
 	return msg, nil
 }
+
+// OnSend performs a bucket addition on a call to Client.SendCMIX or
+// Client.SendManyCMIX, updating the bucket for the amount of messages sent.
+func (c *Client) OnSend(messages uint32) {
+	rateLimitingParam := c.storage.GetBucketParams().Get()
+	c.storage.GetBucket().AddWithExternalParams(messages,
+		rateLimitingParam.Capacity, rateLimitingParam.LeakedTokens,
+		rateLimitingParam.LeakDuration)
+
+}
-- 
GitLab