Files
RenewCMS/test/smtp_server_support_test.go
T

262 lines
6.3 KiB
Go

package test
import (
"bufio"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"encoding/base64"
"math/big"
"net"
"strings"
"sync"
"testing"
"time"
)
type fakeSMTPOptions struct {
authMechanisms string
startTLS bool
rejectRcpt bool
rejectData bool
implicitTLS bool
listenAddress string
}
type fakeSMTPSession struct {
commands []string
from string
to string
data string
username string
password string
usedTLS bool
}
type fakeSMTPServer struct {
listener net.Listener
options fakeSMTPOptions
cert tls.Certificate
roots *x509.CertPool
mu sync.Mutex
sessions []*fakeSMTPSession
wg sync.WaitGroup
}
func newFakeSMTPServer(t *testing.T, options fakeSMTPOptions) *fakeSMTPServer {
t.Helper()
address := options.listenAddress
if address == "" {
address = "127.0.0.1:0"
}
listener, err := net.Listen("tcp", address)
if err != nil {
t.Skipf("cannot listen on %s: %v", address, err)
}
cert, roots := selfSignedLoopbackCert(t)
server := &fakeSMTPServer{listener: listener, options: options, cert: cert, roots: roots}
server.wg.Add(1)
go server.acceptLoop()
t.Cleanup(func() {
_ = listener.Close()
server.wg.Wait()
})
return server
}
func selfSignedLoopbackCert(t *testing.T) (tls.Certificate, *x509.CertPool) {
t.Helper()
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
t.Fatal(err)
}
template := &x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: "127.0.0.1"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(time.Hour),
KeyUsage: x509.KeyUsageDigitalSignature,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
}
der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
if err != nil {
t.Fatal(err)
}
parsed, err := x509.ParseCertificate(der)
if err != nil {
t.Fatal(err)
}
roots := x509.NewCertPool()
roots.AddCert(parsed)
return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}, roots
}
func (s *fakeSMTPServer) port() int {
return s.listener.Addr().(*net.TCPAddr).Port
}
func (s *fakeSMTPServer) clientTLS() *tls.Config {
return &tls.Config{RootCAs: s.roots}
}
func (s *fakeSMTPServer) received() []*fakeSMTPSession {
s.mu.Lock()
defer s.mu.Unlock()
return append([]*fakeSMTPSession(nil), s.sessions...)
}
func (s *fakeSMTPServer) acceptLoop() {
defer s.wg.Done()
for {
connection, err := s.listener.Accept()
if err != nil {
return
}
s.wg.Add(1)
go func() {
defer s.wg.Done()
s.serve(connection)
}()
}
}
func (s *fakeSMTPServer) serve(connection net.Conn) {
defer func() { _ = connection.Close() }()
_ = connection.SetDeadline(time.Now().Add(10 * time.Second))
session := &fakeSMTPSession{}
s.mu.Lock()
s.sessions = append(s.sessions, session)
s.mu.Unlock()
if s.options.implicitTLS {
connection = tls.Server(connection, &tls.Config{Certificates: []tls.Certificate{s.cert}})
session.usedTLS = true
}
reader := bufio.NewReader(connection)
reply := func(line string) { _, _ = connection.Write([]byte(line + "\r\n")) }
reply("220 fake ESMTP")
for {
line, err := reader.ReadString('\n')
if err != nil {
return
}
line = strings.TrimRight(line, "\r\n")
s.record(session, line)
verb := strings.ToUpper(strings.Fields(line + " ")[0])
switch verb {
case "EHLO":
s.replyEhlo(reply, session)
case "STARTTLS":
reply("220 ready")
connection = tls.Server(connection, &tls.Config{Certificates: []tls.Certificate{s.cert}})
reader = bufio.NewReader(connection)
reply = func(line string) { _, _ = connection.Write([]byte(line + "\r\n")) }
session.usedTLS = true
case "AUTH":
if !s.authenticate(line, reader, reply, session) {
return
}
case "MAIL":
session.from = smtpAngleContent(line, "<", ">")
reply("250 ok")
case "RCPT":
if s.options.rejectRcpt {
reply("550 no such user")
continue
}
session.to = smtpAngleContent(line, "<", ">")
reply("250 ok")
case "DATA":
if s.options.rejectData {
reply("554 refused")
continue
}
reply("354 go ahead")
session.data = readSMTPData(reader)
reply("250 queued")
case "QUIT":
reply("221 bye")
return
default:
reply("250 ok")
}
}
}
func (s *fakeSMTPServer) record(session *fakeSMTPSession, line string) {
s.mu.Lock()
session.commands = append(session.commands, line)
s.mu.Unlock()
}
func (s *fakeSMTPServer) replyEhlo(reply func(string), session *fakeSMTPSession) {
lines := []string{"250-fake"}
if s.options.startTLS && !session.usedTLS {
lines = append(lines, "250-STARTTLS")
}
if s.options.authMechanisms != "" {
lines = append(lines, "250-AUTH "+s.options.authMechanisms)
}
lines = append(lines, "250 8BITMIME")
for _, line := range lines {
reply(line)
}
}
func (s *fakeSMTPServer) authenticate(line string, reader *bufio.Reader, reply func(string), session *fakeSMTPSession) bool {
fields := strings.Fields(line)
if len(fields) < 2 {
reply("501 syntax")
return false
}
switch strings.ToUpper(fields[1]) {
case "PLAIN":
raw, _ := base64.StdEncoding.DecodeString(fields[2])
parts := strings.Split(string(raw), "\x00")
session.username, session.password = parts[1], parts[2]
case "LOGIN":
reply("334 " + base64.StdEncoding.EncodeToString([]byte("Username:")))
username, _ := reader.ReadString('\n')
reply("334 " + base64.StdEncoding.EncodeToString([]byte("Password:")))
password, _ := reader.ReadString('\n')
session.username = decodeBase64Line(username)
session.password = decodeBase64Line(password)
default:
reply("504 unsupported")
return false
}
reply("235 ok")
return true
}
func decodeBase64Line(line string) string {
raw, _ := base64.StdEncoding.DecodeString(strings.TrimSpace(line))
return string(raw)
}
func smtpAngleContent(line string, open string, close string) string {
start := strings.Index(line, open)
end := strings.LastIndex(line, close)
if start < 0 || end <= start {
return ""
}
return line[start+1 : end]
}
func readSMTPData(reader *bufio.Reader) string {
var data strings.Builder
for {
line, err := reader.ReadString('\n')
if err != nil || line == ".\r\n" {
return data.String()
}
data.WriteString(strings.TrimPrefix(line, "."))
}
}