From f4ecbc981eeea6aded153c0f526096b7b0115f0c Mon Sep 17 00:00:00 2001
From: Bron Gondwana <brong@fastmailteam.com>
Date: Thu, 12 Sep 2024 18:16:29 +1000
Subject: [PATCH] lmtp: limit connections by host or by user

The sends a sensible 4xx response back to the user, with an error
message saying which limit was hit
---
 imap/imap_err.et  |  6 ++++++
 imap/lmtp_err.et  |  3 +++
 imap/lmtpd.c      | 25 +++++++++++++++++++++++--
 imap/lmtpengine.c |  5 +++++
 4 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/imap/imap_err.et b/imap/imap_err.et
index d1391199db..2fb6cfeb29 100644
--- a/imap/imap_err.et
+++ b/imap/imap_err.et
@@ -56,6 +56,12 @@ ec IMAP_SYS_ERROR,
 ec IMAP_NOSPACE,
    "mail system storage has been exceeded"
 
+ec IMAP_LIMIT_USER,
+   "Too many connections from this user"
+
+ec IMAP_LIMIT_HOST,
+   "Too many connections from this host"
+
 ec IMAP_PERMISSION_DENIED,
    "Permission denied"
 
diff --git a/imap/lmtp_err.et b/imap/lmtp_err.et
index 4e27013035..da87362d76 100644
--- a/imap/lmtp_err.et
+++ b/imap/lmtp_err.et
@@ -62,6 +62,9 @@ ec LMTP_SERVER_FULL,
 ec LMTP_SERVER_FAILURE,
    "451 4.4.3 %s"
 
+ec LMTP_SERVER_BUSY,
+   "451 4.4.5 %s"
+
 ec LMTP_MAILBOX_FULL,
    "452 4.2.2 %s SESSIONID=<%s>"
 
diff --git a/imap/lmtpd.c b/imap/lmtpd.c
index c322a7d3e2..3b6115dd1a 100644
--- a/imap/lmtpd.c
+++ b/imap/lmtpd.c
@@ -88,6 +88,7 @@
 #include "mupdate.h"
 #include "notify.h"
 #include "prometheus.h"
+#include "proc.h"
 #include "prot.h"
 #include "proxy.h"
 #include "slowio.h"
@@ -151,6 +152,8 @@ static int dupelim = 1;         /* eliminate duplicate messages with
 static int singleinstance = 1;  /* attempt single instance store */
 static int isproxy = 0;
 static strarray_t *excluded_specialuse = NULL;
+static const char *lmtpd_clienthost = "[local]";
+static struct proc_handle *proc_handle = NULL;
 
 static struct stagemsg *stage = NULL;
 
@@ -256,6 +259,7 @@ int service_init(int argc __attribute__((unused)),
 int service_main(int argc, char **argv,
                  char **envp __attribute__((unused)))
 {
+    const char *localip, *remoteip;
     int opt;
 
     struct io_count *io_count_start = NULL;
@@ -276,6 +280,9 @@ int service_main(int argc, char **argv,
     prot_setflushonread(deliver_in, deliver_out);
     prot_settimeout(deliver_in, 360);
 
+    lmtpd_clienthost = get_clienthost(0, &localip, &remoteip);
+    proc_register(&proc_handle, 0, config_ident, lmtpd_clienthost, NULL, NULL, NULL);
+
     while ((opt = getopt(argc, argv, "Ha")) != EOF) {
         switch(opt) {
         case 'H': /* expect HAProxy protocol header */
@@ -327,6 +334,9 @@ int service_main(int argc, char **argv,
 
     prometheus_increment(CYRUS_LMTP_READY_LISTENERS);
 
+    proc_cleanup(&proc_handle);
+    lmtpd_clienthost = "[local]";
+
     slowio_reset();
 
     return 0;
@@ -883,11 +893,21 @@ int deliver(message_data_t *msgdata, char *authuser,
             strarray_t flags = STRARRAY_INITIALIZER;
             struct imap4flags imap4flags = { &flags, authstate };
 
+            const char *userid = mbname_userid(mbname);
+            struct proc_limits limits;
+            limits.servicename = config_ident;
+            limits.clienthost = lmtpd_clienthost;
+            limits.userid = userid;
+            r = proc_checklimits(&limits);
+            if (r) goto setstatus;
+
+            proc_register(&proc_handle, 0, config_ident, lmtpd_clienthost, userid, NULL, NULL);
+
             // lock conversations for the duration of delivery, so nothing else can read
             // the state of any mailbox while the delivery is half done
             struct conversations_state *state = NULL;
-            if (mbname_userid(mbname)) {
-                r = conversations_open_user(mbname_userid(mbname), 0/*shared*/, &state);
+            if (userid) {
+                r = conversations_open_user(userid, 0/*shared*/, &state);
                 if (r) goto setstatus;
             }
 
@@ -919,6 +939,7 @@ int deliver(message_data_t *msgdata, char *authuser,
             }
             strarray_fini(&flags);
             conversations_commit(&state);
+            proc_register(&proc_handle, 0, config_ident, lmtpd_clienthost, NULL, NULL, NULL);
         }
 
         telemetry_rusage(mbname_userid(mbname));
diff --git a/imap/lmtpengine.c b/imap/lmtpengine.c
index 3225623b87..5e3d9de15e 100644
--- a/imap/lmtpengine.c
+++ b/imap/lmtpengine.c
@@ -149,6 +149,11 @@ static void send_lmtp_error(struct protstream *pout, int r, strarray_t *resp)
         code = LMTP_OK;
         break;
 
+    case IMAP_LIMIT_HOST:
+    case IMAP_LIMIT_USER:
+        code = LMTP_SERVER_BUSY;
+        break;
+
     case IMAP_SERVER_UNAVAILABLE:
     case MUPDATE_NOCONN:
     case MUPDATE_NOAUTH: