Wrap http.ResponseWriter to greatly increase middleware/logging accesibilty

README.md

@@ -286,6 +286,10 @@ func Logger() gin.HandlerFunc {
         // after request
         latency := time.Since(t)
         log.Print(latency)
+
+        // access the status we are sending
+        status := c.Writer.Status()
+        log.Println(status)
     }
 }
 

examples/logging/example_ansi.go

@@ -0,0 +1,82 @@
+package main
+
+import (
+	"log"
+	"strconv"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/mgutz/ansi"
+)
+
+var (
+	white  = ansi.ColorCode("white+h:black")
+	red    = ansi.ColorCode("red+h:black")
+	green  = ansi.ColorCode("green+h:black")
+	yellow = ansi.ColorCode("yellow+h:black")
+	blue   = ansi.ColorCode("blue+h:black")
+	reset  = ansi.ColorCode("reset")
+)
+
+//
+// Example of an extended ansi-colored logger using the
+// ctx.Writer.Status() function
+func logger(c *gin.Context) {
+	start := time.Now()
+
+	// save the IP of the requester
+	requester := c.Req.Header.Get("X-Real-IP")
+
+	// if the requester-header is empty, check the forwarded-header
+	if requester == "" {
+		requester = c.Req.Header.Get("X-Forwarded-For")
+	}
+
+	// if the requester is still empty, use the hard-coded address from the socket
+	if requester == "" {
+		requester = c.Req.RemoteAddr
+	}
+
+	// ... finally, log the fact we got a request
+	log.Printf("<-- %16s | %6s | %s\n", requester, c.Req.Method, c.Req.URL.Path)
+
+	c.Next()
+
+	var color string
+	if code := c.Writer.Status(); code >= 200 && code <= 299 {
+		color = green
+	} else if code >= 300 && code <= 399 {
+		color = white
+	} else if code >= 400 && code <= 499 {
+		color = yellow
+	} else {
+		color = red
+	}
+
+	log.Printf("--> %s%16s | %6d | %s | %s%s\n",
+		color,
+		requester, c.Writer.Status(), time.Since(start), c.Req.URL.Path,
+		reset,
+	)
+}
+
+func main() {
+	r := gin.New()
+	r.Use(gin.Recovery())
+	r.Use(logger)
+	// or modify func.logger to return a handler and use:
+	// r.Use(logger())
+
+	// Ping test
+	r.GET("/:code", func(c *gin.Context) {
+		asInt, err := strconv.ParseInt(c.Params.ByName("code"), 10, 32)
+		if err != nil {
+			c.String(400, err.Error())
+		} else {
+			c.String(int(asInt), c.Params.ByName("code"))
+		}
+	})
+
+	// Listen and Server in 0.0.0.0:8080
+	r.Run(":8081")
+}

examples/websocket/example_websocket.go

@@ -0,0 +1,58 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"strconv"
+
+	"code.google.com/p/go.net/websocket"
+
+	"github.com/gin-gonic/gin"
+)
+
+func main() {
+	r := gin.New()
+	r.GET("/", func(c *gin.Context) {
+		handler := websocket.Handler(func(conn *websocket.Conn) {
+			// This is the simplest possible echoserver
+			//
+			// For a more advanced handler you can use the
+			// conn.Read and conn.Write methods as the
+			// websocket.Conn type conforms to io.Reader+io.Writer
+
+			io.Copy(conn, conn)
+		})
+		handler.ServeHTTP(&c.Writer, c.Req)
+	})
+
+	go r.Run(":8080")
+
+	lock := make(chan bool)
+	go testServer(100, lock)
+	<-lock
+}
+
+func testServer(count int, done chan bool) {
+	client, err := websocket.Dial("ws://localhost:8080", "", "http://localhost/")
+	if err != nil {
+		panic(err)
+	}
+
+	for i := 0; i < count; i++ {
+		out := []byte(strconv.Itoa(i))
+		_, err = client.Write(out)
+		if err != nil {
+			panic(err)
+		}
+		fmt.Printf("Sent: %s\n", out)
+
+		var in = make([]byte, 512)
+		_, err = client.Read(in)
+		if err != nil {
+			panic(err)
+		}
+		fmt.Printf("Received: %s\n\n", in)
+	}
+
+	done <- true
+}

gin.go

