-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
1, bug fix. 2, add custom ResponseWriter for get Status Code
- Loading branch information
Showing
7 changed files
with
131 additions
and
18 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,96 @@ | ||
package xingyun | ||
|
||
import ( | ||
"bufio" | ||
"fmt" | ||
"net" | ||
"net/http" | ||
) | ||
|
||
// ResponseWriter is a wrapper around http.ResponseWriter that provides extra information about | ||
// the response. It is recommended that middleware handlers use this construct to wrap a responsewriter | ||
// if the functionality calls for it. | ||
type ResponseWriter interface { | ||
http.ResponseWriter | ||
http.Flusher | ||
// Status returns the status code of the response or 0 if the response has not been written. | ||
Status() int | ||
// Written returns whether or not the ResponseWriter has been written. | ||
Written() bool | ||
// Size returns the size of the response body. | ||
Size() int | ||
// Before allows for a function to be called before the ResponseWriter has been written to. This is | ||
// useful for setting headers or any other operations that must happen before a response has been written. | ||
Before(func(ResponseWriter)) | ||
} | ||
|
||
type beforeFunc func(ResponseWriter) | ||
|
||
// NewResponseWriter creates a ResponseWriter that wraps an http.ResponseWriter | ||
func NewResponseWriter(rw http.ResponseWriter) ResponseWriter { | ||
return &responseWriter{rw, 0, 0, nil} | ||
} | ||
|
||
type responseWriter struct { | ||
http.ResponseWriter | ||
status int | ||
size int | ||
beforeFuncs []beforeFunc | ||
} | ||
|
||
func (rw *responseWriter) WriteHeader(s int) { | ||
rw.callBefore() | ||
rw.ResponseWriter.WriteHeader(s) | ||
rw.status = 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) Status() int { | ||
return rw.status | ||
} | ||
|
||
func (rw *responseWriter) Size() int { | ||
return rw.size | ||
} | ||
|
||
func (rw *responseWriter) Written() bool { | ||
return rw.status != 0 | ||
} | ||
|
||
func (rw *responseWriter) Before(before func(ResponseWriter)) { | ||
rw.beforeFuncs = append(rw.beforeFuncs, before) | ||
} | ||
|
||
func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { | ||
hijacker, ok := rw.ResponseWriter.(http.Hijacker) | ||
if !ok { | ||
return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface") | ||
} | ||
return hijacker.Hijack() | ||
} | ||
|
||
func (rw *responseWriter) CloseNotify() <-chan bool { | ||
return rw.ResponseWriter.(http.CloseNotifier).CloseNotify() | ||
} | ||
|
||
func (rw *responseWriter) callBefore() { | ||
for i := len(rw.beforeFuncs) - 1; i >= 0; i-- { | ||
rw.beforeFuncs[i](rw) | ||
} | ||
} | ||
|
||
func (rw *responseWriter) Flush() { | ||
flusher, ok := rw.ResponseWriter.(http.Flusher) | ||
if ok { | ||
flusher.Flush() | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters