github.com/Psiphon-Labs/psiphon-tunnel-core@v2.0.28+incompatible/psiphon/common/quic/gquic-go/h2quic/response_writer.go (about)

     1  package h2quic
     2  
     3  import (
     4  	"bytes"
     5  	"net/http"
     6  	"strconv"
     7  	"strings"
     8  	"sync"
     9  
    10  	quic "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/quic/gquic-go"
    11  	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/quic/gquic-go/internal/protocol"
    12  	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/quic/gquic-go/internal/utils"
    13  	"golang.org/x/net/http2"
    14  	"golang.org/x/net/http2/hpack"
    15  )
    16  
    17  type responseWriter struct {
    18  	dataStreamID protocol.StreamID
    19  	dataStream   quic.Stream
    20  
    21  	headerStream      quic.Stream
    22  	headerStreamMutex *sync.Mutex
    23  
    24  	header        http.Header
    25  	status        int // status code passed to WriteHeader
    26  	headerWritten bool
    27  
    28  	logger utils.Logger
    29  }
    30  
    31  func newResponseWriter(
    32  	headerStream quic.Stream,
    33  	headerStreamMutex *sync.Mutex,
    34  	dataStream quic.Stream,
    35  	dataStreamID protocol.StreamID,
    36  	logger utils.Logger,
    37  ) *responseWriter {
    38  	return &responseWriter{
    39  		header:            http.Header{},
    40  		headerStream:      headerStream,
    41  		headerStreamMutex: headerStreamMutex,
    42  		dataStream:        dataStream,
    43  		dataStreamID:      dataStreamID,
    44  		logger:            logger,
    45  	}
    46  }
    47  
    48  func (w *responseWriter) Header() http.Header {
    49  	return w.header
    50  }
    51  
    52  func (w *responseWriter) WriteHeader(status int) {
    53  	if w.headerWritten {
    54  		return
    55  	}
    56  	w.headerWritten = true
    57  	w.status = status
    58  
    59  	var headers bytes.Buffer
    60  	enc := hpack.NewEncoder(&headers)
    61  	enc.WriteField(hpack.HeaderField{Name: ":status", Value: strconv.Itoa(status)})
    62  
    63  	for k, v := range w.header {
    64  		for index := range v {
    65  			enc.WriteField(hpack.HeaderField{Name: strings.ToLower(k), Value: v[index]})
    66  		}
    67  	}
    68  
    69  	w.logger.Infof("Responding with %d", status)
    70  	w.headerStreamMutex.Lock()
    71  	defer w.headerStreamMutex.Unlock()
    72  	h2framer := http2.NewFramer(w.headerStream, nil)
    73  	err := h2framer.WriteHeaders(http2.HeadersFrameParam{
    74  		StreamID:      uint32(w.dataStreamID),
    75  		EndHeaders:    true,
    76  		BlockFragment: headers.Bytes(),
    77  	})
    78  	if err != nil {
    79  		w.logger.Errorf("could not write h2 header: %s", err.Error())
    80  	}
    81  }
    82  
    83  func (w *responseWriter) Write(p []byte) (int, error) {
    84  	if !w.headerWritten {
    85  		w.WriteHeader(200)
    86  	}
    87  	if !bodyAllowedForStatus(w.status) {
    88  		return 0, http.ErrBodyNotAllowed
    89  	}
    90  	return w.dataStream.Write(p)
    91  }
    92  
    93  func (w *responseWriter) Flush() {}
    94  
    95  // This is a NOP. Use http.Request.Context
    96  func (w *responseWriter) CloseNotify() <-chan bool { return make(<-chan bool) }
    97  
    98  // test that we implement http.Flusher
    99  var _ http.Flusher = &responseWriter{}
   100  
   101  // copied from http2/http2.go
   102  // bodyAllowedForStatus reports whether a given response status code
   103  // permits a body. See RFC 2616, section 4.4.
   104  func bodyAllowedForStatus(status int) bool {
   105  	switch {
   106  	case status >= 100 && status <= 199:
   107  		return false
   108  	case status == 204:
   109  		return false
   110  	case status == 304:
   111  		return false
   112  	}
   113  	return true
   114  }