mirror of
https://github.com/mattermost/mattermost.git
synced 2025-02-25 18:55:24 -06:00
aafea55976
* ResponseWriter wrapper to get status code For our metrics, we need the status code returned by a request so this wrapper includes a new method StatusCode() that includes the desired code * Shadow the responsewriter variable in the handlers In order to avoid confusion to people deciding what variable to use. I've also changed the tests to reflect this change and added a new one that checks the Flush method works
63 lines
1.5 KiB
Go
63 lines
1.5 KiB
Go
// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
|
|
// See LICENSE.txt for license information.
|
|
|
|
package web
|
|
|
|
import (
|
|
"bufio"
|
|
"errors"
|
|
"net"
|
|
"net/http"
|
|
)
|
|
|
|
type responseWriterWrapper struct {
|
|
http.ResponseWriter
|
|
statusCode int
|
|
statusCodeWritten bool
|
|
hijacker http.Hijacker
|
|
flusher http.Flusher
|
|
}
|
|
|
|
func newWrappedWriter(original http.ResponseWriter) *responseWriterWrapper {
|
|
hijacker, _ := original.(http.Hijacker)
|
|
flusher, _ := original.(http.Flusher)
|
|
return &responseWriterWrapper{
|
|
ResponseWriter: original,
|
|
statusCodeWritten: false,
|
|
hijacker: hijacker,
|
|
flusher: flusher,
|
|
}
|
|
}
|
|
|
|
func (rw *responseWriterWrapper) StatusCode() int {
|
|
return rw.statusCode
|
|
}
|
|
|
|
func (rw *responseWriterWrapper) WriteHeader(statusCode int) {
|
|
rw.statusCode = statusCode
|
|
rw.statusCodeWritten = true
|
|
rw.ResponseWriter.WriteHeader(statusCode)
|
|
}
|
|
|
|
func (rw *responseWriterWrapper) Write(data []byte) (int, error) {
|
|
if !rw.statusCodeWritten {
|
|
rw.statusCode = http.StatusOK
|
|
}
|
|
return rw.ResponseWriter.Write(data)
|
|
}
|
|
|
|
// Using as embedded makes the ResponseWrite be stored as interface and that way
|
|
// it loses the access to the implementation for Hijack or Flush
|
|
func (rw *responseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|
if rw.hijacker == nil {
|
|
return nil, nil, errors.New("Hijacker interface not supported by the wrapped ResponseWriter")
|
|
}
|
|
return rw.hijacker.Hijack()
|
|
}
|
|
|
|
func (rw *responseWriterWrapper) Flush() {
|
|
if rw.flusher != nil {
|
|
rw.flusher.Flush()
|
|
}
|
|
}
|