overflow_service_websocket/websocket/config.go

package websocket

import (
	"net/http"
	"time"

	uuid "github.com/satori/go.uuid"
)

const (
	// DefaultWriteTimeout is default value of Write Timeout
	DefaultWriteTimeout = 0
	// DefaultReadTimeout is default value of Read Timeout
	DefaultReadTimeout = 0
	// DefaultPongTimeout is default value of Pong Timeout
	DefaultPongTimeout = 60 * time.Second
	// DefaultPingPeriod is default value of Ping Period
	DefaultPingPeriod = (DefaultPongTimeout * 9) / 10
	// DefaultMaxMessageSize is default value of Max Message Size
	DefaultMaxMessageSize = 1024
	// DefaultReadBufferSize is default value of Read Buffer Size
	DefaultReadBufferSize = 4096
	// DefaultWriteBufferSize is default value of Write Buffer Size
	DefaultWriteBufferSize = 4096
)

var (
	// DefaultIDGenerator returns the UUID of the client
	DefaultIDGenerator = func(*http.Request) string { return uuid.NewV4().String() }
)

type (
	// OptionSetter sets a configuration field to the websocket config
	// used to help developers to write less and configure only what they really want and nothing else
	OptionSetter interface {
		Set(c *Config)
	}

	// OptionSet implements the OptionSetter
	OptionSet func(c *Config)
)

// Set is the func which makes the OptionSet an OptionSetter, this is used mostly
func (o OptionSet) Set(c *Config) {
	o(c)
}

// Config is configuration of the websocket server
type Config struct {
	Error           func(res http.ResponseWriter, req *http.Request, status int, reason error)
	CheckOrigin     func(req *http.Request) bool
	WriteTimeout    time.Duration
	ReadTimeout     time.Duration
	PongTimeout     time.Duration
	PingPeriod      time.Duration
	MaxMessageSize  int64
	BinaryMessages  bool
	ReadBufferSize  int
	WriteBufferSize int
	IDGenerator     func(*http.Request) string
}

// Set is the func which makes the OptionSet an OptionSetter, this is used mostly
func (c Config) Set(main *Config) {
	main.Error = c.Error
	main.CheckOrigin = c.CheckOrigin
	main.WriteTimeout = c.WriteTimeout
	main.ReadTimeout = c.ReadTimeout
	main.PongTimeout = c.PongTimeout
	main.PingPeriod = c.PingPeriod
	main.MaxMessageSize = c.MaxMessageSize
	main.BinaryMessages = c.BinaryMessages
	main.ReadBufferSize = c.ReadBufferSize
	main.WriteBufferSize = c.WriteBufferSize
	main.IDGenerator = c.IDGenerator
}

// Error sets the error handler
func Error(val func(res http.ResponseWriter, req *http.Request, status int, reason error)) OptionSet {
	return func(c *Config) {
		c.Error = val
	}
}

// CheckOrigin sets a handler which will check if different origin(domains) are allowed to contact with
// the websocket server
func CheckOrigin(val func(req *http.Request) bool) OptionSet {
	return func(c *Config) {
		c.CheckOrigin = val
	}
}

// WriteTimeout time allowed to write a message to the connection.
// Default value is 15 * time.Second
func WriteTimeout(val time.Duration) OptionSet {
	return func(c *Config) {
		c.WriteTimeout = val
	}
}

// ReadTimeout time allowed to read a message from the connection.
// Default value is 15 * time.Second
func ReadTimeout(val time.Duration) OptionSet {
	return func(c *Config) {
		c.ReadTimeout = val
	}
}

// PongTimeout allowed to read the next pong message from the connection
// Default value is 60 * time.Second
func PongTimeout(val time.Duration) OptionSet {
	return func(c *Config) {
		c.PongTimeout = val
	}
}

// PingPeriod send ping messages to the connection with this period. Must be less than PongTimeout
// Default value is (PongTimeout * 9) / 10
func PingPeriod(val time.Duration) OptionSet {
	return func(c *Config) {
		c.PingPeriod = val
	}
}

// MaxMessageSize max message size allowed from connection
// Default value is 1024
func MaxMessageSize(val int64) OptionSet {
	return func(c *Config) {
		c.MaxMessageSize = val
	}
}

// BinaryMessages set it to true in order to denotes binary data messages instead of utf-8 text
// compatible if you wanna use the Connection's EmitMessage to send a custom binary data to the client,
// like a native server-client communication.
// defaults to false
func BinaryMessages(val bool) OptionSet {
	return func(c *Config) {
		c.BinaryMessages = val
	}
}

// ReadBufferSize is the buffer size for the underline reader
func ReadBufferSize(val int) OptionSet {
	return func(c *Config) {
		c.ReadBufferSize = val
	}
}

// WriteBufferSize is the buffer size for the underline writer
func WriteBufferSize(val int) OptionSet {
	return func(c *Config) {
		c.WriteBufferSize = val
	}
}

// IDGenerator used to create (and later on, set)
// an ID for each incoming websocket connections (clients).
// The request is an argument which you can use to generate the ID (from headers for example).
// If empty then the ID is generated by func: uuid.NewV4().String()
func IDGenerator(val func(*http.Request) string) OptionSet {
	return func(c *Config) {
		c.IDGenerator = val
	}
}

// Validate validates the configuration
func (c Config) Validate() Config {

	if c.WriteTimeout < 0 {
		c.WriteTimeout = DefaultWriteTimeout
	}

	if c.ReadTimeout < 0 {
		c.ReadTimeout = DefaultReadTimeout
	}

	if c.PongTimeout < 0 {
		c.PongTimeout = DefaultPongTimeout
	}

	if c.PingPeriod <= 0 {
		c.PingPeriod = DefaultPingPeriod
	}

	if c.MaxMessageSize <= 0 {
		c.MaxMessageSize = DefaultMaxMessageSize
	}

	if c.ReadBufferSize <= 0 {
		c.ReadBufferSize = DefaultReadBufferSize
	}

	if c.WriteBufferSize <= 0 {
		c.WriteBufferSize = DefaultWriteBufferSize
	}

	if c.Error == nil {
		c.Error = func(res http.ResponseWriter, req *http.Request, status int, reason error) {
		}
	}

	if c.CheckOrigin == nil {
		c.CheckOrigin = func(req *http.Request) bool {
			return true
		}
	}

	if c.IDGenerator == nil {
		c.IDGenerator = DefaultIDGenerator
	}

	return c
}