websocket.go

package iris

import (
	irisWebsocket "github.com/iris-contrib/websocket"
	"github.com/kataras/go-websocket"
)

// conversionals
const (
	// All is the string which the Emmiter use to send a message to all
	All = websocket.All
	// NotMe is the string which the Emmiter use to send a message to all except this websocket.Connection
	NotMe = websocket.NotMe
	// Broadcast is the string which the Emmiter use to send a message to all except this websocket.Connection, same as 'NotMe'
	Broadcast = websocket.Broadcast
)

// Note I keep this code only to no change the front-end API, we could only use the go-websocket and set our custom upgrader

type (
	// WebsocketServer is the iris websocket server, expose the websocket.Server
	// the below code is a wrapper and bridge between iris-contrib/websocket and kataras/go-websocket
	WebsocketServer struct {
		websocket.Server
		upgrader irisWebsocket.Upgrader

		// the only fields we need at runtime here for iris-specific error and check origin funcs
		// they comes from WebsocketConfiguration

		// Error specifies the function for generating HTTP error responses.
		Error func(ctx *Context, status int, reason string)
		// CheckOrigin returns true if the request Origin header is acceptable. If
		// CheckOrigin is nil, the host in the Origin header must not be set or
		// must match the host of the request.
		CheckOrigin func(ctx *Context) bool
	}
)

// NewWebsocketServer returns an empty WebsocketServer, nothing special here.
func NewWebsocketServer() *WebsocketServer {
	return &WebsocketServer{}
}

// Upgrade upgrades the HTTP server connection to the WebSocket protocol.
//
// The responseHeader is included in the response to the client's upgrade
// request. Use the responseHeader to specify cookies (Set-Cookie) and the
// application negotiated subprotocol (Sec-Websocket-Protocol).
//
// If the upgrade fails, then Upgrade replies to the client with an HTTP error
// response.
func (s *WebsocketServer) Upgrade(ctx *Context) error {
	return s.upgrader.Upgrade(ctx.RequestCtx)
}

// Handler is the iris Handler to upgrade the request
// used inside RegisterRoutes
func (s *WebsocketServer) Handler(ctx *Context) {
	// first, check origin
	if !s.CheckOrigin(ctx) {
		s.Error(ctx, StatusForbidden, "websocket: origin not allowed")
		return
	}

	// all other errors comes from the underline iris-contrib/websocket
	if err := s.Upgrade(ctx); err != nil {
		if ctx.framework.Config.IsDevelopment {
			ctx.Log("Websocket error while trying to Upgrade the connection. Trace: %s", err.Error())
		}

		statusErrCode := StatusBadRequest
		if herr, isHandshake := err.(irisWebsocket.HandshakeError); isHandshake {
			statusErrCode = herr.Status()
		}
		// if not handshake error just fire the custom(if any) StatusBadRequest
		// with the websocket's error message in the ctx.Get("WsError")
		DefaultWebsocketError(ctx, statusErrCode, err.Error())

	}
}

// RegisterTo creates the client side source route and the route path Endpoint with the correct Handler
// receives the websocket configuration and  the iris station
func (s *WebsocketServer) RegisterTo(station *Framework, c WebsocketConfiguration) {

	// Note: s.Server should be initialize on the first OnConnection, which is called before this func  when Default websocket server.
	// When not: when calling this function before OnConnection, when we have more than one websocket server running
	if s.Server == nil {
		s.Server = websocket.New()
	}
	// is just a conversional type for kataras/go-websocket.Connection
	s.upgrader = irisWebsocket.Custom(s.Server.HandleConnection, c.ReadBufferSize, c.WriteBufferSize, c.Headers)

	// set the routing for client-side source (javascript) (optional)
	clientSideLookupName := "iris-websocket-client-side"
	station.Get(c.Endpoint, s.Handler)
	// check if client side already exists
	if station.Lookup(clientSideLookupName) == nil {
		// serve the client side on domain:port/iris-ws.js
		station.StaticContent("/iris-ws.js", contentJavascript, websocket.ClientSource)(clientSideLookupName)
	}

	s.Server.Set(websocket.Config{
		WriteTimeout:    c.WriteTimeout,
		PongTimeout:     c.PongTimeout,
		PingPeriod:      c.PingPeriod,
		MaxMessageSize:  c.MaxMessageSize,
		BinaryMessages:  c.BinaryMessages,
		ReadBufferSize:  c.ReadBufferSize,
		WriteBufferSize: c.WriteBufferSize,
	})

	s.Error = c.Error
	s.CheckOrigin = c.CheckOrigin

	if s.Error == nil {
		s.Error = DefaultWebsocketError
	}

	if s.CheckOrigin == nil {
		s.CheckOrigin = DefaultWebsocketCheckOrigin
	}

	// run the ws server
	s.Server.Serve()

}

// WebsocketConnection is the front-end API that you will use to communicate with the client side
type WebsocketConnection interface {
	websocket.Connection
}

// OnConnection this is the main event you, as developer, will work with each of the websocket connections
func (s *WebsocketServer) OnConnection(connectionListener func(WebsocketConnection)) {
	if s.Server == nil {
		// for default webserver this is the time when the websocket server will be init
		// let's initialize here the ws server, the user/dev is free to change its config before this step.
		s.Server = websocket.New() // we need that in order to use the Iris' WebsocketConnnection, which
		// config is empty here because are setted on the RegisterTo
		// websocket's configuration is optional on New because it doesn't really used before the websocket.Serve
	}
	s.Server.OnConnection(func(c websocket.Connection) {
		connectionListener(c)
	})
}