Fast and extensible multi-platform HTTP/1-2-3 web server with automatic HTTPS
Mirrored from https://github.com/Slazee/caddy.git
Synced Nov 12, 14:30 UTC
caddyhttp: record num. bytes read when response writer is hijacked (#6173)
* record the number of bytes read when response writer is hijacked * record body size when not nil
WeidiDeng committed
Apr 17, 2024 at 23:00 UTC
e0daa39cd3373d26dcf78ae0d30ac24d1df5dd57
2 files changed
+42
modules/caddyhttp/responsewriter.go
+37
@@ -66,6 +66,8 @@ type responseRecorder struct {
66
size int
67
wroteHeader bool
68
stream bool
69
+
70
+ readSize *int
71
}
72
73
// NewResponseRecorder returns a new ResponseRecorder that can be
@@ -240,6 +242,12 @@ func (rr *responseRecorder) FlushError() error {
242
return nil
243
}
244
245
+// Private interface so it can only be used in this package
246
+// #TODO: maybe export it later
247
+func (rr *responseRecorder) setReadSize(size *int) {
248
+ rr.readSize = size
249
+}
250
+
251
func (rr *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
252
//nolint:bodyclose
253
conn, brw, err := http.NewResponseController(rr.ResponseWriterWrapper).Hijack()
@@ -249,6 +257,15 @@ func (rr *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
257
// Per http documentation, returned bufio.Writer is empty, but bufio.Read maybe not
258
conn = &hijackedConn{conn, rr}
259
brw.Writer.Reset(conn)
260
+
261
+ buffered := brw.Reader.Buffered()
262
+ if buffered != 0 {
263
+ conn.(*hijackedConn).updateReadSize(buffered)
264
+ data, _ := brw.Peek(buffered)
265
+ brw.Reader.Reset(io.MultiReader(bytes.NewReader(data), conn))
266
+ } else {
267
+ brw.Reader.Reset(conn)
268
+ }
269
return conn, brw, nil
270
}
271
@@ -258,6 +275,24 @@ type hijackedConn struct {
275
rr *responseRecorder
276
}
277
278
+func (hc *hijackedConn) updateReadSize(n int) {
279
+ if hc.rr.readSize != nil {
280
+ *hc.rr.readSize += n
281
+ }
282
+}
283
+
284
+func (hc *hijackedConn) Read(p []byte) (int, error) {
285
+ n, err := hc.Conn.Read(p)
286
+ hc.updateReadSize(n)
287
+ return n, err
288
+}
289
+
290
+func (hc *hijackedConn) WriteTo(w io.Writer) (int64, error) {
291
+ n, err := io.Copy(w, hc.Conn)
292
+ hc.updateReadSize(int(n))
293
+ return n, err
294
+}
295
+
296
func (hc *hijackedConn) Write(p []byte) (int, error) {
297
n, err := hc.Conn.Write(p)
298
hc.rr.size += n
@@ -298,4 +333,6 @@ var (
333
_ io.ReaderFrom = (*ResponseWriterWrapper)(nil)
334
_ io.ReaderFrom = (*responseRecorder)(nil)
335
_ io.ReaderFrom = (*hijackedConn)(nil)
336
+
337
+ _ io.WriterTo = (*hijackedConn)(nil)
338
)
modules/caddyhttp/server.go
+5
@@ -326,6 +326,11 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
326
if r.Body != nil {
327
bodyReader = &lengthReader{Source: r.Body}
328
r.Body = bodyReader
329
+
330
+ // should always be true, private interface can only be referenced in the same package
331
+ if setReadSizer, ok := wrec.(interface{ setReadSize(*int) }); ok {
332
+ setReadSizer.setReadSize(&bodyReader.Length)
333
+ }
334
}
335
336
// capture the original version of the request