github.com/TugasAkhir-QUIC/quic-go@v0.0.2-0.20240215011318-d20e25a9054c/http3/body.go (about)

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