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, } }