WebSocket API

Signed-off-by: Vishal Rana <vr@labstack.com>

context.go

@@ -15,6 +15,8 @@ import (
 	"strings"
 	"time"
 
+	"golang.org/x/net/websocket"
+
 	"github.com/labstack/echo/log"
 
 	"bytes"
@@ -40,6 +42,12 @@ type (
 		// Request returns `*Response`.
 		Response() *Response
 
+		// WebSocket returns `*websocket.Conn`.
+		WebSocket() *websocket.Conn
+
+		// SetWebSocket sets `*websocket.Conn`.
+		SetWebSocket(*websocket.Conn)
+
 		// IsTLS returns true if HTTP connection is TLS otherwise false.
 		IsTLS() bool
 
@@ -194,15 +202,16 @@ type (
 	}
 
 	echoContext struct {
-		context  context.Context
-		request  *http.Request
-		response *Response
-		path     string
-		pnames   []string
-		pvalues  []string
-		query    url.Values
-		handler  HandlerFunc
-		echo     *Echo
+		context   context.Context
+		request   *http.Request
+		response  *Response
+		webSocket *websocket.Conn
+		path      string
+		pnames    []string
+		pvalues   []string
+		query     url.Values
+		handler   HandlerFunc
+		echo      *Echo
 	}
 )
 
@@ -243,6 +252,14 @@ func (c *echoContext) Response() *Response {
 	return c.response
 }
 
+func (c *echoContext) WebSocket() *websocket.Conn {
+	return c.webSocket
+}
+
+func (c *echoContext) SetWebSocket(ws *websocket.Conn) {
+	c.webSocket = ws
+}
+
 func (c *echoContext) IsTLS() bool {
 	return c.request.TLS != nil
 }

echo.go

@@ -52,6 +52,7 @@ import (
 	"time"
 
 	"golang.org/x/net/context"
+	"golang.org/x/net/websocket"
 
 	"github.com/labstack/echo/log"
 	glog "github.com/labstack/gommon/log"
@@ -376,6 +377,22 @@ func (e *Echo) File(path, file string) {
 	})
 }
 
+// WebSocket registers a new WebSocket route for a path with matching handler in
+// the router with optional route-level middleware.
+func (e *Echo) WebSocket(path string, h HandlerFunc, m ...MiddlewareFunc) {
+	e.GET(path, func(c Context) (err error) {
+		wss := websocket.Server{
+			Handler: func(ws *websocket.Conn) {
+				c.SetWebSocket(ws)
+				c.Response().Status = http.StatusSwitchingProtocols
+				err = h(c)
+			},
+		}
+		wss.ServeHTTP(c.Response(), c.Request())
+		return err
+	}, m...)
+}
+
 func (e *Echo) add(method, path string, handler HandlerFunc, middleware ...MiddlewareFunc) {
 	name := handlerName(handler)
 	e.router.Add(method, path, func(c Context) error {

echo_test.go

@@ -2,10 +2,13 @@ package echo
 
 import (
 	"bytes"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
+	"golang.org/x/net/websocket"
+
 	"reflect"
 	"strings"
 
@@ -246,6 +249,27 @@ func TestEchoMatch(t *testing.T) { // JFC
 	})
 }
 
+func TestEchoWebSocket(t *testing.T) {
+	e := New()
+	e.WebSocket("/ws", func(c Context) error {
+		c.WebSocket().Write([]byte("test"))
+		return nil
+	})
+	srv := httptest.NewServer(e)
+	defer srv.Close()
+	addr := srv.Listener.Addr().String()
+	origin := "http://localhost"
+	url := fmt.Sprintf("ws://%s/ws", addr)
+	ws, err := websocket.Dial(url, "", origin)
+	if assert.NoError(t, err) {
+		ws.Write([]byte("test\n"))
+		defer ws.Close()
+		buf := new(bytes.Buffer)
+		buf.ReadFrom(ws)
+		assert.Equal(t, "test", buf.String())
+	}
+}
+
 func TestEchoURL(t *testing.T) {
 	e := New()
 	static := func(Context) error { return nil }