iris/websocket.go

package iris

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

// ---------------------------------------------------------------------------------------------------------
// ---------------------------------------------------------------------------------------------------------
// --------------------------------Websocket implementation-------------------------------------------------
// Global functions in order to be able to use unlimitted number of websocket servers on each iris station--
// ---------------------------------------------------------------------------------------------------------

// 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

// NewWebsocketServer creates a websocket server and returns it
func NewWebsocketServer(c *config.Websocket) *WebsocketServer {
	wsConfig := websocket.Config{
		WriteTimeout:    c.WriteTimeout,
		PongTimeout:     c.PongTimeout,
		PingPeriod:      c.PingPeriod,
		MaxMessageSize:  c.MaxMessageSize,
		BinaryMessages:  c.BinaryMessages,
		ReadBufferSize:  c.ReadBufferSize,
		WriteBufferSize: c.WriteBufferSize,
	}

	wsServer := websocket.New(wsConfig)

	upgrader := irisWebsocket.Custom(wsServer.HandleConnection, c.ReadBufferSize, c.WriteBufferSize, false)

	srv := &WebsocketServer{Server: wsServer, Config: c, upgrader: upgrader}

	return srv
}

// RegisterWebsocketServer registers the handlers for the websocket server
// it's a bridge between station and websocket server
func RegisterWebsocketServer(station FrameworkAPI, server *WebsocketServer, logger *logger.Logger) {
	c := server.Config
	if c.Endpoint == "" {
		return
	}

	websocketHandler := func(ctx *Context) {

		if err := server.Upgrade(ctx); err != nil {
			if ctx.framework.Config.IsDevelopment {
				logger.Printf("Websocket error while trying to Upgrade the connection. Trace: %s", err.Error())
			}
			ctx.EmitError(StatusBadRequest)
		}
	}

	if c.Headers != nil && len(c.Headers) > 0 { // only for performance matter just re-create the websocketHandler if we have headers to set
		websocketHandler = func(ctx *Context) {
			for k, v := range c.Headers {
				ctx.SetHeader(k, v)
			}

			if err := server.Upgrade(ctx); err != nil {
				if ctx.framework.Config.IsDevelopment {
					logger.Printf("Websocket error while trying to Upgrade the connection. Trace: %s", err.Error())
				}
				ctx.EmitError(StatusBadRequest)
			}
		}
	}
	clientSideLookupName := "iris-websocket-client-side"
	station.Get(c.Endpoint, websocketHandler)
	// 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)
	}

	// run the ws server
	server.Serve()
}

// 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
)

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
		Config   *config.Websocket
		upgrader irisWebsocket.Upgrader
	}
)

// 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)
}

// 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)) {
	s.Server.OnConnection(func(c websocket.Connection) {
		connectionListener(c)
	})
}