smtp/pipeline: Skip body buffering when delivering to a single streaming target

framework/module/streaming_delivery.go

@@ -0,0 +1,44 @@
+/*
+Maddy Mail Server - Composable all-in-one email server.
+Copyright © 2019-2025 Max Mazurov <fox.cpp@disroot.org>, Maddy Mail Server contributors
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+*/
+
+package module
+
+import (
+	"context"
+	"io"
+
+	"github.com/emersion/go-message/textproto"
+)
+
+// StreamingDelivery is an optional interface that may be implemented by the
+// object returned by DeliveryTarget.StartDelivery to receive the message body
+// while it is still being read from the source, without buffering it first.
+type StreamingDelivery interface {
+	// CanStreamBody reports whether BodyStream can be used instead of Body for
+	// the current message. It is called after all AddRcpt calls.
+	CanStreamBody() bool
+
+	// BodyStream is similar to Body of the regular Delivery interface except
+	// that the body is read exactly once, while the message source is still
+	// receiving it.
+	//
+	// If reading body fails, BodyStream must return an error and must make
+	// sure the partially read message is not delivered. The message source
+	// calls Abort after a failed BodyStream.
+	BodyStream(ctx context.Context, header textproto.Header, body io.Reader) error
+}

internal/endpoint/smtp/session.go

@@ -428,7 +428,7 @@ func (s *Session) Logout() error {
 	return nil
 }
 
