sacrificial-socket/socket.go

package ss

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"github.com/gorilla/websocket"
	"github.com/raz-varren/log"
	"sync"
)

var (
	socketRNG = NewRNG()
)

//Socket represents a websocket connection
type Socket struct {
	l      *sync.RWMutex
	id     string
	ws     *websocket.Conn
	closed bool
	serv   *SocketServer
	roomsl *sync.RWMutex
	rooms  map[string]bool
}

const (
	idLen int = 24

	typeJSON string = "J"
	typeBin         = "B"
	typeStr         = "S"
)

func newSocket(serv *SocketServer, ws *websocket.Conn) *Socket {
	s := &Socket{
		l:      &sync.RWMutex{},
		id:     newSocketID(),
		ws:     ws,
		closed: false,
		serv:   serv,
		roomsl: &sync.RWMutex{},
		rooms:  make(map[string]bool),
	}
	serv.hub.addSocket(s)
	return s
}

func newSocketID() string {
	idBuf := make([]byte, idLen)
	socketRNG.Read(idBuf)
	return base64.StdEncoding.EncodeToString(idBuf)
}

func (s *Socket) receive() ([]byte, error) {
	_, data, err := s.ws.ReadMessage()
	return data, err
}

func (s *Socket) send(msgType int, data []byte) error {
	s.l.Lock()
	defer s.l.Unlock()
	return s.ws.WriteMessage(msgType, data)
}

//InRoom returns true if s is currently a member of roomName
func (s *Socket) InRoom(roomName string) bool {
	s.roomsl.RLock()
	defer s.roomsl.RUnlock()
	inRoom := s.rooms[roomName]
	return inRoom
}

//GetRooms returns a list of rooms that s is a member of
func (s *Socket) GetRooms() []string {
	s.roomsl.RLock()
	defer s.roomsl.RUnlock()

	var roomList []string
	for room := range s.rooms {
		roomList = append(roomList, room)
	}
	return roomList
}

//Join adds s to the specified room. If the room does
//not exist, it will be created
func (s *Socket) Join(roomName string) {
	s.roomsl.Lock()
	defer s.roomsl.Unlock()
	s.serv.hub.joinRoom(&joinRequest{roomName, s})
	s.rooms[roomName] = true
}

//Leave removes s from the specified room. If s
//is not a member of the room, nothing will happen. If the room is
//empty upon removal of s, the room will be closed
func (s *Socket) Leave(roomName string) {
	s.roomsl.Lock()
	defer s.roomsl.Unlock()
	s.serv.hub.leaveRoom(&leaveRequest{roomName, s})
	delete(s.rooms, roomName)
}

//Roomcast dispatches an event to all Sockets in the specified room.
func (s *Socket) Roomcast(roomName, eventName string, data interface{}) {
	s.serv.hub.roomcast(&RoomMsg{roomName, eventName, data})
}

//Broadcast dispatches an event to all Sockets on the SocketServer.
func (s *Socket) Broadcast(eventName string, data interface{}) {
	s.serv.hub.broadcast(&BroadcastMsg{eventName, data})
}

//Socketcast dispatches an event to the specified socket ID.
func (s *Socket) Socketcast(socketID, eventName string, data interface{}) {
	s.serv.Roomcast("__socket_id:"+socketID, eventName, data)
}

//Emit dispatches an event to s.
func (s *Socket) Emit(eventName string, data interface{}) error {
	d, msgType := emitData(eventName, data)
	return s.send(msgType, d)
}

//ID returns the unique ID of s
func (s *Socket) ID() string {
	return s.id
}

//emitData combines the eventName and data into a payload that is understood
//by the sac-sock protocol.
func emitData(eventName string, data interface{}) ([]byte, int) {
	buf := bytes.NewBuffer(nil)
	buf.WriteString(eventName)
	buf.WriteByte(startOfHeaderByte)

	switch d := data.(type) {
	case string:
		buf.WriteString(typeStr)
		buf.WriteByte(startOfDataByte)
		buf.WriteString(d)
		return buf.Bytes(), websocket.TextMessage

	case []byte:
		buf.WriteString(typeBin)
		buf.WriteByte(startOfDataByte)
		buf.Write(d)
		return buf.Bytes(), websocket.BinaryMessage

	default:
		buf.WriteString(typeJSON)
		buf.WriteByte(startOfDataByte)
		jsonData, err := json.Marshal(d)
		if err != nil {
			log.Err.Println(err)
		} else {
			buf.Write(jsonData)
		}
		return buf.Bytes(), websocket.TextMessage
	}
}

//Close closes the Socket connection and removes the Socket
//from any rooms that it was a member of
func (s *Socket) Close() {
	s.l.Lock()
	isAlreadyClosed := s.closed
	s.closed = true
	s.l.Unlock()

	if isAlreadyClosed { //can't reclose the socket
		return
	}

	defer log.Debug.Println(s.ID(), "disconnected")

	s.ws.Close()

	rooms := s.GetRooms()

	for _, room := range rooms {
		s.Leave(room)
	}

	s.serv.l.RLock()
	event := s.serv.onDisconnectFunc
	s.serv.l.RUnlock()

	if event != nil {
		event(s)
	}

	s.serv.hub.removeSocket(s)
}