[CritFix] Fix milter commands pipelining

diff --git a/src/libserver/milter.c b/src/libserver/milter.c
index d7e7a6a3dc67db44647da71c1a88e2c5a49008d0..2cced96eca897ff4d32623b0c828bc8c6de5cb5c 100644 (file)
--- a/src/libserver/milter.c
+++ b/src/libserver/milter.c
@@ -644,13 +644,24 @@ rspamd_milter_process_command (struct rspamd_milter_session *session,
                        version, actions, protocol);
                break;
        case RSPAMD_MILTER_CMD_QUIT:
-               msg_debug_milter ("quit command, refcount: %d", session->ref.refcount);
-               priv->state = RSPAMD_MILTER_WANNA_DIE;
-               REF_RETAIN (session);
-               priv->fin_cb (priv->fd, session, priv->ud);
-               REF_RELEASE (session);
+               if (priv->out_chain) {
+                       msg_debug_milter ("quit command, refcount: %d, "
+                                       "some output buffers left - draining",
+                                       session->ref.refcount);
+
+                       priv->state = RSPAMD_MILTER_WRITE_AND_DIE;
+               }
+               else {
+                       msg_debug_milter ("quit command, refcount: %d",
+                                       session->ref.refcount);
 
-               return FALSE;
+                       priv->state = RSPAMD_MILTER_WANNA_DIE;
+                       REF_RETAIN (session);
+                       priv->fin_cb (priv->fd, session, priv->ud);
+                       REF_RELEASE (session);
+
+                       return FALSE;
+               }
                break;
        case RSPAMD_MILTER_CMD_RCPT:
                msg_debug_milter ("rcpt command");
@@ -935,10 +946,25 @@ rspamd_milter_handle_session (struct rspamd_milter_session *session,
                        return rspamd_milter_consume_input (session, priv);
                }
        case RSPAMD_MILTER_WRITE_REPLY:
+       case RSPAMD_MILTER_WRITE_AND_DIE:
                if (priv->out_chain == NULL) {
-                       /* We have written everything, so we can read something */
-                       priv->state = RSPAMD_MILTER_READ_MORE;
-                       rspamd_milter_plan_io (session, priv, EV_READ);
+                       if (priv->state == RSPAMD_MILTER_WRITE_AND_DIE) {
+                               /* Finished writing, let's die finally */
+                               msg_debug_milter ("output drained, terminating, refcount: %d",
+                                               session->ref.refcount);
+
+                               /* Session should be destroyed by fin_cb... */
+                               REF_RETAIN (session);
+                               priv->fin_cb (priv->fd, session, priv->ud);
+                               REF_RELEASE (session);
+
+                               return FALSE;
+                       }
+                       else {
+                               /* We have written everything, so we can read something */
+                               priv->state = RSPAMD_MILTER_READ_MORE;
+                               rspamd_milter_plan_io (session, priv, EV_READ);
+                       }
                }
                else {
                        DL_FOREACH_SAFE (priv->out_chain, obuf, obuf_tmp) {

diff --git a/src/libserver/milter_internal.h b/src/libserver/milter_internal.h
index b2a2dc712c58b5f5b5844680cc5b73d55f6fcff7..9087bdbfb1ce00b799e65fec52be6176281586f7 100644 (file)
--- a/src/libserver/milter_internal.h
+++ b/src/libserver/milter_internal.h
@@ -48,7 +48,8 @@ struct rspamd_milter_outbuf {
 enum rspamd_milter_io_state {
        RSPAMD_MILTER_READ_MORE,
        RSPAMD_MILTER_WRITE_REPLY,
-       RSPAMD_MILTER_WANNA_DIE
+       RSPAMD_MILTER_WANNA_DIE,
+       RSPAMD_MILTER_WRITE_AND_DIE,
 };
 
 struct rspamd_milter_private {