Use golang websockets

Author: thani-sh

server/method_context.go

@@ -2,21 +2,23 @@ package server
 
 import (
 	"errors"
+
+	"golang.org/x/net/websocket"
 )
 
 type MethodContext struct {
 	ID      string
 	Args    []interface{}
-	Res     Response
+	Conn    *websocket.Conn
 	Done    bool
 	Updated bool
 }
 
-func NewMethodContext(m Message, res Response) MethodContext {
+func NewMethodContext(m Message, ws *websocket.Conn) MethodContext {
 	ctx := MethodContext{}
 	ctx.ID = m.ID
 	ctx.Args = m.Params
-	ctx.Res = res
+	ctx.Conn = ws
 	return ctx
 }
 
@@ -27,7 +29,7 @@ func (ctx *MethodContext) SendResult(r interface{}) error {
 	}
 
 	ctx.Done = true
-	return ctx.Res.WriteJSON(map[string]interface{}{
+	return websocket.JSON.Send(ctx.Conn, map[string]interface{}{
 		"msg":    "result",
 		"id":     ctx.ID,
 		"result": r,
@@ -41,7 +43,7 @@ func (ctx *MethodContext) SendError(e string) error {
 	}
 
 	ctx.Done = true
-	return ctx.Res.WriteJSON(map[string]interface{}{
+	return websocket.JSON.Send(ctx.Conn, map[string]interface{}{
 		"msg": "result",
 		"id":  ctx.ID,
 		"error": map[string]string{
@@ -57,7 +59,7 @@ func (ctx *MethodContext) SendUpdated() error {
 	}
 
 	ctx.Updated = true
-	return ctx.Res.WriteJSON(map[string]interface{}{
+	return websocket.JSON.Send(ctx.Conn, map[string]interface{}{
 		"msg":     "updated",
 		"methods": []string{ctx.ID},
 	})

server/server.go

@@ -1,53 +1,45 @@
 package server
 
 import (
-	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 
-	"github.com/gorilla/websocket"
 	"github.com/meteorhacks/goddp/utils/random"
+	"golang.org/x/net/websocket"
 )
 
 func New() Server {
 	s := Server{}
+	s.wsserver = websocket.Server{Handler: s.wsHandler, Handshake: handshake}
 	s.methods = make(map[string]MethodFn)
-	s.upgrader = websocket.Upgrader{
-		ReadBufferSize:  1024,
-		WriteBufferSize: 1024,
-		CheckOrigin:     checkOrigin,
-	}
-
 	return s
 }
 
 type Server struct {
 	methods  map[string]MethodFn
-	upgrader websocket.Upgrader
+	wsserver websocket.Server
 }
 
 func (s *Server) Listen(addr string) error {
-	http.HandleFunc("/websocket", s.Handler)
-	http.HandleFunc("/sockjs/websocket", s.Handler)
+	http.Handle("/websocket", s.wsserver)
+	http.Handle("/sockjs/websocket", s.wsserver)
 	return http.ListenAndServe(addr, nil)
 }
 
 func (s *Server) Method(name string, fn MethodFn) {
 	s.methods[name] = fn
 }
 
-func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
-	ws, err := s.upgrader.Upgrade(w, r, nil)
-	if err != nil {
-		// TODO => handle non-websocket requests
-		return
-	}
-
+func (s *Server) wsHandler(ws *websocket.Conn) {
 	for {
-		msg, err := readMessage(ws)
+		var msg Message
+		if err := websocket.JSON.Receive(ws, &msg); err != nil {
+			if err != io.EOF {
+				fmt.Println("Read Error: ", err)
+			}
 
-		if err != nil {
 			break
 		}
 
@@ -59,6 +51,7 @@ func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
 		case "method":
 			handleMethod(s, ws, msg)
 		default:
+			fmt.Println("Error: unknown message type", msg)
 			// TODO => send "error" ddp message
 			break
 		}
@@ -67,52 +60,33 @@ func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
 	ws.Close()
 }
 
-func checkOrigin(r *http.Request) bool {
-	return true
-}
-
-func readMessage(req Request) (Message, error) {
-	t, str, err := req.ReadMessage()
-	msg := Message{}
-
-	if err != nil {
-		return msg, err
-	}
-
-	if t != 1 {
-		err = errors.New("DDP does not supports binary streams yet")
-		return msg, err
-	}
-
-	if err := json.Unmarshal(str, &msg); err != nil {
-		return msg, err
-	}
-
-	return msg, nil
+func handshake(config *websocket.Config, req *http.Request) error {
+	// accept all connections
+	return nil
 }
 
-func handleConnect(s *Server, res Response, m Message) error {
-	return res.WriteJSON(map[string]string{
+func handleConnect(s *Server, ws *websocket.Conn, m Message) error {
+	return websocket.JSON.Send(ws, map[string]string{
 		"msg":     "connected",
 		"session": random.Id(17),
 	})
 }
 
-func handleMethod(s *Server, res Response, m Message) error {
+func handleMethod(s *Server, ws *websocket.Conn, m Message) error {
 	fn, ok := s.methods[m.Method]
 
 	if !ok {
 		err := errors.New(fmt.Sprintf("method %s not found", m.Method))
 		return err
 	}
 
-	ctx := NewMethodContext(m, res)
+	ctx := NewMethodContext(m, ws)
 	go fn(ctx)
 
 	return nil
 }
 
-func handlePing(s *Server, res Response, m Message) error {
+func handlePing(s *Server, ws *websocket.Conn, m Message) error {
 	msg := map[string]string{
 		"msg": "pong",
 	}
@@ -121,5 +95,5 @@ func handlePing(s *Server, res Response, m Message) error {
 		msg["id"] = m.ID
 	}
 
-	return res.WriteJSON(msg)
+	return websocket.JSON.Send(ws, msg)
 }