File size: 1,470 Bytes
1f10f31 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | package oasmiddleware
import (
"bytes"
"net/http"
)
// ResponseWriterWrapper leverage to the original response write to be able to access
// the written body and the status code of the response. This is used in the response middleware
// but can also be used in a logging middleware to log the status code of the response.
type ResponseWriterWrapper struct {
w *http.ResponseWriter
body *bytes.Buffer
statusCode *int
}
func (rww ResponseWriterWrapper) Body() *bytes.Buffer {
return rww.body
}
func (rww ResponseWriterWrapper) StatusCode() *int {
return rww.statusCode
}
func NewResponseWriterWrapper(w http.ResponseWriter) ResponseWriterWrapper {
var (
buf bytes.Buffer
statusCode = 200
)
return ResponseWriterWrapper{
w: &w,
body: &buf,
statusCode: &statusCode,
}
}
// Write function overwrites the http.ResponseWriter Header() function
func (rww ResponseWriterWrapper) Write(buf []byte) (int, error) {
rww.body.Write(buf)
return (*rww.w).Write(buf)
}
// Header function overwrites the http.ResponseWriter Header() function
func (rww ResponseWriterWrapper) Header() http.Header {
return (*rww.w).Header()
}
// WriteHeader function overwrites the http.ResponseWriter WriteHeader() function
func (rww ResponseWriterWrapper) WriteHeader(statusCode int) {
*rww.statusCode = statusCode
(*rww.w).WriteHeader(statusCode)
}
func (rww ResponseWriterWrapper) Unwrap() http.ResponseWriter {
return *rww.w
}
|