go-omni/httpx/trackingWriter.go

package httpx

import (
	"bufio"
	"io"
	"net"
	"net/http"
)

// trackingWriter does not buffer response bodies. Each wrapper has its own
// commit flag, so a standard middleware may safely wrap an already tracked writer.
type trackingWriter struct {
	http.ResponseWriter
	context   *Context
	committed bool
}

func (w *trackingWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *trackingWriter) inheritCommit(status int)    { w.mark(status) }
func (w *trackingWriter) mark(status int) {
	w.committed = true
	w.context.responded = true
	w.context.status = status
}
func (w *trackingWriter) WriteHeader(status int) {
	if w.committed {
		return
	}
	// Let the server validate the status before updating the observation.
	w.ResponseWriter.WriteHeader(status)
	if status == 101 || status >= 200 {
		w.mark(status)
	}
}
func (w *trackingWriter) Write(p []byte) (int, error) {
	if !w.committed {
		w.mark(http.StatusOK)
	}
	return w.ResponseWriter.Write(p)
}

type flushWriter struct{ *trackingWriter }

func (w flushWriter) Flush() { _ = w.FlushError() }
func (w flushWriter) FlushError() error {
	if !w.committed {
		w.WriteHeader(http.StatusOK)
	}
	return http.NewResponseController(w.ResponseWriter).Flush()
}

type hijackWriter struct{ *trackingWriter }

func (w hijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	conn, rw, err := w.ResponseWriter.(http.Hijacker).Hijack()
	if err == nil {
		w.context.responded = true
		w.committed = true
	}
	return conn, rw, err
}

type pushWriter struct{ *trackingWriter }

func (w pushWriter) Push(target string, options *http.PushOptions) error {
	return w.ResponseWriter.(http.Pusher).Push(target, options)
}

type readFromWriter struct{ *trackingWriter }

func (w readFromWriter) ReadFrom(r io.Reader) (int64, error) {
	// A writer-only view makes every write observable, including partial failures,
	// while io.Copy retains bounded streaming memory and source WriterTo support.
	return io.Copy(struct{ io.Writer }{w.trackingWriter}, r)
}

type closeNotifyWriter struct{ *trackingWriter }

func (w closeNotifyWriter) CloseNotify() <-chan bool {
	return w.ResponseWriter.(http.CloseNotifier).CloseNotify()
}

