From 8f28babc473bf970567ad7cf99fcfb255c00e001 Mon Sep 17 00:00:00 2001
From: Winni Neessen <wn@neessen.net>
Date: Sat, 9 Nov 2024 14:58:23 +0100
Subject: [PATCH] Add tests for Client StartTLS functionality

Introduce new tests to cover the Client's behavior when initiating a STARTTLS session under different conditions: normal operation, failure on EHLO/HELO, and a server not supporting STARTTLS. This ensures robustness in handling STARTTLS interactions.
---
 smtp/smtp_test.go | 102 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 102 insertions(+)

diff --git a/smtp/smtp_test.go b/smtp/smtp_test.go
index d7a2a43..9dc3c10 100644
--- a/smtp/smtp_test.go
+++ b/smtp/smtp_test.go
@@ -1691,6 +1691,108 @@ func TestClient_cmd(t *testing.T) {
 	})
 }
 
+func TestClient_StartTLS(t *testing.T) {
+	t.Run("normal STARTTLS should succeed", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		PortAdder.Add(1)
+		serverPort := int(TestServerPortBase + PortAdder.Load())
+		featureSet := "250-STARTTLS\r\n250-8BITMIME\r\n250-DSN\r\n250 SMTPUTF8"
+		go func() {
+			if err := simpleSMTPServer(ctx, t, &serverProps{
+				FeatureSet: featureSet,
+				ListenPort: serverPort,
+			},
+			); err != nil {
+				t.Errorf("failed to start test server: %s", err)
+				return
+			}
+		}()
+		time.Sleep(time.Millisecond * 30)
+
+		client, err := Dial(fmt.Sprintf("%s:%d", TestServerAddr, serverPort))
+		if err != nil {
+			t.Errorf("failed to dial to test server: %s", err)
+		}
+		t.Cleanup(func() {
+			if err = client.Close(); err != nil {
+				t.Errorf("failed to close client: %s", err)
+			}
+		})
+		tlsConfig := &tls.Config{InsecureSkipVerify: true}
+		if err = client.StartTLS(tlsConfig); err != nil {
+			t.Errorf("failed to initialize STARTTLS session: %s", err)
+		}
+	})
+	t.Run("STARTTLS fails on EHLO/HELO", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		PortAdder.Add(1)
+		serverPort := int(TestServerPortBase + PortAdder.Load())
+		featureSet := "250-STARTTLS\r\n250-8BITMIME\r\n250-DSN\r\n250 SMTPUTF8"
+		go func() {
+			if err := simpleSMTPServer(ctx, t, &serverProps{
+				FailOnEhlo: true,
+				FailOnHelo: true,
+				FeatureSet: featureSet,
+				ListenPort: serverPort,
+			},
+			); err != nil {
+				t.Errorf("failed to start test server: %s", err)
+				return
+			}
+		}()
+		time.Sleep(time.Millisecond * 30)
+
+		client, err := Dial(fmt.Sprintf("%s:%d", TestServerAddr, serverPort))
+		if err != nil {
+			t.Errorf("failed to dial to test server: %s", err)
+		}
+		t.Cleanup(func() {
+			if err = client.Close(); err != nil {
+				t.Errorf("failed to close client: %s", err)
+			}
+		})
+		tlsConfig := &tls.Config{InsecureSkipVerify: true}
+		if err = client.StartTLS(tlsConfig); err == nil {
+			t.Error("STARTTLS should fail on EHLO")
+		}
+	})
+	t.Run("STARTTLS fails on server not supporting STARTTLS", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		PortAdder.Add(1)
+		serverPort := int(TestServerPortBase + PortAdder.Load())
+		featureSet := "250-8BITMIME\r\n250-DSN\r\n250 SMTPUTF8"
+		go func() {
+			if err := simpleSMTPServer(ctx, t, &serverProps{
+				FailOnSTARTTLS: true,
+				FeatureSet:     featureSet,
+				ListenPort:     serverPort,
+			},
+			); err != nil {
+				t.Errorf("failed to start test server: %s", err)
+				return
+			}
+		}()
+		time.Sleep(time.Millisecond * 30)
+
+		client, err := Dial(fmt.Sprintf("%s:%d", TestServerAddr, serverPort))
+		if err != nil {
+			t.Errorf("failed to dial to test server: %s", err)
+		}
+		t.Cleanup(func() {
+			if err = client.Close(); err != nil {
+				t.Errorf("failed to close client: %s", err)
+			}
+		})
+		tlsConfig := &tls.Config{InsecureSkipVerify: true}
+		if err = client.StartTLS(tlsConfig); err == nil {
+			t.Error("STARTTLS should fail for server not supporting it")
+		}
+	})
+}
+
 // Issue 17794: don't send a trailing space on AUTH command when there's no password.
 func TestClient_Auth_trimSpace(t *testing.T) {
 	server := "220 hello world\r\n" +