Fix: Return 404 response when Lua script is not found

contrib/imap-server/imap-server.go

@@ -18,9 +18,11 @@ package main
 import (
 	"crypto/tls"
 	"errors"
+	"fmt"
 	"log"
 	"net"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/emersion/go-imap"
@@ -78,10 +80,13 @@ func (u *User) Logout() error {
 	return nil
 }
 
+var _ backend.User = (*User)(nil)
+
 type Backend struct{}
 
 func (b *Backend) Login(connInfo *imap.ConnInfo, username string, password string) (user backend.User, err error) {
 	_ = password
+
 	log.Println("Connect from", connInfo.RemoteAddr.String(), "username", username)
 
 	user = &User{name: username}
@@ -91,50 +96,176 @@ func (b *Backend) Login(connInfo *imap.ConnInfo, username string, password strin
 	return user, errContactSupport
 }
 
-func main() {
-	be := &Backend{}
+var _ backend.Backend = (*Backend)(nil)
 
-	// Create a new server
-	s := server.New(be)
+type ProxyAndTLSListener struct {
+	ProxyListener *proxyproto.Listener
+	TLSConfig     *tls.Config
+}
 
-	address := os.Getenv("FAKE_IMAP_SERVER_ADDRESS")
-	tlsCert := os.Getenv("FAKE_IMAP_SERVER_TLSCERT")
-	tlsKey := os.Getenv("FAKE_IMAP_SERVER_TLSKEY")
+func (p *ProxyAndTLSListener) Accept() (net.Conn, error) {
+	rawConn, err := p.ProxyListener.Accept()
+	if err != nil {
+		return nil, fmt.Errorf("failed to accept connection: %w", err)
+	}
 
-	if tlsCert != "" && tlsKey != "" {
-		cer, err := tls.LoadX509KeyPair(tlsCert, tlsKey)
-		if err != nil {
-			log.Println(err)
+	tlsConn := tls.Server(rawConn, p.TLSConfig)
 
-			return
-		}
+	return tlsConn, nil
+}
+
+func (p *ProxyAndTLSListener) Close() error {
+	return p.ProxyListener.Close()
+}
+
+func (p *ProxyAndTLSListener) Addr() net.Addr {
+	return p.ProxyListener.Addr()
+}
+
+var _ net.Listener = (*ProxyAndTLSListener)(nil)
 
-		s.TLSConfig = &tls.Config{
-			MinVersion:   tls.VersionTLS12,
-			Certificates: []tls.Certificate{cer},
+func NewProxyAndTLSListener(rawListener net.Listener, tlsConfig *tls.Config) net.Listener {
+	proxyListener := &proxyproto.Listener{
+		Listener: rawListener,
+		ConnPolicy: func(opts proxyproto.ConnPolicyOptions) (proxyproto.Policy, error) {
+			return proxyproto.REQUIRE, nil
+		},
+	}
+
+	return &ProxyAndTLSListener{
+		ProxyListener: proxyListener,
+		TLSConfig:     tlsConfig,
+	}
+}
+
+type ServerConfig struct {
+	Address   string
+	TLSConfig *tls.Config
+}
+
+type IMAPType uint
+
+const (
+	IMAP IMAPType = iota
+	IMAPS
+)
+
+type IMAPServer struct {
+	serverDescription string
+	serverType        IMAPType
+	config            *ServerConfig
+	server            *server.Server
+}
+
+func NewIMAPServer(serverType IMAPType, serverDescription string, address string, backend backend.Backend) *IMAPServer {
+	return &IMAPServer{
+		serverDescription: serverDescription,
+		serverType:        serverType,
+		config: &ServerConfig{
+			Address:   address,
+			TLSConfig: configureTLS(),
+		},
+		server: server.New(backend),
+	}
+}
+
+func (s *IMAPServer) Start(wg *sync.WaitGroup) {
+	var listener net.Listener
+
+	defer wg.Done()
+
+	log.Printf("Starting %s server at %s", s.serverDescription, s.config.Address)
+
+	rawListener, err := net.Listen("tcp", s.config.Address)
+	if err != nil {
+		log.Fatalf("Failed to start %s server: %v", s.serverDescription, err)
+	}
+
+	if s.serverType == IMAPS {
+		listener = NewProxyAndTLSListener(rawListener, s.config.TLSConfig)
+	} else {
+		listener = &proxyproto.Listener{
+			Listener: rawListener,
+			ConnPolicy: func(opts proxyproto.ConnPolicyOptions) (proxyproto.Policy, error) {
+				return proxyproto.REQUIRE, nil
+			},
 		}
+
+		s.server.TLSConfig = s.config.TLSConfig
 	}
 
-	if address == "" {
-		address = "127.0.0.1:10143"
+	if err := s.server.Serve(listener); err != nil {
+		log.Fatalf("%s server stopped unexpectedly: %v", s.serverDescription, err)
 	}
+}
+
+type ServerManager struct {
+	wg    *sync.WaitGroup
+	imap  *IMAPServer
+	imaps *IMAPServer
+}
+
+func NewServerManager(backend *Backend) *ServerManager {
+	imapServer := NewIMAPServer(
+		IMAP,
+		"IMAP (StartTLS)",
+		getEnvWithDefault("FAKE_IMAP_SERVER_ADDRESS", "127.0.0.1:10143"),
+		backend,
+	)
+	imapsServer := NewIMAPServer(
+		IMAPS,
+		"IMAPS",
+		getEnvWithDefault("FAKE_IMAPS_SERVER_ADDRESS", "127.0.0.1:10993"),
+		backend,
+	)
+
+	return &ServerManager{
+		wg:    &sync.WaitGroup{},
+		imap:  imapServer,
+		imaps: imapsServer,
+	}
+}
 
-	s.Addr = address
-	s.AllowInsecureAuth = true
+func (m *ServerManager) StartAll() {
+	m.wg.Add(2)
 
-	log.Println("Starting IMAP server at", s.Addr)
+	go m.imap.Start(m.wg)
+	go m.imaps.Start(m.wg)
 
-	// Set up listener
-	listener, err := net.Listen("tcp", s.Addr)
+	m.wg.Wait()
+}
+
+func configureTLS() *tls.Config {
+	serverName := os.Getenv("FAKE_IMAP_SERVER_NAME")
+	tlsCert := os.Getenv("FAKE_IMAP_SERVER_TLSCERT")
+	tlsKey := os.Getenv("FAKE_IMAP_SERVER_TLSKEY")
+
+	if tlsCert == "" || tlsKey == "" {
+		log.Fatal("TLS certificate and key must be provided for IMAPS")
+	}
+
+	tlsCertificate, err := tls.LoadX509KeyPair(tlsCert, tlsKey)
 	if err != nil {
-		log.Fatal(err)
+		log.Fatalf("Failed to load TLS certificate: %v", err)
 	}
 
-	// Wrap listener in proxyproto
-	proxyListener := &proxyproto.Listener{Listener: listener}
+	return &tls.Config{
+		MinVersion:   tls.VersionTLS12,
+		Certificates: []tls.Certificate{tlsCertificate},
+		ServerName:   serverName,
+	}
+}
 
-	// Start server
-	if err := s.Serve(proxyListener); err != nil {
-		log.Fatal(err)
+func getEnvWithDefault(envVar, defaultValue string) string {
+	if value := os.Getenv(envVar); value != "" {
+		return value
 	}
+
+	return defaultValue
+}
+
+func main() {
+	manager := NewServerManager(&Backend{})
+
+	manager.StartAll()
 }