1package http3
2
3import (
4	"bufio"
5	"bytes"
6	"net/http"
7	"strconv"
8	"strings"
9
10	"github.com/lucas-clemente/quic-go"
11	"github.com/lucas-clemente/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.
21type DataStreamer interface {
22	DataStream() quic.Stream
23}
24
25type 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
37var (
38	_ http.ResponseWriter = &responseWriter{}
39	_ http.Flusher        = &responseWriter{}
40	_ DataStreamer        = &responseWriter{}
41)
42
43func 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
52func (w *responseWriter) Header() http.Header {
53	return w.header
54}
55
56func (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
90func (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
106func (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
112func (w *responseWriter) usedDataStream() bool {
113	return w.dataStreamUsed
114}
115
116func (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.
125func 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}
136