@@ -35,7 +35,7 @@ type (
 	// manage the flow, validate the JSON of a request and render a JSON response for example.
 	Context struct {
 		Req      *http.Request
-		Writer   http.ResponseWriter
+		Writer   ResponseWriter
 		Keys     map[string]interface{}
 		Errors   ErrorMsgs
 		Params   httprouter.Params
@@ -101,7 +101,7 @@ func (engine *Engine) handle404(w http.ResponseWriter, req *http.Request) {
 	handlers := engine.combineHandlers(engine.handlers404)
 	c := engine.createContext(w, req, nil, handlers)
 	if engine.handlers404 == nil {
-		http.NotFound(c.Writer, c.Req)
+		http.NotFound(&c.Writer, c.Req)
 	} else {
 		c.Writer.WriteHeader(404)
 	}
@@ -125,7 +125,7 @@ func (engine *Engine) ServeFiles(path string, root http.FileSystem) {
 
 // ServeHTTP makes the router implement the http.Handler interface.
 func (engine *Engine) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	engine.router.ServeHTTP(w, req)
+	engine.router.ServeHTTP(&ResponseWriter{w, -1}, req)
 }
 
 func (engine *Engine) Run(addr string) {
@@ -138,7 +138,7 @@ func (engine *Engine) Run(addr string) {
 
 func (group *RouterGroup) createContext(w http.ResponseWriter, req *http.Request, params httprouter.Params, handlers []HandlerFunc) *Context {
 	return &Context{
-		Writer:   w,
+		Writer:   ResponseWriter{w, StatusUnset},
 		Req:      req,
 		index:    -1,
 		engine:   group.engine,
@@ -178,7 +178,7 @@ func (group *RouterGroup) Handle(method, p string, handlers []HandlerFunc) {
 	p = path.Join(group.prefix, p)
 	handlers = group.combineHandlers(handlers)
 	group.engine.router.Handle(method, p, func(w http.ResponseWriter, req *http.Request, params httprouter.Params) {
-		group.createContext(w, req, params, handlers).Next()
+		group.createContext(&ResponseWriter{w, -1}, req, params, handlers).Next()
 	})
 }
 
@@ -327,10 +327,10 @@ func (c *Context) JSON(code int, obj interface{}) {
 	if code >= 0 {
 		c.Writer.WriteHeader(code)
 	}
-	encoder := json.NewEncoder(c.Writer)
+	encoder := json.NewEncoder(&c.Writer)
 	if err := encoder.Encode(obj); err != nil {
 		c.Error(err, obj)
-		http.Error(c.Writer, err.Error(), 500)
+		http.Error(&c.Writer, err.Error(), 500)
 	}
 }
 
@@ -341,10 +341,10 @@ func (c *Context) XML(code int, obj interface{}) {
 	if code >= 0 {
 		c.Writer.WriteHeader(code)
 	}
-	encoder := xml.NewEncoder(c.Writer)
+	encoder := xml.NewEncoder(&c.Writer)
 	if err := encoder.Encode(obj); err != nil {
 		c.Error(err, obj)
-		http.Error(c.Writer, err.Error(), 500)
+		http.Error(&c.Writer, err.Error(), 500)
 	}
 }
 
@@ -356,12 +356,12 @@ func (c *Context) HTML(code int, name string, data interface{}) {
 	if code >= 0 {
 		c.Writer.WriteHeader(code)
 	}
-	if err := c.engine.HTMLTemplates.ExecuteTemplate(c.Writer, name, data); err != nil {
+	if err := c.engine.HTMLTemplates.ExecuteTemplate(&c.Writer, name, data); err != nil {
 		c.Error(err, map[string]interface{}{
 			"name": name,
 			"data": data,
 		})
-		http.Error(c.Writer, err.Error(), 500)
+		http.Error(&c.Writer, err.Error(), 500)
 	}
 }
 

response_writer.go

@@ -0,0 +1,60 @@
+package gin
+
+import (
+	"bufio"
+	"errors"
+	"net"
+	"net/http"
+)
+
+const (
+	StatusUnset int = -1
+)
+
+type ResponseWriter struct {
+	http.ResponseWriter
+	status int
+}
+
+func (w *ResponseWriter) WriteHeader(code int) {
+	w.status = code
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *ResponseWriter) Write(data []byte) (int, error) {
+	// net/http.Response.Write only has two options: 200 or 500
+	// we will follow that lead and defer to their logic
+
+	// check if the write gave an error and set status accordingly
+	size, err := w.ResponseWriter.Write(data)
+	if err != nil {
+		// error on write, we give a 500
+		w.status = http.StatusInternalServerError
+	} else if w.WasWritten() == false {
+		// everything went okay and we never set a custom
+		// status so 200 it is
+		w.status = http.StatusOK
+	}
+
+	// can easily tap into Content-Length here with 'size'
+	return size, err
+}
+
+// returns the status of the given response
+func (w *ResponseWriter) Status() int {
+	return w.status
+}
+
+// return a boolean acknowledging if a status code has all ready been set
+func (w *ResponseWriter) WasWritten() bool {
+	return w.status == StatusUnset
+}
+
+// allow connection hijacking
+func (w *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hijacker, ok := w.ResponseWriter.(http.Hijacker)
+	if !ok {
+		return nil, nil, errors.New("the ResponseWriter doesn't support the Hijacker interface")
+	}
+	return hijacker.Hijack()
+}