-func (s *Session) prepareBody(r io.Reader) (textproto.Header, buffer.Buffer, error) {
+func (s *Session) readHeader(r io.Reader) (textproto.Header, io.Reader, error) {
 	limitr := limitReader(r, s.endp.maxHeaderBytes, &exterrors.SMTPError{
 		Code:         552,
 		EnhancedCode: exterrors.EnhancedCode{5, 3, 4},
@@ -451,9 +451,26 @@ func (s *Session) prepareBody(r io.Reader) (textproto.Header, buffer.Buffer, err
 	// the header size check is done. The message size will be checked by go-smtp
 	limitr.Enabled = false
 
-	buf, err := s.endp.buffer(bufr)
+	return header, bufr, nil
+}
+
+func (s *Session) bufferBody(body io.Reader) (buffer.Buffer, error) {
+	buf, err := s.endp.buffer(body)
+	if err != nil {
+		return nil, fmt.Errorf("I/O error while writing buffer: %w", err)
+	}
+	return buf, nil
+}
+
+func (s *Session) prepareBody(r io.Reader) (textproto.Header, buffer.Buffer, error) {
+	header, body, err := s.readHeader(r)
+	if err != nil {
+		return textproto.Header{}, nil, err
+	}
+
+	buf, err := s.bufferBody(body)
 	if err != nil {
-		return textproto.Header{}, nil, fmt.Errorf("I/O error while writing buffer: %w", err)
+		return textproto.Header{}, nil, err
 	}
 
 	return header, buf, nil
@@ -471,27 +488,39 @@ func (s *Session) Data(r io.Reader) error {
 		return s.endp.wrapErr(s.msgMeta.ID, !s.opts.UTF8, "DATA", err)
 	}
 
-	header, buf, err := s.prepareBody(r)
+	header, body, err := s.readHeader(r)
 	if err != nil {
 		return wrapErr(err)
 	}
-	defer func() {
-		if err := buf.Remove(); err != nil {
-			s.log.Error("failed to remove buffered body", err)
-		}
-
-		// go-smtp will call Reset, but it will call Abort if delivery is non-nil.
-		s.cleanSession()
-	}()
 
 	if err := s.checkRoutingLoops(header); err != nil {
 		return wrapErr(err)
 	}
 
-	if err := s.delivery.Body(bodyCtx, header, buf); err != nil {
+	if streamDelivery, ok := s.delivery.(module.StreamingDelivery); ok && streamDelivery.CanStreamBody() {
+		err = streamDelivery.BodyStream(bodyCtx, header, body)
+	} else {
+		var buf buffer.Buffer
+		buf, err = s.bufferBody(body)
+		if err != nil {
+			return wrapErr(err)
+		}
+		defer func() {
+			if err := buf.Remove(); err != nil {
+				s.log.Error("failed to remove buffered body", err)
+			}
+		}()
+
+		err = s.delivery.Body(bodyCtx, header, buf)
+	}
+	if err != nil {
+		// s.delivery is kept so that Reset aborts it.
 		return wrapErr(err)
 	}
 
+	// go-smtp will call Reset, but it will call Abort if delivery is non-nil.
+	defer s.cleanSession()
+
 	if err := s.delivery.Commit(bodyCtx); err != nil {
 		return wrapErr(err)
 	}

internal/endpoint/smtp/smtp_test.go

@@ -451,6 +451,78 @@ func TestSMTPDelivery_AbortData(t *testing.T) {
 	}
 }
 
+func TestSMTPDelivery_StreamBody(t *testing.T) {
+	tgt := testutils.Target{Streaming: true}
+	endp := testEndpoint(t, "smtp", nil, &tgt, nil, nil)
+	defer func() {
+		assert.NoError(t, endp.Stop())
+	}()
+
+	cl, err := smtp.Dial("127.0.0.1:" + testPort)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		assert.NoError(t, cl.Close())
+	}()
+
+	err = submitMsg(t, cl, "sender@example.org", []string{"rcpt1@example.com", "rcpt2@example.com"}, testMsg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(tgt.Messages) != 1 {
+		t.Fatal("Expected a message, got", len(tgt.Messages))
+	}
+	msg := tgt.Messages[0]
+	if !msg.Streamed {
+		t.Error("Message body was not streamed")
+	}
+	testutils.CheckMsgID(t, &msg, "sender@example.org", []string{"rcpt1@example.com", "rcpt2@example.com"}, "")
+	if msg.Header.Get("Received") == "" {
+		t.Error("Missing Received field")
+	}
+}
+
+func TestSMTPDelivery_StreamBody_AbortData(t *testing.T) {
+	tgt := testutils.Target{Streaming: true}
+	endp := testEndpoint(t, "smtp", nil, &tgt, nil, nil)
+	defer func() {
+		assert.NoError(t, endp.Stop())
+	}()
+
+	conn, err := net.Dial("tcp", "127.0.0.1:"+testPort)
+	if err != nil {
+		t.Fatal(err)
+	}
+	cl := smtp.NewClient(conn)
+
+	if err := cl.Hello("mx.example.org"); err != nil {
+		t.Fatal(err)
+	}
+	if err := cl.Mail("sender@example.org", nil); err != nil {
+		t.Fatal(err)
+	}
+	if err := cl.Rcpt("test@example.com", &smtp.RcptOptions{}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := cl.Data(); err != nil {
+		t.Fatal(err)
+	}
+	// Send the header and a part of the body, then close the connection
+	// without sending the final dot.
+	if _, err := conn.Write([]byte(testMsg)); err != nil {
+		t.Fatal(err)
+	}
+	require.NoError(t, conn.Close())
+
+	time.Sleep(250 * time.Millisecond)
+
+	if len(tgt.Messages) != 0 {
+		t.Fatal("Expected no messages, got", len(tgt.Messages))
+	}
+}
+
 func TestSMTPDelivery_EmptyMessage(t *testing.T) {
 	tgt := testutils.Target{}
 	endp := testEndpoint(t, "smtp", nil, &tgt, nil, nil)

internal/msgpipeline/bodystream_test.go

@@ -0,0 +1,128 @@
+/*
+Maddy Mail Server - Composable all-in-one email server.
+Copyright © 2019-2025 Max Mazurov <fox.cpp@disroot.org>, Maddy Mail Server contributors
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+*/
+
+package msgpipeline
+
+import (
+	"context"
+	"testing"
+
+	"github.com/emersion/go-smtp"
+	"github.com/foxcpp/maddy/framework/module"
+	"github.com/foxcpp/maddy/internal/modify"
+	"github.com/foxcpp/maddy/internal/testutils"
+)
+
+func streamTestPipeline(t *testing.T, targets ...module.DeliveryTarget) *MsgPipeline {
+	return &MsgPipeline{
+		msgpipelineCfg: msgpipelineCfg{
+			perSource: map[string]sourceBlock{},
+			defaultSource: sourceBlock{
+				perRcpt: map[string]*rcptBlock{},
+				defaultRcpt: &rcptBlock{
+					targets: targets,
+				},
+			},
+		},
+		Log: testutils.Logger(t, "msgpipeline"),
+	}
+}
+
+func TestMsgPipeline_BodyStream(t *testing.T) {
+	target := testutils.Target{Streaming: true}
+	d := streamTestPipeline(t, &target)
+
+	testutils.DoTestDeliveryStream(t, d, "sender@example.com", []string{"rcpt1@example.com", "rcpt2@example.com"})
+
+	if len(target.Messages) != 1 {
+		t.Fatalf("wrong amount of messages received, want %d, got %d", 1, len(target.Messages))
+	}
+	testutils.CheckTestMessage(t, &target, 0, "sender@example.com", []string{"rcpt1@example.com", "rcpt2@example.com"})
+	if !target.Messages[0].Streamed {
+		t.Error("message was not streamed")
+	}
+}
+
+func TestMsgPipeline_CanStreamBody(t *testing.T) {
+	canStream := func(t *testing.T, d *MsgPipeline) bool {
+		t.Helper()
+
+		delivery, err := d.StartDelivery(context.Background(), &module.MsgMetadata{ID: "test"}, "sender@example.com")
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := delivery.Abort(context.Background()); err != nil {
+				t.Error(err)
+			}
+		}()
+		if err := delivery.AddRcpt(context.Background(), "rcpt@example.com", smtp.RcptOptions{}); err != nil {
+			t.Fatal(err)
+		}
+		return delivery.(module.StreamingDelivery).CanStreamBody()
+	}
+
+	t.Run("single streaming target", func(t *testing.T) {
+		d := streamTestPipeline(t, &testutils.Target{Streaming: true})
+		if !canStream(t, d) {
+			t.Error("expected streaming to be possible")
+		}
+	})
+	t.Run("target without streaming", func(t *testing.T) {
+		d := streamTestPipeline(t, &testutils.Target{})
+		if canStream(t, d) {
+			t.Error("expected streaming to be impossible")
+		}
+	})
+	t.Run("two targets", func(t *testing.T) {
+		d := streamTestPipeline(t,
+			&testutils.Target{Streaming: true, InstName: "a"},
+			&testutils.Target{Streaming: true, InstName: "b"})
+		if canStream(t, d) {
+			t.Error("expected streaming to be impossible")
+		}
+	})
+	t.Run("global check", func(t *testing.T) {
+		d := streamTestPipeline(t, &testutils.Target{Streaming: true})
+		d.globalChecks = []module.Check{&testutils.Check{}}
+		if canStream(t, d) {
+			t.Error("expected streaming to be impossible")
+		}
+	})
+	t.Run("per-recipient check", func(t *testing.T) {
+		d := streamTestPipeline(t, &testutils.Target{Streaming: true})
+		d.defaultSource.defaultRcpt.checks = []module.Check{&testutils.Check{}}
+		if canStream(t, d) {
+			t.Error("expected streaming to be impossible")
+		}
+	})
+	t.Run("source modifier", func(t *testing.T) {
+		d := streamTestPipeline(t, &testutils.Target{Streaming: true})
+		d.defaultSource.modifiers = modify.Group{Modifiers: []module.Modifier{testutils.Modifier{}}}
+		if canStream(t, d) {
+			t.Error("expected streaming to be impossible")
+		}
+	})
+	t.Run("dmarc", func(t *testing.T) {
+		d := streamTestPipeline(t, &testutils.Target{Streaming: true})
+		d.doDMARC = true
+		if canStream(t, d) {
+			t.Error("expected streaming to be impossible")
+		}
+	})
+}