From 57461e01b5345e0dbd65512608465e402575cad0 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Tue, 22 Mar 2022 09:34:11 +0100
Subject: [PATCH] add a http3.Hijacker that allows stream creation on a QUIC
 session from a http.Response.Body

---
 http3/body.go      | 40 +++++++++++++++++++++++++++++++++++-----
 http3/body_test.go |  4 ++--
 http3/client.go    |  2 +-
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/http3/body.go b/http3/body.go
index 851eaa1f..b6363dc6 100644
--- a/http3/body.go
+++ b/http3/body.go
@@ -1,12 +1,28 @@
 package http3
 
 import (
+	"context"
 	"fmt"
 	"io"
 
 	"github.com/lucas-clemente/quic-go"
 )
 
+type StreamCreator interface {
+	OpenStream() (quic.Stream, error)
+	OpenStreamSync(context.Context) (quic.Stream, error)
+	OpenUniStream() (quic.SendStream, error)
+	OpenUniStreamSync(context.Context) (quic.SendStream, error)
+}
+
+var _ StreamCreator = quic.Connection(nil)
+
+// A Hijacker allows hijacking of the stream creating part of a quic.Session from a http.Response.Body.
+// It is used by WebTransport to create WebTransport streams after a session has been established.
+type Hijacker interface {
+	StreamCreator() StreamCreator
+}
+
 // The body of a http.Request or http.Response.
 type body struct {
 	str quic.Stream
@@ -24,6 +40,13 @@ type body struct {
 
 var _ io.ReadCloser = &body{}
 
+type hijackableBody struct {
+	body
+	conn quic.Connection // only needed to implement Hijacker
+}
+
+var _ Hijacker = &hijackableBody{}
+
 func newRequestBody(str quic.Stream, onFrameError func()) *body {
 	return &body{
 		str:          str,
@@ -31,14 +54,21 @@ func newRequestBody(str quic.Stream, onFrameError func()) *body {
 	}
 }
 
-func newResponseBody(str quic.Stream, done chan<- struct{}, onFrameError func()) *body {
-	return &body{
-		str:          str,
-		onFrameError: onFrameError,
-		reqDone:      done,
+func newResponseBody(str quic.Stream, conn quic.Connection, done chan<- struct{}, onFrameError func()) *hijackableBody {
+	return &hijackableBody{
+		body: body{
+			str:          str,
+			onFrameError: onFrameError,
+			reqDone:      done,
+		},
+		conn: conn,
 	}
 }
 
+func (r *hijackableBody) StreamCreator() StreamCreator {
+	return r.conn
+}
+
 func (r *body) Read(b []byte) (int, error) {
 	n, err := r.readImpl(b)
 	if err != nil {
diff --git a/http3/body_test.go b/http3/body_test.go
index d9d5c780..f50004dc 100644
--- a/http3/body_test.go
+++ b/http3/body_test.go
@@ -29,7 +29,7 @@ func (t bodyType) String() string {
 
 var _ = Describe("Body", func() {
 	var (
-		rb            *body
+		rb            io.ReadCloser
 		str           *mockquic.MockStream
 		buf           *bytes.Buffer
 		reqDone       chan struct{}
@@ -68,7 +68,7 @@ var _ = Describe("Body", func() {
 					rb = newRequestBody(str, errorCb)
 				case bodyTypeResponse:
 					reqDone = make(chan struct{})
-					rb = newResponseBody(str, reqDone, errorCb)
+					rb = newResponseBody(str, nil, reqDone, errorCb)
 				}
 			})
 
diff --git a/http3/client.go b/http3/client.go
index 498e2e75..afd50dd2 100644
--- a/http3/client.go
+++ b/http3/client.go
@@ -316,7 +316,7 @@ func (c *client) doRequest(
 			res.Header.Add(hf.Name, hf.Value)
 		}
 	}
-	respBody := newResponseBody(str, reqDone, func() {
+	respBody := newResponseBody(str, c.conn, reqDone, func() {
 		c.conn.CloseWithError(quic.ApplicationErrorCode(errorFrameUnexpected), "")
 	})