Support tcp and http proxy on same port (dask#203)

* Support tcp and http proxy on same port

Adds support for running both the TCP and HTTP proxies on the same port.

This relies on the SNI only being used for requests made by dask itself.
If the dask-gateway-server is running behind another proxy that makes
use of SNIs, then this method won't work and the proxies will have to
use separate ports (which is still supported).

Running all on the same point is nice from an administrative end (only
need to expose one port, not two), as well as a configuration end (users
only need to remember the main address, not the proxy address).

* Filter valid SNIs

Previously we'd accept any SNI, we now filter to only proxy SNIs that
start with `daskgateway-`. This helps prevent accidentally mixing HTTPS
and dask comms, since browsers will send an SNI on connect if connecting
via HTTPS to a domain with a non-ip hostname.

Author: jcrist

dask-gateway-server/dask-gateway-proxy/main.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"bufio"
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
@@ -89,20 +88,18 @@ func (r *Route) Validate(full bool) error {
 }
 
 type Proxy struct {
-	tlsTimeout time.Duration
 	logger     *Logger
 	sniRoutes  map[string]string
 	router     *Router
-	proxy      *httputil.ReverseProxy
 	routesLock sync.RWMutex
+	proxy      *httputil.ReverseProxy
 }
 
-func NewProxy(tlsTimeout time.Duration, logLevel LogLevel) *Proxy {
+func NewProxy(logLevel LogLevel) *Proxy {
 	out := Proxy{
-		tlsTimeout: tlsTimeout,
-		logger:     NewLogger("Proxy", logLevel),
-		sniRoutes:  make(map[string]string),
-		router:     NewRouter(),
+		logger:    NewLogger("Proxy", logLevel),
+		sniRoutes: make(map[string]string),
+		router:    NewRouter(),
 	}
 	out.proxy = &httputil.ReverseProxy{
 		Director:      out.director,
@@ -122,12 +119,38 @@ func awaitShutdown() {
 	}
 }
 
-func (p *Proxy) run(httpAddress, tlsAddress, apiURL, apiToken, tlsCert, tlsKey string, isChildProcess bool) {
+func (p *Proxy) run(address, tcpAddress, apiURL, apiToken, tlsCert, tlsKey string, isChildProcess bool) {
 	if isChildProcess {
 		go awaitShutdown()
 	}
-	go p.runHTTP(httpAddress, tlsCert, tlsKey)
-	go p.runTLS(tlsAddress)
+
+	if tcpAddress == "" {
+		tcpAddress = address
+	}
+
+	tlsListener, err := net.Listen("tcp", tcpAddress)
+	if err != nil {
+		p.logger.Errorf("%s", err)
+		os.Exit(1)
+	}
+	p.logger.Infof("TCP Proxy serving at %s", tcpAddress)
+
+	var httpListener net.Listener
+	var forwarder *Forwarder = nil
+	if address == tcpAddress {
+		forwarder = newForwarder(tlsListener)
+		httpListener = forwarder
+	} else {
+		httpListener, err = net.Listen("tcp", address)
+		if err != nil {
+			p.logger.Errorf("%s", err)
+			os.Exit(1)
+		}
+	}
+	p.logger.Infof("HTTP Proxy serving at %s", address)
+
+	go p.runHTTP(httpListener, tlsCert, tlsKey)
+	go p.runTCP(tlsListener, forwarder)
 	p.watchRoutes(apiURL, apiToken, 15*time.Second)
 }
 
@@ -408,110 +431,97 @@ func (p *Proxy) director(req *http.Request) {
 	}
 }
 
-func (p *Proxy) runHTTP(address, tlsCert, tlsKey string) {
+func (p *Proxy) runHTTP(ln net.Listener, tlsCert, tlsKey string) {
 	server := &http.Server{
-		Addr:    address,
 		Handler: p,
 	}
 	var err error
 	if tlsCert == "" {
-		p.logger.Infof("HTTP Proxy serving at http://%s", address)
-		err = server.ListenAndServe()
+		err = server.Serve(ln)
 	} else {
-		p.logger.Infof("HTTP Proxy serving at https://%s", address)
-		err = server.ListenAndServeTLS(tlsCert, tlsKey)
+		err = server.ServeTLS(ln, tlsCert, tlsKey)
 	}
 	if err != nil && err != http.ErrServerClosed {
 		p.logger.Errorf("%s", err)
 		os.Exit(1)
 	}
 }
 
-// TLS proxy implementation
+// TCP proxy implementation
 
-type tlsConn struct {
-	inConn   *net.TCPConn
-	outConn  *net.TCPConn
-	sni      string
-	outAddr  string
-	tlsMinor int
+type Forwarder struct {
+	net.Listener
+	conns chan net.Conn
 }
 
-type tlsAlert int8
-
-const (
-	internalError    tlsAlert = 80
-	unrecognizedName          = 112
-)
-
-func (p *Proxy) sendAlert(c *tlsConn, alert tlsAlert, format string, args ...interface{}) {
-	p.logger.Debugf(format, args...)
-
-	alertMsg := []byte{21, 3, byte(c.tlsMinor), 0, 2, 2, byte(alert)}
-
-	if err := c.inConn.SetWriteDeadline(time.Now().Add(p.tlsTimeout)); err != nil {
-		p.logger.Debugf("Error while setting write deadline during abort: %s", err)
-		return
+func newForwarder(ln net.Listener) *Forwarder {
+	return &Forwarder{
+		Listener: ln,
+		conns:    make(chan net.Conn),
 	}
+}
 
-	if _, err := c.inConn.Write(alertMsg); err != nil {
-		p.logger.Debugf("Error while sending alert: %s", err)
-	}
+func (h *Forwarder) Forward(conn net.Conn) {
+	h.conns <- conn
 }
 
-func (p *Proxy) handleConnection(c *tlsConn) {
-	defer c.inConn.Close()
+func (h *Forwarder) Accept() (net.Conn, error) {
+	conn := <-h.conns
+	return conn, nil
+}
 
+func (p *Proxy) handleConnection(inConn *net.TCPConn, forwarder *Forwarder) {
 	var err error
 
-	if err = c.inConn.SetReadDeadline(time.Now().Add(p.tlsTimeout)); err != nil {
-		p.sendAlert(c, internalError, "Setting read deadline for handshake: %s", err)
-		return
-	}
-
-	var handshakeBuf bytes.Buffer
-	c.sni, c.tlsMinor, err = readVerAndSNI(io.TeeReader(c.inConn, &handshakeBuf))
+	sni, isTLS, pInConn, err := readSNI(inConn)
 	if err != nil {
-		p.sendAlert(c, internalError, "Extracting SNI: %s", err)
+		p.logger.Debugf("Error extracting SNI: %s", err)
+		inConn.Close()
 		return
 	}
 
-	if err = c.inConn.SetReadDeadline(time.Time{}); err != nil {
-		p.sendAlert(c, internalError, "Clearing read deadline for handshake: %s", err)
-		return
+	const sniPrefix = "daskgateway-"
+	if !isTLS || !strings.HasPrefix(sni, sniPrefix) {
+		if forwarder != nil {
+			forwarder.Forward(pInConn)
+			return
+		} else {
+			p.logger.Debug("Invalid connection attempt, closing")
+			inConn.Close()
+			return
+		}
 	}
 
+	sni = sni[len(sniPrefix):]
+
+	defer inConn.Close()
+
 	p.routesLock.RLock()
-	c.outAddr = p.sniRoutes[c.sni]
+	outAddr := p.sniRoutes[sni]
 	p.routesLock.RUnlock()
-	if c.outAddr == "" {
-		p.sendAlert(c, unrecognizedName, "SNI %q not found", c.sni)
+
+	if outAddr == "" {
+		p.logger.Infof("SNI %q not found", sni)
 		return
 	}
 
-	p.logger.Infof("SNI %q -> %q", c.sni, c.outAddr)
+	p.logger.Infof("SNI %q -> %q", sni, outAddr)
 
-	outConn, err := net.DialTimeout("tcp", c.outAddr, 10*time.Second)
+	outConn, err := net.DialTimeout("tcp", outAddr, 10*time.Second)
 	if err != nil {
-		p.sendAlert(c, internalError, "Failed to connect to destination %q: %s", c.outAddr, err)
-		return
-	}
-	c.outConn = outConn.(*net.TCPConn)
-	defer c.outConn.Close()
-
-	if _, err = io.Copy(c.outConn, &handshakeBuf); err != nil {
-		p.sendAlert(c, internalError, "Failed to replay handshake to %q: %s", c.outAddr, err)
+		p.logger.Debugf("Failed to connect to destination %q: %s", outAddr, err)
 		return
 	}
+	defer outConn.Close()
 
 	var wg sync.WaitGroup
 	wg.Add(2)
-	go p.proxyConnections(&wg, c.inConn, c.outConn)
-	go p.proxyConnections(&wg, c.outConn, c.inConn)
+	go p.proxyConnections(&wg, pInConn, outConn.(*net.TCPConn))
+	go p.proxyConnections(&wg, outConn.(*net.TCPConn), pInConn)
 	wg.Wait()
 }
 
-func (p *Proxy) proxyConnections(wg *sync.WaitGroup, in, out *net.TCPConn) {
+func (p *Proxy) proxyConnections(wg *sync.WaitGroup, in, out tcpConn) {
 	defer wg.Done()
 	if _, err := io.Copy(in, out); err != nil {
 		p.logger.Debugf("Error proxying %q -> %q: %s", in.RemoteAddr(), out.RemoteAddr(), err)
@@ -520,47 +530,38 @@ func (p *Proxy) proxyConnections(wg *sync.WaitGroup, in, out *net.TCPConn) {
 	out.CloseWrite()
 }
 
-func (p *Proxy) runTLS(tlsAddress string) {
-	p.logger.Infof("TLS Proxy serving at %s", tlsAddress)
-	l, err := net.Listen("tcp", tlsAddress)
-	if err != nil {
-		p.logger.Errorf("%s", err)
-		os.Exit(1)
-	}
+func (p *Proxy) runTCP(ln net.Listener, forwarder *Forwarder) {
 	for {
-		c, err := l.Accept()
+		c, err := ln.Accept()
 		if err != nil {
 			p.logger.Debugf("Failed to accept new connection: %s", err)
 		}
 
-		conn := &tlsConn{inConn: c.(*net.TCPConn)}
-		go p.handleConnection(conn)
+		go p.handleConnection(c.(*net.TCPConn), forwarder)
 	}
 }
 
 func main() {
 	var (
-		httpAddress    string
-		tlsAddress     string
+		address        string
+		tcpAddress     string
 		apiURL         string
 		logLevelString string
 		isChildProcess bool
 		tlsCert        string
 		tlsKey         string
-		tlsTimeout     time.Duration
 	)
 
 	command := flag.NewFlagSet("dask-gateway-proxy", flag.ExitOnError)
-	command.StringVar(&httpAddress, "address", ":8000", "HTTP proxy listening address")
-	command.StringVar(&tlsAddress, "tls-address", ":8786", "TLS proxy listening address")
+	command.StringVar(&address, "address", ":8000", "HTTP proxy listening address")
+	command.StringVar(&tcpAddress, "tcp-address", "", "TCP proxy listening address. If empty, `address` will be used.")
 	command.StringVar(&apiURL, "api-url", "", "The URL the proxy should watch for updating the routing table")
 	command.StringVar(&logLevelString, "log-level", "INFO",
 		"The log level. One of {DEBUG, INFO, WARN, ERROR}")
 	command.BoolVar(&isChildProcess, "is-child-process", false,
 		"If set, will exit when stdin EOFs. Useful when running as a child process.")
 	command.StringVar(&tlsCert, "tls-cert", "", "TLS cert to use, if any.")
 	command.StringVar(&tlsKey, "tls-key", "", "TLS key to use, if any.")
-	command.DurationVar(&tlsTimeout, "tls-timeout", 3*time.Second, "Timeout for TLS Handshake initialization")
 
 	command.Parse(os.Args[1:])
 
@@ -582,7 +583,7 @@ func main() {
 		panic("Invalid -api-url: " + err.Error())
 	}
 
-	proxy := NewProxy(tlsTimeout, logLevel)
+	proxy := NewProxy(logLevel)
 
-	proxy.run(httpAddress, tlsAddress, apiURL, token, tlsCert, tlsKey, isChildProcess)
+	proxy.run(address, tcpAddress, apiURL, token, tlsCert, tlsKey, isChildProcess)
 }

dask-gateway-server/dask-gateway-proxy/sni.go

@@ -1,154 +1,81 @@
 package main
 
 import (
-	"encoding/binary"
-	"errors"
-	"fmt"
+	"bufio"
+	"bytes"
+	"crypto/tls"
 	"io"
+	"net"
 )
 
-const extensionID = 0
-const hostnameID = 0
-
-// Parses a Vector object (length prefixed bytes) per the TLS spec
-func parseVector(buf []byte, lenBytes int) ([]byte, []byte, error) {
-	if len(buf) < lenBytes {
-		return nil, nil, errors.New("Not enough space in packet for vector")
-	}
-	var l int
-	for _, b := range buf[:lenBytes] {
-		l = (l << 8) + int(b)
-	}
-	if len(buf) < l+lenBytes {
-		return nil, nil, errors.New("Not enough space in packet for vector")
-	}
-	return buf[lenBytes : l+lenBytes], buf[l+lenBytes:], nil
+type tcpConn interface {
+	net.Conn
+	CloseWrite() error
+	CloseRead() error
 }
 
-func readSNI(buf []byte) (string, error) {
-	if len(buf) == 0 {
-		return "", errors.New("Zero length handshake record")
-	}
-	if buf[0] != 1 {
-		return "", fmt.Errorf("Non-ClientHello handshake record type %d", buf[0])
-	}
+type peekedTCPConn struct {
+	peeked []byte
+	*net.TCPConn
+}
 
-	buf, _, err := parseVector(buf[1:], 3)
-	if err != nil {
-		return "", fmt.Errorf("Reading ClientHello: %s", err)
+func (c *peekedTCPConn) Read(p []byte) (n int, err error) {
+	if len(c.peeked) > 0 {
+		n = copy(p, c.peeked)
+		c.peeked = c.peeked[n:]
+		if len(c.peeked) == 0 {
+			c.peeked = nil
+		}
+		return n, nil
 	}
+	return c.TCPConn.Read(p)
+}
 
-	if len(buf) < 34 {
-		return "", errors.New("ClientHello packet too short")
-	}
+func wrapPeeked(inConn *net.TCPConn, br *bufio.Reader) tcpConn {
+	peeked, _ := br.Peek(br.Buffered())
+	return &peekedTCPConn{TCPConn: inConn, peeked: peeked}
+}
 
-	if buf[0] != 3 || buf[1] < 1 || buf[1] > 3 {
-		return "", fmt.Errorf("ClientHello has unsupported version %d.%d", buf[0], buf[1])
-	}
+type readonly struct {
+	r io.Reader
+	net.Conn
+}
 
-	// Skip version and random struct
-	buf = buf[34:]
+func (c readonly) Read(p []byte) (int, error) { return c.r.Read(p) }
+func (readonly) Write(p []byte) (int, error)  { return 0, io.EOF }
 
-	vec, buf, err := parseVector(buf, 1)
+func readSNI(inConn *net.TCPConn) (string, bool, tcpConn, error) {
+	br := bufio.NewReader(inConn)
+	hdr, err := br.Peek(1)
 	if err != nil {
-		return "", fmt.Errorf("Reading ClientHello SessionID: %s", err)
-	}
-	if len(vec) > 32 {
-		return "", fmt.Errorf("ClientHello SessionID too long (%db)", len(vec))
+		return "", false, nil, err
 	}
 
-	vec, buf, err = parseVector(buf, 2)
-	if err != nil {
-		return "", fmt.Errorf("Reading ClientHello CipherSuites: %s", err)
-	}
-	if len(vec) < 2 || len(vec)%2 != 0 {
-		return "", fmt.Errorf("ClientHello CipherSuites invalid length %d", len(vec))
+	if hdr[0] != 0x16 {
+		// Not a TLS handshake
+		return "", false, wrapPeeked(inConn, br), nil
 	}
 
-	vec, buf, err = parseVector(buf, 1)
+	const headerLen = 5
+	hdr, err = br.Peek(headerLen)
 	if err != nil {
-		return "", fmt.Errorf("Reading ClientHello CompressionMethods: %s", err)
-	}
-	if len(vec) < 1 {
-		return "", fmt.Errorf("ClientHello CompressionMethods invalid length %d", len(vec))
-	}
-
-	if len(buf) != 0 {
-		// Check vector is proper length for remaining msg
-		vec, buf, err = parseVector(buf, 2)
-		if err != nil {
-			return "", fmt.Errorf("Reading ClientHello extensions: %s", err)
-		}
-		if len(buf) != 0 {
-			return "", fmt.Errorf("%d bytes of trailing garbage in ClientHello", len(buf))
-		}
-		buf = vec
-
-		for len(buf) >= 4 {
-			typ := binary.BigEndian.Uint16(buf[:2])
-			vec, buf, err = parseVector(buf[2:], 2)
-			if err != nil {
-				return "", fmt.Errorf("Reading ClientHello extension %d: %s", typ, err)
-			}
-			if typ == extensionID {
-				// We found an SNI extension, attempt to extract the hostname
-				buf, _, err := parseVector(vec, 2)
-				if err != nil {
-					return "", err
-				}
-
-				for len(buf) >= 3 {
-					typ := buf[0]
-					vec, buf, err = parseVector(buf[1:], 2)
-					if err != nil {
-						return "", errors.New("Truncated SNI extension")
-					}
-
-					// This vec is a hostname, return
-					if typ == hostnameID {
-						return string(vec), nil
-					}
-				}
-
-				if len(buf) != 0 {
-					return "", errors.New("Trailing garbage at end of SNI extension")
-				}
-
-				return "", nil
-			}
-		}
+		return "", false, wrapPeeked(inConn, br), nil
 	}
-	return "", errors.New("No SNI found")
-}
 
-func readVerAndSNI(r io.Reader) (string, int, error) {
-	var header struct {
-		Type         uint8
-		Major, Minor uint8
-		Length       uint16
-	}
-	if err := binary.Read(r, binary.BigEndian, &header); err != nil {
-		return "", 0, fmt.Errorf("Error reading TLS record header: %s", err)
-	}
-
-	if header.Type != 22 {
-		return "", 0, fmt.Errorf("TLS record is not a handshake")
-	}
-
-	if header.Major != 3 || header.Minor < 1 || header.Minor > 3 {
-		return "", 0, fmt.Errorf("TLS record has unsupported version %d.%d",
-			header.Major, header.Minor)
-	}
-
-	if header.Length > 16384 {
-		return "", 0, errors.New("TLS record is malformed")
+	recLen := int(hdr[3])<<8 | int(hdr[4])
+	helloBytes, err := br.Peek(headerLen + recLen)
+	if err != nil {
+		return "", true, wrapPeeked(inConn, br), nil
 	}
 
-	buf := make([]byte, header.Length)
-	if _, err := io.ReadFull(r, buf); err != nil {
-		return "", 0, err
-	}
+	sni := ""
+	server := tls.Server(readonly{r: bytes.NewReader(helloBytes)}, &tls.Config{
+		GetConfigForClient: func(hello *tls.ClientHelloInfo) (*tls.Config, error) {
+			sni = hello.ServerName
+			return nil, nil
+		},
+	})
+	server.Handshake()
 
-	sni, err := readSNI(buf)
-	return sni, int(header.Minor), err
+	return sni, true, wrapPeeked(inConn, br), nil
 }

dask-gateway-server/dask_gateway_server/proxy/core.py

@@ -53,10 +53,9 @@ class Proxy(LoggingConfigurable):
         config=True,
     )
 
-    scheduler_address = Unicode(
-        ":8786",
+    tcp_address = Unicode(
         help="""
-        The address the scheduler proxy should *listen* at.
+        The address the TCP (scheduler) proxy should *listen* at.
 
         Should be of the form ``{hostname}:{port}``
 
@@ -65,11 +64,17 @@ class Proxy(LoggingConfigurable):
         - ``hostname`` sets the hostname to *listen* at. Set to ``""`` or
           ``"0.0.0.0"`` to listen on all interfaces.
         - ``port`` sets the port to *listen* at.
+
+        If not specified, will default to `address`.
         """,
         config=True,
     )
 
-    @validate("address", "scheduler_address")
+    @default("tcp_address")
+    def _default_tcp_address(self):
+        return self.address
+
+    @validate("address", "tcp_address")
     def _validate_addresses(self, proposal):
         return normalize_address(proposal.value)
 
@@ -192,8 +197,8 @@ def get_start_command(self, is_child_process=True):
             _PROXY_EXE,
             "-address",
             self.address,
-            "-tls-address",
-            self.scheduler_address,
+            "-tcp-address",
+            self.tcp_address,
             "-api-url",
             self.gateway_url + "/api/routes",
             "-log-level",
@@ -232,9 +237,7 @@ def start_proxy_process(self):
             "https" if self.tls_cert else "http",
             self.address,
         )
-        self.log.info(
-            "- Scheduler routes listening at tls://%s", self.scheduler_address
-        )
+        self.log.info("- Scheduler routes listening at gateway://%s", self.tcp_address)
 
     async def monitor_proxy_process(self):
         backoff = default_backoff = 0.5

dask-gateway/dask_gateway/client.py

@@ -244,9 +244,10 @@ class Gateway(object):
     address : str, optional
         The address to the gateway server.
     proxy_address : str, int, optional
-        The address of the scheduler proxy server. If an int, it's used as the
-        port, with the host/ip taken from ``address``. Provide a full address
-        if a different host/ip should be used.
+        The address of the scheduler proxy server. Defaults to `address` if not
+        provided. If an int, it's used as the port, with the host/ip taken from
+        ``address``. Provide a full address if a different host/ip should be
+        used.
     auth : GatewayAuth, optional
         The authentication method to use.
     asynchronous : bool, optional
@@ -277,9 +278,7 @@ def __init__(
         if proxy_address is None:
             proxy_address = format_template(dask.config.get("gateway.proxy-address"))
         if proxy_address is None:
-            raise ValueError(
-                "No dask-gateway proxy address provided or found in configuration"
-            )
+            proxy_address = address
         if isinstance(proxy_address, int):
             parsed = urlparse(address)
             proxy_netloc = "%s:%d" % (parsed.hostname, proxy_address)

dask-gateway/dask_gateway/comm.py

@@ -33,7 +33,8 @@ class GatewayConnector(Connector):
     client = TCPClient(resolver=_resolver)
 
     async def connect(self, address, deserialize=True, **connection_args):
-        ip, port, sni = parse_gateway_address(address)
+        ip, port, path = parse_gateway_address(address)
+        sni = "daskgateway-" + path
         ctx = connection_args.get("ssl_context")
         if not isinstance(ctx, ssl.SSLContext):
             raise TypeError(

dask-gateway/dask_gateway/gateway.yaml

@@ -9,9 +9,10 @@ gateway:
                         # If `None` (default), `gateway.address` will be used.
                         # May be a template string.
 
-  proxy-address: 8786   # The full address or port to the dask-gateway
+  proxy-address: null   # The full address or port to the dask-gateway
                         # scheduler proxy. If a port, the host/ip is taken from
-                        # ``address``. May also be a template string.
+                        # ``address``. If null, defaults to `address`.
+                        # May also be a template string.
 
   auth:
     type: basic         # The authentication type to use. Options are basic,

tests/test_cli.py

@@ -46,7 +46,7 @@ def test_proxy_cli(tmpdir, monkeypatch):
     text = (
         "c.DaskGateway.address = '127.0.0.1:8888'\n"
         "c.Proxy.address = '127.0.0.1:8866'\n"
-        "c.Proxy.scheduler_address = '127.0.0.1:8867'\n"
+        "c.Proxy.tcp_address = '127.0.0.1:8867'\n"
         "c.Proxy.log_level = 'debug'\n"
         "c.Proxy.api_token = 'abcde'"
     )
@@ -71,7 +71,7 @@ def mock_execle(*args):
         "dask-gateway-proxy",
         "-address",
         "127.0.0.1:8866",
-        "-tls-address",
+        "-tcp-address",
         "127.0.0.1:8867",
         "-api-url",
         "http://127.0.0.1:8888/api/routes",

tests/test_client.py

@@ -133,9 +133,8 @@ def test_client_init():
     config["gateway"]["proxy-address"] = None
 
     with dask.config.set(config):
-        # No proxy-address provided
-        with pytest.raises(ValueError):
-            Gateway()
+        g = Gateway()
+        assert g.proxy_address == "gateway://127.0.0.1:8888"
 
 
 def test_gateway_addresses_template_environment_vars(monkeypatch):

tests/test_db_backend.py

@@ -748,7 +748,6 @@ async def test_gateway_resume_clusters_after_shutdown(tmpdir):
     config.LocalTestingBackend.check_timeouts_period = 0.5
     config.DaskGateway.address = "127.0.0.1:%d" % random_port()
     config.Proxy.address = "127.0.0.1:%d" % random_port()
-    config.Proxy.scheduler_address = "127.0.0.1:%d" % random_port()
 
     async with temp_gateway(config=config) as g:
         async with g.gateway_client() as gateway:

tests/test_proxies.py

@@ -21,7 +21,6 @@ def __init__(self, **kwargs):
         self.proxy = Proxy(
             address="127.0.0.1:0",
             prefix="/foobar",
-            scheduler_address="127.0.0.1:0",
             gateway_address=f"127.0.0.1:{self._port}",
             log_level="debug",
             proxy_status_period=0.5,
@@ -43,9 +42,10 @@ async def __aexit__(self, *args):
         await self.proxy.cleanup()
 
 
-@pytest.fixture
-async def proxy():
-    async with temp_proxy() as proxy:
+@pytest.fixture(params=["separate", "shared"])
+async def proxy(request):
+    kwargs = {"tcp_address": "127.0.0.1:0"} if request.param == "separate" else {}
+    async with temp_proxy(**kwargs) as proxy:
         yield proxy
 
 
@@ -140,8 +140,8 @@ async def test_502():
         await with_retries(test_502, 5)
 
 
-@pytest.fixture
-async def ca_and_tls_proxy(tmpdir_factory):
+@pytest.fixture(params=["separate", "shared"])
+async def ca_and_tls_proxy(request, tmpdir_factory):
     trustme = pytest.importorskip("trustme")
     ca = trustme.CA()
     cert = ca.issue_cert("127.0.0.1")
@@ -153,7 +153,8 @@ async def ca_and_tls_proxy(tmpdir_factory):
     cert.private_key_pem.write_to_path(tls_key)
     cert.cert_chain_pems[0].write_to_path(tls_cert)
 
-    async with temp_proxy(tls_key=tls_key, tls_cert=tls_cert) as proxy:
+    kwargs = {"tcp_address": "127.0.0.1:0"} if request.param == "separate" else {}
+    async with temp_proxy(tls_key=tls_key, tls_cert=tls_cert, **kwargs) as proxy:
         yield ca, proxy
 
 
@@ -196,7 +197,7 @@ async def test_fails():
 async def test_scheduler_proxy(proxy, cluster_and_security):
     cluster, security = cluster_and_security
 
-    proxied_addr = f"gateway://{proxy.scheduler_address}/temp"
+    proxied_addr = f"gateway://{proxy.tcp_address}/temp"
 
     # Add a route
     await proxy.add_route(kind="SNI", sni="temp", target=cluster.scheduler_address)

tests/utils_test.py

@@ -46,7 +46,6 @@ def __init__(self, **kwargs):
 
         c.DaskGateway.address = "127.0.0.1:0"
         c.Proxy.address = "127.0.0.1:0"
-        c.Proxy.scheduler_address = "127.0.0.1:0"
         c.DaskGateway.authenticator_class = (
             "dask_gateway_server.auth.SimpleAuthenticator"
         )
@@ -63,7 +62,7 @@ async def __aenter__(self):
         await self.gateway.setup()
         await self.gateway.backend.proxy._proxy_contacted
         self.address = f"http://{self.gateway.backend.proxy.address}"
-        self.proxy_address = f"tls://{self.gateway.backend.proxy.scheduler_address}"
+        self.proxy_address = f"gateway://{self.gateway.backend.proxy.tcp_address}"
         return self
 
     async def __aexit__(self, *args):