multiplex multiple clients on one packet conn

client.go

@@ -47,6 +47,8 @@ type client struct {
 	logger utils.Logger
 }
 
+var _ packetHandler = &client{}
+
 var (
 	// make it possible to mock connection ID generation in the tests
 	generateConnectionID         = protocol.GenerateConnectionID
@@ -79,7 +81,15 @@ func DialAddrContext(
 	if err != nil {
 		return nil, err
 	}
-	return DialContext(ctx, udpConn, udpAddr, addr, tlsConf, config)
+	c, err := newClient(udpConn, udpAddr, config, tlsConf, addr)
+	if err != nil {
+		return nil, err
+	}
+	go c.listen()
+	if err := c.dial(ctx); err != nil {
+		return nil, err
+	}
+	return c.session, nil
 }
 
 // Dial establishes a new QUIC connection to a server using a net.PacketConn.
@@ -104,6 +114,18 @@ func DialContext(
 	tlsConf *tls.Config,
 	config *Config,
 ) (Session, error) {
+	c, err := newClient(pconn, remoteAddr, config, tlsConf, host)
+	if err != nil {
+		return nil, err
+	}
+	getClientMultiplexer().Add(pconn, c.srcConnID, c)
+	if err := c.dial(ctx); err != nil {
+		return nil, err
+	}
+	return c.session, nil
+}
+
+func newClient(pconn net.PacketConn, remoteAddr net.Addr, config *Config, tlsConf *tls.Config, host string) (*client, error) {
 	clientConfig := populateClientConfig(config)
 	version := clientConfig.Versions[0]
 	srcConnID, err := generateConnectionID()
@@ -137,8 +159,7 @@ func DialContext(
 			}
 		}
 	}
-
-	c := &client{
+	return &client{
 		conn:          &conn{pconn: pconn, currentAddr: remoteAddr},
 		srcConnID:     srcConnID,
 		destConnID:    destConnID,
@@ -148,14 +169,7 @@ func DialContext(
 		version:       version,
 		handshakeChan: make(chan struct{}),
 		logger:        utils.DefaultLogger.WithPrefix("client"),
-	}
-
-	c.logger.Infof("Starting new connection to %s (%s -> %s), source connection ID %s, destination connection ID %s, version %s", hostname, c.conn.LocalAddr(), c.conn.RemoteAddr(), c.srcConnID, c.destConnID, c.version)
-
-	if err := c.dial(ctx); err != nil {
-		return nil, err
-	}
-	return c.session, nil
+	}, nil
 }
 
 // populateClientConfig populates fields in the quic.Config with their default values, if none are set
@@ -213,6 +227,8 @@ func populateClientConfig(config *Config) *Config {
 }
 
 func (c *client) dial(ctx context.Context) error {
+	c.logger.Infof("Starting new connection to %s (%s -> %s), source connection ID %s, destination connection ID %s, version %s", c.hostname, c.conn.LocalAddr(), c.conn.RemoteAddr(), c.srcConnID, c.destConnID, c.version)
+
 	var err error
 	if c.version.UsesTLS() {
 		err = c.dialTLS(ctx)
@@ -229,7 +245,6 @@ func (c *client) dialGQUIC(ctx context.Context) error {
 	if err := c.createNewGQUICSession(); err != nil {
 		return err
 	}
-	go c.listen()
 	return c.establishSecureConnection(ctx)
 }
 
@@ -255,7 +270,6 @@ func (c *client) dialTLS(ctx context.Context) error {
 	if err := c.createNewTLSSession(extHandler.GetPeerParams(), c.version); err != nil {
 		return err
 	}
-	go c.listen()
 	if err := c.establishSecureConnection(ctx); err != nil {
 		if err != handshake.ErrCloseSessionForRetry {
 			return err
@@ -530,3 +544,12 @@ func (c *client) createNewTLSSession(
 	)
 	return err
 }
+
+func (c *client) Close(err error) error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+	if c.session == nil {
+		return nil
+	}
+	return c.session.Close(err)
+}