func trackWriter(writer http.ResponseWriter, c *Context) http.ResponseWriter {
	w := &trackingWriter{ResponseWriter: writer, context: c}
	mask := 0
	if _, ok := writer.(http.Flusher); ok {
		mask |= 1
	} else if _, ok := writer.(interface{ FlushError() error }); ok {
		mask |= 1
	}
	if _, ok := writer.(http.Hijacker); ok {
		mask |= 2
	}
	if _, ok := writer.(http.Pusher); ok {
		mask |= 4
	}
	if _, ok := writer.(io.ReaderFrom); ok {
		mask |= 8
	}
	if _, ok := writer.(http.CloseNotifier); ok {
		mask |= 16
	}
	switch mask {
	case 1:
		return struct {
			*trackingWriter
			flushCapability
		}{w, flushWriter{w}}
	case 2:
		return struct {
			*trackingWriter
			http.Hijacker
		}{w, hijackWriter{w}}
	case 3:
		return struct {
			*trackingWriter
			flushCapability
			http.Hijacker
		}{w, flushWriter{w}, hijackWriter{w}}
	case 4:
		return struct {
			*trackingWriter
			http.Pusher
		}{w, pushWriter{w}}
	case 5:
		return struct {
			*trackingWriter
			flushCapability
			http.Pusher
		}{w, flushWriter{w}, pushWriter{w}}
	case 6:
		return struct {
			*trackingWriter
			http.Hijacker
			http.Pusher
		}{w, hijackWriter{w}, pushWriter{w}}
	case 7:
		return struct {
			*trackingWriter
			flushCapability
			http.Hijacker
			http.Pusher
		}{w, flushWriter{w}, hijackWriter{w}, pushWriter{w}}
	case 8:
		return struct {
			*trackingWriter
			io.ReaderFrom
		}{w, readFromWriter{w}}
	case 9:
		return struct {
			*trackingWriter
			flushCapability
			io.ReaderFrom
		}{w, flushWriter{w}, readFromWriter{w}}
	case 10:
		return struct {
			*trackingWriter
			http.Hijacker
			io.ReaderFrom
		}{w, hijackWriter{w}, readFromWriter{w}}
	case 11:
		return struct {
			*trackingWriter
			flushCapability
			http.Hijacker
			io.ReaderFrom
		}{w, flushWriter{w}, hijackWriter{w}, readFromWriter{w}}
	case 12:
		return struct {
			*trackingWriter
			http.Pusher
			io.ReaderFrom
		}{w, pushWriter{w}, readFromWriter{w}}
	case 13:
		return struct {
			*trackingWriter
			flushCapability
			http.Pusher
			io.ReaderFrom
		}{w, flushWriter{w}, pushWriter{w}, readFromWriter{w}}
	case 14:
		return struct {
			*trackingWriter
			http.Hijacker
			http.Pusher
			io.ReaderFrom
		}{w, hijackWriter{w}, pushWriter{w}, readFromWriter{w}}
	case 15:
		return struct {
			*trackingWriter
			flushCapability
			http.Hijacker
			http.Pusher
			io.ReaderFrom
		}{w, flushWriter{w}, hijackWriter{w}, pushWriter{w}, readFromWriter{w}}
	case 16:
		return struct {
			*trackingWriter
			http.CloseNotifier
		}{w, closeNotifyWriter{w}}
	case 17:
		return struct {
			*trackingWriter
			flushCapability
			http.CloseNotifier
		}{w, flushWriter{w}, closeNotifyWriter{w}}
	case 18:
		return struct {
			*trackingWriter
			http.Hijacker
			http.CloseNotifier
		}{w, hijackWriter{w}, closeNotifyWriter{w}}
	case 19:
		return struct {
			*trackingWriter
			flushCapability
			http.Hijacker
			http.CloseNotifier
		}{w, flushWriter{w}, hijackWriter{w}, closeNotifyWriter{w}}
	case 20:
		return struct {
			*trackingWriter
			http.Pusher
			http.CloseNotifier
		}{w, pushWriter{w}, closeNotifyWriter{w}}
	case 21:
		return struct {
			*trackingWriter
			flushCapability
			http.Pusher
			http.CloseNotifier
		}{w, flushWriter{w}, pushWriter{w}, closeNotifyWriter{w}}
	case 22:
		return struct {
			*trackingWriter
			http.Hijacker
			http.Pusher
			http.CloseNotifier
		}{w, hijackWriter{w}, pushWriter{w}, closeNotifyWriter{w}}
	case 23:
		return struct {
			*trackingWriter
			flushCapability
			http.Hijacker
			http.Pusher
			http.CloseNotifier
		}{w, flushWriter{w}, hijackWriter{w}, pushWriter{w}, closeNotifyWriter{w}}
	case 24:
		return struct {
			*trackingWriter
			io.ReaderFrom
			http.CloseNotifier
		}{w, readFromWriter{w}, closeNotifyWriter{w}}
	case 25:
		return struct {
			*trackingWriter
			flushCapability
			io.ReaderFrom
			http.CloseNotifier
		}{w, flushWriter{w}, readFromWriter{w}, closeNotifyWriter{w}}
	case 26:
		return struct {
			*trackingWriter
			http.Hijacker
			io.ReaderFrom
			http.CloseNotifier
		}{w, hijackWriter{w}, readFromWriter{w}, closeNotifyWriter{w}}
	case 27:
		return struct {
			*trackingWriter
			flushCapability
			http.Hijacker
			io.ReaderFrom
			http.CloseNotifier
		}{w, flushWriter{w}, hijackWriter{w}, readFromWriter{w}, closeNotifyWriter{w}}
	case 28:
		return struct {
			*trackingWriter
			http.Pusher
			io.ReaderFrom
			http.CloseNotifier
		}{w, pushWriter{w}, readFromWriter{w}, closeNotifyWriter{w}}
	case 29:
		return struct {
			*trackingWriter
			flushCapability
			http.Pusher
			io.ReaderFrom
			http.CloseNotifier
		}{w, flushWriter{w}, pushWriter{w}, readFromWriter{w}, closeNotifyWriter{w}}
	case 30:
		return struct {
			*trackingWriter
			http.Hijacker
			http.Pusher
			io.ReaderFrom
			http.CloseNotifier
		}{w, hijackWriter{w}, pushWriter{w}, readFromWriter{w}, closeNotifyWriter{w}}
	case 31:
		return struct {
			*trackingWriter
			flushCapability
			http.Hijacker
			http.Pusher
			io.ReaderFrom
			http.CloseNotifier
		}{w, flushWriter{w}, hijackWriter{w}, pushWriter{w}, readFromWriter{w}, closeNotifyWriter{w}}
	}
	return w
}

type flushCapability interface {
	Flush()
	FlushError() error
}