From b29ad0afa15f8f34f9825cadc31ee97559ebcfd7 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Giovanini" <mail@gabrielgio.me>
Date: Fri, 22 Mar 2024 18:28:31 +0100
Subject: feat: Add metrics

---
 ext.go | 75 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)
 create mode 100644 ext.go

(limited to 'ext.go')

diff --git a/ext.go b/ext.go
new file mode 100644
index 0000000..f6a0af8
--- /dev/null
+++ b/ext.go
@@ -0,0 +1,75 @@
+package main
+
+import (
+	"io"
+	"net/http"
+)
+
+type ResponseWriter interface {
+	http.ResponseWriter
+	Status() int
+}
+
+type responseWriter struct {
+	http.ResponseWriter
+	pendingStatus int
+	status        int
+	size          int
+}
+
+func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
+	return &responseWriter{
+		ResponseWriter: rw,
+	}
+}
+
+func (rw *responseWriter) WriteHeader(s int) {
+	if rw.Written() {
+		return
+	}
+
+	rw.pendingStatus = s
+
+	if rw.Written() {
+		return
+	}
+
+	rw.status = s
+	rw.ResponseWriter.WriteHeader(s)
+}
+
+func (rw *responseWriter) Write(b []byte) (int, error) {
+	if !rw.Written() {
+		// The status will be StatusOK if WriteHeader has not been called yet
+		rw.WriteHeader(http.StatusOK)
+	}
+	size, err := rw.ResponseWriter.Write(b)
+	rw.size += size
+	return size, err
+}
+
+func (rw *responseWriter) ReadFrom(r io.Reader) (n int64, err error) {
+	if !rw.Written() {
+		// The status will be StatusOK if WriteHeader has not been called yet
+		rw.WriteHeader(http.StatusOK)
+	}
+	n, err = io.Copy(rw.ResponseWriter, r)
+	rw.size += int(n)
+	return
+}
+
+func (rw *responseWriter) Unwrap() http.ResponseWriter {
+	return rw.ResponseWriter
+}
+
+func (rw *responseWriter) Status() int {
+	if rw.Written() {
+		return rw.status
+	}
+
+	return rw.pendingStatus
+}
+
+func (rw *responseWriter) Written() bool {
+	return rw.status != 0
+}
-- 
cgit v1.2.3