123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114 |
- package macaron
- import (
- "bufio"
- "fmt"
- "net"
- "net/http"
- )
- type ResponseWriter interface {
- http.ResponseWriter
- http.Flusher
-
- Status() int
-
- Written() bool
-
- Size() int
-
-
- Before(BeforeFunc)
- }
- type BeforeFunc func(ResponseWriter)
- func NewResponseWriter(method string, rw http.ResponseWriter) ResponseWriter {
- return &responseWriter{method, rw, 0, 0, nil}
- }
- type responseWriter struct {
- method string
- 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) (size int, err error) {
- if !rw.Written() {
-
- rw.WriteHeader(http.StatusOK)
- }
- if rw.method != "HEAD" {
- 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 BeforeFunc) {
- 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()
- }
- }
|