Semaphore/api/sockets/handler.go
2016-05-24 10:55:48 +01:00

123 lines
2.4 KiB
Go

package sockets
import (
"fmt"
"log"
"time"
"github.com/ansible-semaphore/semaphore/models"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
const (
// Time allowed to write a message to the peer.
writeWait = 10 * time.Second
// Time allowed to read the next pong message from the peer.
pongWait = 60 * time.Second
// Send pings to peer with this period. Must be less than pongWait.
pingPeriod = (pongWait * 9) / 10
// Maximum message size allowed from peer.
maxMessageSize = 512
)
type connection struct {
ws *websocket.Conn
send chan []byte
userID int
}
// readPump pumps messages from the websocket connection to the hub.
func (c *connection) readPump() {
defer func() {
h.unregister <- c
c.ws.Close()
}()
c.ws.SetReadLimit(maxMessageSize)
c.ws.SetReadDeadline(time.Now().Add(pongWait))
c.ws.SetPongHandler(func(string) error {
c.ws.SetReadDeadline(time.Now().Add(pongWait))
return nil
})
for {
_, message, err := c.ws.ReadMessage()
fmt.Println(string(message))
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
log.Printf("error: %v", err)
}
break
}
}
}
// write writes a message with the given message type and payload.
func (c *connection) write(mt int, payload []byte) error {
c.ws.SetWriteDeadline(time.Now().Add(writeWait))
return c.ws.WriteMessage(mt, payload)
}
// writePump pumps messages from the hub to the websocket connection.
func (c *connection) writePump() {
ticker := time.NewTicker(pingPeriod)
defer func() {
ticker.Stop()
c.ws.Close()
}()
for {
select {
case message, ok := <-c.send:
if !ok {
c.write(websocket.CloseMessage, []byte{})
return
}
if err := c.write(websocket.TextMessage, message); err != nil {
return
}
case <-ticker.C:
if err := c.write(websocket.PingMessage, []byte{}); err != nil {
return
}
}
}
}
func Handler(context *gin.Context) {
user := context.MustGet("user").(*models.User)
ws, err := upgrader.Upgrade(context.Writer, context.Request, nil)
if err != nil {
panic(err)
}
c := &connection{
send: make(chan []byte, 256),
ws: ws,
userID: user.ID,
}
h.register <- c
go c.writePump()
c.readPump()
}
func Message(userID int, message []byte) {
h.broadcast <- &sendRequest{
userID: userID,
msg: message,
}
}