From 36add0e7a0d1d1a8390e5409d826f142ec680d08 Mon Sep 17 00:00:00 2001
From: Yingrong Zhao <yingrong.zhao@gmail.com>
Date: Mon, 5 Oct 2020 13:10:57 -0400
Subject: [PATCH] client: Add DialEarlyContext and DialAddrEarlyContext API

---
 client.go                             | 28 +++++++++++++++++++++++++--
 integrationtests/self/timeout_test.go | 19 ++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/client.go b/client.go
index 3663cf88..53d85dc0 100644
--- a/client.go
+++ b/client.go
@@ -68,7 +68,18 @@ func DialAddrEarly(
 	tlsConf *tls.Config,
 	config *Config,
 ) (EarlySession, error) {
-	sess, err := dialAddrContext(context.Background(), addr, tlsConf, config, true)
+	return DialAddrEarlyContext(context.Background(), addr, tlsConf, config)
+}
+
+// DialAddrEarlyContext establishes a new 0-RTT QUIC connection to a server using provided context.
+// See DialAddrEarly for details
+func DialAddrEarlyContext(
+	ctx context.Context,
+	addr string,
+	tlsConf *tls.Config,
+	config *Config,
+) (EarlySession, error) {
+	sess, err := dialAddrContext(ctx, addr, tlsConf, config, true)
 	if err != nil {
 		return nil, err
 	}
@@ -134,7 +145,20 @@ func DialEarly(
 	tlsConf *tls.Config,
 	config *Config,
 ) (EarlySession, error) {
-	return dialContext(context.Background(), pconn, remoteAddr, host, tlsConf, config, true, false)
+	return DialEarlyContext(context.Background(), pconn, remoteAddr, host, tlsConf, config)
+}
+
+// DialEarlyContext establishes a new 0-RTT QUIC connection to a server using a net.PacketConn using the provided context.
+// See DialEarly for details.
+func DialEarlyContext(
+	ctx context.Context,
+	pconn net.PacketConn,
+	remoteAddr net.Addr,
+	host string,
+	tlsConf *tls.Config,
+	config *Config,
+) (EarlySession, error) {
+	return dialContext(ctx, pconn, remoteAddr, host, tlsConf, config, true, false)
 }
 
 // DialContext establishes a new QUIC connection to a server using a net.PacketConn using the provided context.
diff --git a/integrationtests/self/timeout_test.go b/integrationtests/self/timeout_test.go
index 6143d993..898dd415 100644
--- a/integrationtests/self/timeout_test.go
+++ b/integrationtests/self/timeout_test.go
@@ -94,6 +94,25 @@ var _ = Describe("Timeout tests", func() {
 		Expect(err).To(MatchError(context.DeadlineExceeded))
 	})
 
+	It("returns the context error when the context expires with 0RTT enabled", func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+		defer cancel()
+		errChan := make(chan error)
+		go func() {
+			_, err := quic.DialAddrEarlyContext(
+				ctx,
+				"localhost:12345",
+				getTLSClientConfig(),
+				getQuicConfig(nil),
+			)
+			errChan <- err
+		}()
+		var err error
+		Eventually(errChan).Should(Receive(&err))
+		// This is not a net.Error timeout error
+		Expect(err).To(MatchError(context.DeadlineExceeded))
+	})
+
 	It("returns net.Error timeout errors when an idle timeout occurs", func() {
 		const idleTimeout = 100 * time.Millisecond