github.com/ooni/psiphon/tunnel-core@v0.0.0-20230105123940-fe12a24c96ee/oovendor/quic-go/http3/response_writer.go (about)

     1  package http3
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"net/http"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/ooni/psiphon/tunnel-core/oovendor/quic-go"
    11  	"github.com/ooni/psiphon/tunnel-core/oovendor/quic-go/internal/utils"
    12  	"github.com/marten-seemann/qpack"
    13  )
    14  
    15  // DataStreamer lets the caller take over the stream. After a call to DataStream
    16  // the HTTP server library will not do anything else with the connection.
    17  //
    18  // It becomes the caller's responsibility to manage and close the stream.
    19  //
    20  // After a call to DataStream, the original Request.Body must not be used.
    21  type DataStreamer interface {
    22  	DataStream() quic.Stream
    23  }
    24  
    25  type responseWriter struct {
    26  	stream         quic.Stream // needed for DataStream()
    27  	bufferedStream *bufio.Writer
    28  
    29  	header         http.Header
    30  	status         int // status code passed to WriteHeader
    31  	headerWritten  bool
    32  	dataStreamUsed bool // set when DataSteam() is called
    33  
    34  	logger utils.Logger
    35  }
    36  
    37  var (
    38  	_ http.ResponseWriter = &responseWriter{}
    39  	_ http.Flusher        = &responseWriter{}
    40  	_ DataStreamer        = &responseWriter{}
    41  )
    42  
    43  func newResponseWriter(stream quic.Stream, logger utils.Logger) *responseWriter {
    44  	return &responseWriter{
    45  		header:         http.Header{},
    46  		stream:         stream,
    47  		bufferedStream: bufio.NewWriter(stream),
    48  		logger:         logger,
    49  	}
    50  }
    51  
    52  func (w *responseWriter) Header() http.Header {
    53  	return w.header
    54  }
    55  
    56  func (w *responseWriter) WriteHeader(status int) {
    57  	if w.headerWritten {
    58  		return
    59  	}
    60  
    61  	if status < 100 || status >= 200 {
    62  		w.headerWritten = true
    63  	}
    64  	w.status = status
    65  
    66  	var headers bytes.Buffer
    67  	enc := qpack.NewEncoder(&headers)
    68  	enc.WriteField(qpack.HeaderField{Name: ":status", Value: strconv.Itoa(status)})
    69  
    70  	for k, v := range w.header {
    71  		for index := range v {
    72  			enc.WriteField(qpack.HeaderField{Name: strings.ToLower(k), Value: v[index]})
    73  		}
    74  	}
    75  
    76  	buf := &bytes.Buffer{}
    77  	(&headersFrame{Length: uint64(headers.Len())}).Write(buf)
    78  	w.logger.Infof("Responding with %d", status)
    79  	if _, err := w.bufferedStream.Write(buf.Bytes()); err != nil {
    80  		w.logger.Errorf("could not write headers frame: %s", err.Error())
    81  	}
    82  	if _, err := w.bufferedStream.Write(headers.Bytes()); err != nil {
    83  		w.logger.Errorf("could not write header frame payload: %s", err.Error())
    84  	}
    85  	if !w.headerWritten {
    86  		w.Flush()
    87  	}
    88  }
    89  
    90  func (w *responseWriter) Write(p []byte) (int, error) {
    91  	if !w.headerWritten {
    92  		w.WriteHeader(200)
    93  	}
    94  	if !bodyAllowedForStatus(w.status) {
    95  		return 0, http.ErrBodyNotAllowed
    96  	}
    97  	df := &dataFrame{Length: uint64(len(p))}
    98  	buf := &bytes.Buffer{}
    99  	df.Write(buf)
   100  	if _, err := w.bufferedStream.Write(buf.Bytes()); err != nil {
   101  		return 0, err
   102  	}
   103  	return w.bufferedStream.Write(p)
   104  }
   105  
   106  func (w *responseWriter) Flush() {
   107  	if err := w.bufferedStream.Flush(); err != nil {
   108  		w.logger.Errorf("could not flush to stream: %s", err.Error())
   109  	}
   110  }
   111  
   112  func (w *responseWriter) usedDataStream() bool {
   113  	return w.dataStreamUsed
   114  }
   115  
   116  func (w *responseWriter) DataStream() quic.Stream {
   117  	w.dataStreamUsed = true
   118  	w.Flush()
   119  	return w.stream
   120  }
   121  
   122  // copied from http2/http2.go
   123  // bodyAllowedForStatus reports whether a given response status code
   124  // permits a body. See RFC 2616, section 4.4.
   125  func bodyAllowedForStatus(status int) bool {
   126  	switch {
   127  	case status >= 100 && status <= 199:
   128  		return false
   129  	case status == 204:
   130  		return false
   131  	case status == 304:
   132  		return false
   133  	}
   134  	return true
   135  }