43 lines
717 B
Go
43 lines
717 B
Go
package middleware
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
)
|
|
|
|
type ResponseWrapper struct {
|
|
http.ResponseWriter
|
|
|
|
status int
|
|
}
|
|
|
|
func WrapResponseWriter(w http.ResponseWriter) *ResponseWrapper {
|
|
return &ResponseWrapper{w, 0}
|
|
}
|
|
|
|
func (w *ResponseWrapper) WriteHeader(code int) {
|
|
w.status = code
|
|
w.ResponseWriter.WriteHeader(code)
|
|
}
|
|
|
|
func (w *ResponseWrapper) Status() int {
|
|
return w.status
|
|
}
|
|
|
|
func (w *ResponseWrapper) Write(b []byte) (int, error) {
|
|
if w.status == 0 {
|
|
w.status = http.StatusOK
|
|
}
|
|
|
|
size, err := w.ResponseWriter.Write(b)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("error writing response: %w", err)
|
|
}
|
|
|
|
return size, nil
|
|
}
|
|
|
|
func (w *ResponseWrapper) Unwrap() http.ResponseWriter {
|
|
return w.ResponseWriter
|
|
}
|