github.com/MerlinKodo/quic-go@v0.39.2/http3/body.go (about)

     1  package http3
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/MerlinKodo/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  	n, err := r.str.Read(b)
    67  	return n, maybeReplaceError(err)
    68  }
    69  
    70  func (r *body) Close() error {
    71  	r.str.CancelRead(quic.StreamErrorCode(ErrCodeRequestCanceled))
    72  	return nil
    73  }
    74  
    75  type hijackableBody struct {
    76  	body
    77  	conn quic.Connection // only needed to implement Hijacker
    78  
    79  	// only set for the http.Response
    80  	// The channel is closed when the user is done with this response:
    81  	// either when Read() errors, or when Close() is called.
    82  	reqDone       chan<- struct{}
    83  	reqDoneClosed bool
    84  }
    85  
    86  var (
    87  	_ Hijacker     = &hijackableBody{}
    88  	_ HTTPStreamer = &hijackableBody{}
    89  )
    90  
    91  func newResponseBody(str Stream, conn quic.Connection, done chan<- struct{}) *hijackableBody {
    92  	return &hijackableBody{
    93  		body: body{
    94  			str: str,
    95  		},
    96  		reqDone: done,
    97  		conn:    conn,
    98  	}
    99  }
   100  
   101  func (r *hijackableBody) StreamCreator() StreamCreator {
   102  	return r.conn
   103  }
   104  
   105  func (r *hijackableBody) Read(b []byte) (int, error) {
   106  	n, err := r.str.Read(b)
   107  	if err != nil {
   108  		r.requestDone()
   109  	}
   110  	return n, maybeReplaceError(err)
   111  }
   112  
   113  func (r *hijackableBody) requestDone() {
   114  	if r.reqDoneClosed || r.reqDone == nil {
   115  		return
   116  	}
   117  	if r.reqDone != nil {
   118  		close(r.reqDone)
   119  	}
   120  	r.reqDoneClosed = true
   121  }
   122  
   123  func (r *body) StreamID() quic.StreamID {
   124  	return r.str.StreamID()
   125  }
   126  
   127  func (r *hijackableBody) Close() error {
   128  	r.requestDone()
   129  	// If the EOF was read, CancelRead() is a no-op.
   130  	r.str.CancelRead(quic.StreamErrorCode(ErrCodeRequestCanceled))
   131  	return nil
   132  }
   133  
   134  func (r *hijackableBody) HTTPStream() Stream {
   135  	return r.str
   136  }