109 lines
1.9 KiB
Go
109 lines
1.9 KiB
Go
|
package main
|
||
|
|
||
|
import (
|
||
|
"bytes"
|
||
|
"log"
|
||
|
"net/http"
|
||
|
"time"
|
||
|
|
||
|
"github.com/gin-gonic/gin"
|
||
|
"github.com/gorilla/websocket"
|
||
|
)
|
||
|
|
||
|
const (
|
||
|
// Time allowed to write message to peer
|
||
|
writeWait = 10 * time.Second
|
||
|
|
||
|
// Time allowed to read the next pong message from the client.
|
||
|
pongWait = 5 * time.Second
|
||
|
|
||
|
// Send pings to client with this period. Must be less than pongWait.
|
||
|
pingPeriod = (pongWait * 9) / 10
|
||
|
//pingPeriod = 5 * time.Second
|
||
|
|
||
|
// Maximum message size
|
||
|
maxMessageSize = 512
|
||
|
)
|
||
|
|
||
|
var upgrader = websocket.Upgrader{
|
||
|
ReadBufferSize: 1024,
|
||
|
WriteBufferSize: 1024,
|
||
|
CheckOrigin: func(r *http.Request) bool {
|
||
|
return true
|
||
|
},
|
||
|
}
|
||
|
|
||
|
var (
|
||
|
conn *websocket.Conn
|
||
|
)
|
||
|
|
||
|
var (
|
||
|
newline = []byte{'\n'}
|
||
|
space = []byte{' '}
|
||
|
)
|
||
|
|
||
|
func reader() {
|
||
|
defer conn.Close()
|
||
|
conn.SetReadLimit(maxMessageSize)
|
||
|
conn.SetReadDeadline(time.Now().Add(pongWait))
|
||
|
conn.SetPongHandler(func(string) error {
|
||
|
//log.Println("pong")
|
||
|
conn.SetReadDeadline(time.Now().Add(pongWait))
|
||
|
return nil
|
||
|
})
|
||
|
|
||
|
for {
|
||
|
_, message, err := conn.ReadMessage()
|
||
|
if err != nil {
|
||
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway,
|
||
|
websocket.CloseAbnormalClosure) {
|
||
|
log.Printf("error: %v", err)
|
||
|
}
|
||
|
break
|
||
|
}
|
||
|
message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func writer() {
|
||
|
pingTicker := time.NewTicker(pingPeriod)
|
||
|
defer func() {
|
||
|
log.Println("quit ws server")
|
||
|
|
||
|
pingTicker.Stop()
|
||
|
conn.Close()
|
||
|
}()
|
||
|
|
||
|
for {
|
||
|
select {
|
||
|
case <-pingTicker.C:
|
||
|
//log.Println("ping")
|
||
|
|
||
|
conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||
|
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||
|
log.Println(err)
|
||
|
return
|
||
|
}
|
||
|
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func serveWebsocket(c *gin.Context) {
|
||
|
var err error
|
||
|
log.Println("websocket request")
|
||
|
|
||
|
conn, err = upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
|
if err != nil {
|
||
|
if _, ok := err.(websocket.HandshakeError); !ok {
|
||
|
log.Printf("handshake error: %s", err)
|
||
|
}
|
||
|
|
||
|
log.Println(err)
|
||
|
return
|
||
|
}
|
||
|
|
||
|
go writer()
|
||
|
reader()
|
||
|
}
|