add custom dialer

internal/alerts/alerts_api_test.go

@@ -531,7 +531,7 @@ func TestSendTestNotification(t *testing.T) {
 		},
 	}
 
-	for _, url := range []string{localURL, "smtp://user:pass@consul", "mqtt://consul/topic"} {
+	for _, url := range []string{localURL, "smtp://user:pass@127.0.0.1/?fromAddress=sender@example.com&toAddresses=recipient@example.com", "mqtt://127.0.0.1/topic"} {
 		scenarios = append(scenarios, beszelTests.ApiScenario{
 			Name:            "readonly cannot send to " + url,
 			Method:          http.MethodPost,

internal/alerts/notification_client.go

@@ -1,8 +1,10 @@
 package alerts
 
 import (
+	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"net/netip"
@@ -16,25 +18,22 @@ import (
 
 var (
 	errInternalDestination   = errors.New("Only admins can send to internal destinations")
-	errUnrestrictedService   = errors.New("Only admins can use notification services without HTTP client support")
+	errUnrestrictedService   = errors.New("Only admins can use this notification service") // Restrict services w/o custom connection support
+	publicNotificationDialer = &net.Dialer{
+		Timeout: 10 * time.Second,
+		// Control checks each resolved address immediately before connecting.
+		Control: func(_, address string, _ syscall.RawConn) error { return checkNotificationAddress(address) },
+	}
 	publicNotificationClient = newPublicNotificationClient()
 )
 
 func newPublicNotificationClient() *http.Client {
-	dialer := &net.Dialer{
-		Timeout: 10 * time.Second,
-		// Control receives the resolved IP, immediately before connect. Every
-		// address attempted (including DNS retries and redirects) is checked.
-		Control: func(_, address string, _ syscall.RawConn) error {
-			return checkNotificationAddress(address)
-		},
-	}
 	return &http.Client{
 		Timeout: 15 * time.Second,
 		Transport: &http.Transport{
 			// Do not use proxies: they can resolve the target themselves and
 			// bypass the destination check on our socket.
-			DialContext:         dialer.DialContext,
+			DialContext:         publicNotificationDialer.DialContext,
 			TLSHandshakeTimeout: 10 * time.Second,
 			IdleConnTimeout:     90 * time.Second,
 		},
@@ -55,8 +54,11 @@ func checkNotificationAddress(address string) error {
 
 func sendPublicNotification(rawURL, message string) error {
 	client := &notificationClient{Client: publicNotificationClient}
-	service, err := newPublicNotificationService(rawURL, client)
+	service, err := newPublicNotificationService(rawURL, types.SenderOptions{HTTPClient: client, DialContext: client.dialContext})
 	if err == nil {
+		if closer, ok := service.(io.Closer); ok {
+			defer closer.Close()
+		}
 		err = service.Send(message, &types.Params{})
 	}
 	// Some services format errors without preserving their error chain.
@@ -79,7 +81,15 @@ func (c *notificationClient) Do(req *http.Request) (*http.Response, error) {
 	return response, err
 }
 
-func newPublicNotificationService(rawURL string, client types.HTTPClient) (types.Service, error) {
+func (c *notificationClient) dialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	conn, err := publicNotificationDialer.DialContext(ctx, network, address)
+	if errors.Is(err, errInternalDestination) {
+		c.blocked.Store(true)
+	}
+	return conn, err
+}
+
+func newPublicNotificationService(rawURL string, opts types.SenderOptions) (types.Service, error) {
 	r := &router.ServiceRouter{}
 	scheme, serviceURL, err := r.ExtractServiceName(rawURL)
 	if err != nil {
@@ -89,8 +99,9 @@ func newPublicNotificationService(rawURL string, client types.HTTPClient) (types
 	if err != nil {
 		return nil, err
 	}
-	setter, ok := service.(types.HTTPClientSetter)
+	httpSetter, httpOK := service.(types.HTTPClientSetter)
-	if !ok {
+	dialSetter, dialOK := service.(types.DialContextSetter)
+	if (!httpOK || opts.HTTPClient == nil) && (!dialOK || opts.DialContext == nil) {
 		return nil, errUnrestrictedService
 	}
 	if serviceURL.Scheme != scheme {
@@ -103,14 +114,24 @@ func newPublicNotificationService(rawURL string, client types.HTTPClient) (types
 			return nil, err
 		}
 	}
-	// Shoutrrr v0.19.0 CreateSenderWithOptions injects only AFTER Initialize.
+	// Shoutrrr v0.20.0 CreateSenderWithOptions injects only AFTER Initialize.
 	// Matrix can log in during Initialize, so inject before it as well.
-	setter.SetHTTPClient(client)
+	if httpOK {
+		httpSetter.SetHTTPClient(opts.HTTPClient)
+	}
+	if dialOK {
+		dialSetter.SetDialContext(opts.DialContext)
+	}
 	if err := service.Initialize(serviceURL, nil); err != nil {
 		return nil, err
 	}
 	// Some initializers replace their HTTP client with a default client.
-	setter.SetHTTPClient(client)
+	if httpOK {
+		httpSetter.SetHTTPClient(opts.HTTPClient)
+	}
+	if dialOK {
+		dialSetter.SetDialContext(opts.DialContext)
+	}
 	return service, nil
 }
 

internal/alerts/notification_client_test.go

@@ -54,11 +54,7 @@ func TestPublicNotificationBlocksInternalRequests(t *testing.T) {
 	if hits.Load() != 0 {
 		t.Fatal("internal server received a request")
 	}
-	for _, rawURL := range []string{"smtp://user:pass@consul", "mqtt://consul/topic", "mqtts://consul/topic"} {
+
-		if err := sendPublicNotification(rawURL, "test"); !errors.Is(err, errUnrestrictedService) {
-			t.Errorf("expected unsupported transport rejection for %s, got %v", rawURL, err)
-		}
-	}
 }
 
 type notificationRoundTripper func(*http.Request) (*http.Response, error)
@@ -93,7 +89,7 @@ func TestPublicNotificationServiceClient(t *testing.T) {
 				}
 				return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
 			})}
-			service, err := newPublicNotificationService(rawURL, client)
+			service, err := newPublicNotificationService(rawURL, types.SenderOptions{HTTPClient: client})
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -158,9 +154,64 @@ func TestPublicNotificationDNS(t *testing.T) {
 	client := newPublicNotificationClient()
 	defer client.CloseIdleConnections()
 	for _, host := range []string{"rebind.example", "consul"} {
+		guarded := &notificationClient{Client: client}
+		conn, dialErr := guarded.dialContext(context.Background(), "tcp", net.JoinHostPort(host, "25"))
+		if conn != nil {
+			conn.Close()
+		}
+		if !errors.Is(dialErr, errInternalDestination) || !guarded.blocked.Load() {
+			t.Errorf("expected TCP dial-time rejection for %s, got %v", host, dialErr)
+		}
 		_, err := client.Get("http://" + host + "/")
 		if !errors.Is(err, errInternalDestination) {
 			t.Errorf("expected dial-time rejection for %s, got %v", host, err)
 		}
 	}
 }
+
+func TestPublicNotificationTCP(t *testing.T) {
+	for _, rawURL := range []string{
+		"smtp://user:pass@HOST:25/?fromAddress=sender@example.com&toAddresses=recipient@example.com",
+		"smtp://user:pass@HOST:465/?fromAddress=sender@example.com&toAddresses=recipient@example.com",
+		"mqtt://HOST:1883/topic",
+		"mqtts://HOST:8883/topic",
+	} {
+		t.Run(rawURL, func(t *testing.T) {
+			t.Parallel()
+			t.Run("internal destination", func(t *testing.T) {
+				err := sendPublicNotification(strings.ReplaceAll(rawURL, "HOST", "127.0.0.1"), "test")
+				if !errors.Is(err, errInternalDestination) {
+					t.Fatalf("expected blocked destination, got %v", err)
+				}
+			})
+			t.Run("public destination uses injected dialer", func(t *testing.T) {
+				var calls atomic.Int32
+				stopped := errors.New("test dial stopped")
+				service, err := newPublicNotificationService(strings.ReplaceAll(rawURL, "HOST", "8.8.8.8"), types.SenderOptions{
+					DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
+						calls.Add(1)
+						if network != "tcp" || !strings.HasPrefix(address, "8.8.8.8:") {
+							t.Errorf("unexpected dial: %s %s", network, address)
+						}
+						if err := checkNotificationAddress(address); err != nil {
+							t.Error(err)
+						}
+						return nil, stopped
+					},
+				})
+				if err != nil {
+					t.Fatal(err)
+				}
+				if closer, ok := service.(io.Closer); ok {
+					defer closer.Close()
+				}
+				if err := service.Send("test", &types.Params{}); err == nil {
+					t.Fatal("expected dial failure")
+				}
+				if calls.Load() == 0 {
+					t.Fatal("custom dialer was not used")
+				}
+			})
+		})
+	}
+}