context aware connectSourceConnection

internal/restrictednet/client.go

@@ -2,6 +2,7 @@ package restrictednet
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net/http"
 
@@ -48,12 +49,15 @@ func (c *Client) OpenHTTPSByDomain(ctx context.Context, domain string) (
 		return nil, nil, fmt.Errorf("no IP address found for name %q", domain)
 	}
 
-	selectedIP := resolvedIPs[0]
-
-	httpClient, cleanup, err = c.OpenHTTPS(ctx, domain, selectedIP)
-	if err != nil {
-		return nil, nil, fmt.Errorf("opening HTTPS: %w", err)
+	errs := make([]error, 0, len(resolvedIPs))
+	for _, ip := range resolvedIPs {
+		httpClient, cleanup, err := c.OpenHTTPS(ctx, domain, ip)
+		if err != nil {
+			errs = append(errs, fmt.Errorf("for %s: %w", ip, err))
+			continue
+		}
+		return httpClient, cleanup, nil
 	}
 
-	return httpClient, cleanup, nil
+	return nil, nil, fmt.Errorf("opening HTTPS to %s: %w", domain, errors.Join(errs...))
 }

internal/restrictednet/https.go

@@ -34,7 +34,7 @@ func (c *Client) OpenHTTPS(ctx context.Context, destinationTLSName string, desti
 		return nil, nil, fmt.Errorf("allowing output traffic through firewall: %w", err)
 	}
 
-	connection, err := connectSourceConnection(fd, destinationAddrPort)
+	connection, err := connectSourceConnection(ctx, fd, destinationAddrPort)
 	if err != nil {
 		const remove = true
 		_ = c.firewall.AcceptOutputFromIPPortToIPPort(ctx, "tcp", c.outboundInterface,
@@ -47,7 +47,7 @@ func (c *Client) OpenHTTPS(ctx context.Context, destinationTLSName string, desti
 		var errs []error
 		httpClient.CloseIdleConnections()
 		const remove = true
-		err := c.firewall.AcceptOutputFromIPPortToIPPort(ctx, "tcp", c.outboundInterface,
+		err := c.firewall.AcceptOutputFromIPPortToIPPort(context.Background(), "tcp", c.outboundInterface,
 			sourceAddrPort, destinationAddrPort, remove)
 		if err != nil {
 			errs = append(errs, fmt.Errorf("removing output traffic rule: %w", err))
@@ -76,7 +76,8 @@ func newHTTPSClient(destinationTLSName string, connection net.Conn) *http.Client
 		ServerName: destinationTLSName,
 	}
 
-	expectedAddress := net.JoinHostPort(destinationTLSName, "443")
+	_, destinationPort, _ := net.SplitHostPort(connection.RemoteAddr().String())
+	expectedAddress := net.JoinHostPort(destinationTLSName, destinationPort)
 	httpTransport.DialContext = func(_ context.Context, network, address string) (net.Conn, error) {
 		switch network {
 		case "tcp", "tcp4", "tcp6":
@@ -128,10 +129,27 @@ func bindSourceConnection(destinationIP netip.Addr) (fd int, sourceAddr netip.Ad
 	return fd, sourceAddr, nil
 }
 
-func connectSourceConnection(fd int, destinationAddrPort netip.AddrPort) (connection net.Conn, err error) {
-	err = connectFD(fd, destinationAddrPort)
-	if err != nil {
+func connectSourceConnection(ctx context.Context, fd int, destinationAddrPort netip.AddrPort) (
+	connection net.Conn, err error,
+) {
+	errCh := make(chan error)
+	go func() {
+		errCh <- connectFD(fd, destinationAddrPort)
+	}()
+
+	select {
+	case err = <-errCh:
+		if err != nil {
+			closeFD(fd)
+			return nil, fmt.Errorf("connecting socket: %w", err)
+		}
+	case <-ctx.Done():
+		err = ctx.Err()
 		closeFD(fd)
+		connectErr := <-errCh
+		if connectErr != nil {
+			err = fmt.Errorf("%w (%w)", connectErr, err)
+		}
 		return nil, fmt.Errorf("connecting socket: %w", err)
 	}
 

internal/restrictednet/https_test.go

@@ -65,7 +65,7 @@ func Test_Client_OpenHTTPS(t *testing.T) {
 		return nil
 	})
 	firewall.EXPECT().AcceptOutputFromIPPortToIPPort(
-		ctx, "tcp", "eth0", sourceMatcher, destination, true,
+		context.Background(), "tcp", "eth0", sourceMatcher, destination, true,
 	)
 
 	const ipv6Supported = false

internal/restrictednet/resolve.go

@@ -142,8 +142,8 @@ func (c *Client) doHQuery(ctx context.Context, queryWire []byte,
 	}
 
 	if response.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("response status code is %s, data: %s",
-			response.Status, responseData)
+		return nil, fmt.Errorf("response status code is %s (data length %d)",
+			response.Status, len(responseData))
 	}
 
 	responseMessage = new(dns.Msg)