github.com/tumi8/quic-go@v0.37.4-tum/http3/body.go (about)

     1  package http3
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/tumi8/quic-go"
     9  )
    10  
    11  // The HTTPStreamer allows taking over a HTTP/3 stream. The interface is implemented by:
    12  // * for the server: the http.Request.Body
    13  // * for the client: the http.Response.Body
    14  // On the client side, the stream will be closed for writing, unless the DontCloseRequestStream RoundTripOpt was set.
    15  // When a stream is taken over, it's the caller's responsibility to close the stream.
    16  type HTTPStreamer interface {
    17  	HTTPStream() Stream
    18  }
    19  
    20  type StreamCreator interface {
    21  	// Context returns a context that is cancelled when the underlying connection is closed.
    22  	Context() context.Context
    23  	OpenStream() (quic.Stream, error)
    24  	OpenStreamSync(context.Context) (quic.Stream, error)
    25  	OpenUniStream() (quic.SendStream, error)
    26  	OpenUniStreamSync(context.Context) (quic.SendStream, error)
    27  	LocalAddr() net.Addr
    28  	RemoteAddr() net.Addr
    29  	ConnectionState() quic.ConnectionState
    30  }
    31  
    32  var _ StreamCreator = quic.Connection(nil)
    33  
    34  // A Hijacker allows hijacking of the stream creating part of a quic.Session from a http.Response.Body.
    35  // It is used by WebTransport to create WebTransport streams after a session has been established.
    36  type Hijacker interface {
    37  	StreamCreator() StreamCreator
    38  }
    39  
    40  // The body of a http.Request or http.Response.
    41  type body struct {
    42  	str quic.Stream
    43  
    44  	wasHijacked bool // set when HTTPStream is called
    45  }
    46  
    47  var (
    48  	_ io.ReadCloser = &body{}
    49  	_ HTTPStreamer  = &body{}
    50  )
    51  
    52  func newRequestBody(str Stream) *body {
    53  	return &body{str: str}
    54  }
    55  
    56  func (r *body) HTTPStream() Stream {
    57  	r.wasHijacked = true
    58  	return r.str
    59  }
    60  
    61  func (r *body) wasStreamHijacked() bool {
    62  	return r.wasHijacked
    63  }
    64  
    65  func (r *body) Read(b []byte) (int, error) {
    66  	return r.str.Read(b)
    67  }
    68  
    69  func (r *body) Close() error {
    70  	r.str.CancelRead(quic.StreamErrorCode(ErrCodeRequestCanceled))
    71  	return nil
    72  }
    73  
    74  type hijackableBody struct {
    75  	body
    76  	conn quic.Connection // only needed to implement Hijacker
    77  
    78  	// only set for the http.Response
    79  	// The channel is closed when the user is done with this response:
    80  	// either when Read() errors, or when Close() is called.
    81  	reqDone       chan<- struct{}
    82  	reqDoneClosed bool
    83  }
    84  
    85  var (
    86  	_ Hijacker     = &hijackableBody{}
    87  	_ HTTPStreamer = &hijackableBody{}
    88  )
    89  
    90  func newResponseBody(str Stream, conn quic.Connection, done chan<- struct{}) *hijackableBody {
    91  	return &hijackableBody{
    92  		body: body{
    93  			str: str,
    94  		},
    95  		reqDone: done,
    96  		conn:    conn,
    97  	}
    98  }
    99  
   100  func (r *hijackableBody) StreamCreator() StreamCreator {
   101  	return r.conn
   102  }
   103  
   104  func (r *hijackableBody) Read(b []byte) (int, error) {
   105  	n, err := r.str.Read(b)
   106  	if err != nil {
   107  		r.requestDone()
   108  	}
   109  	return n, err
   110  }
   111  
   112  func (r *hijackableBody) requestDone() {
   113  	if r.reqDoneClosed || r.reqDone == nil {
   114  		return
   115  	}
   116  	if r.reqDone != nil {
   117  		close(r.reqDone)
   118  	}
   119  	r.reqDoneClosed = true
   120  }
   121  
   122  func (r *body) StreamID() quic.StreamID {
   123  	return r.str.StreamID()
   124  }
   125  
   126  func (r *hijackableBody) Close() error {
   127  	r.requestDone()
   128  	// If the EOF was read, CancelRead() is a no-op.
   129  	r.str.CancelRead(quic.StreamErrorCode(ErrCodeRequestCanceled))
   130  	return nil
   131  }
   132  
   133  func (r *hijackableBody) HTTPStream() Stream {
   134  	return r.str
   135  }