mirror of https://github.com/schollz/croc.git
117 lines
2.6 KiB
Go
117 lines
2.6 KiB
Go
package relay
|
|
|
|
import (
|
|
"net/http"
|
|
"time"
|
|
|
|
log "github.com/cihub/seelog"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
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 = 6000 * 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 = 1024 * 1024
|
|
)
|
|
|
|
var upgrader = websocket.Upgrader{
|
|
ReadBufferSize: 1024 * 1024,
|
|
WriteBufferSize: 1024 * 1024,
|
|
}
|
|
|
|
// connection is an middleman between the websocket connection and the hub.
|
|
type connection struct {
|
|
// The websocket connection.
|
|
ws *websocket.Conn
|
|
|
|
// Buffered channel of outbound messages.
|
|
send chan messageChannel
|
|
}
|
|
|
|
type messageChannel struct {
|
|
data []byte
|
|
messageType int
|
|
}
|
|
|
|
// readPump pumps messages from the websocket connection to the hub.
|
|
func (s subscription) readPump() {
|
|
c := s.conn
|
|
defer func() {
|
|
h.unregister <- s
|
|
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 {
|
|
messageType, msg, err := c.ws.ReadMessage()
|
|
if err != nil {
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
|
|
log.Errorf("error: %v", err)
|
|
}
|
|
break
|
|
}
|
|
h.broadcast <- message{messageChannel{msg, messageType}, s.room, c.ws.RemoteAddr().String()}
|
|
}
|
|
}
|
|
|
|
// 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 (s *subscription) writePump() {
|
|
c := s.conn
|
|
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(message.messageType, message.data); err != nil {
|
|
return
|
|
}
|
|
case <-ticker.C:
|
|
if err := c.write(websocket.PingMessage, []byte{}); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// serveWs handles websocket requests from the peer.
|
|
func serveWs(w http.ResponseWriter, r *http.Request) {
|
|
ws, err := upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
log.ErrorStr(err)
|
|
}
|
|
vals := r.URL.Query()
|
|
room := "default"
|
|
rooms, ok := vals["room"]
|
|
if ok {
|
|
room = rooms[0]
|
|
}
|
|
|
|
c := &connection{send: make(chan messageChannel, 256), ws: ws}
|
|
s := subscription{c, room}
|
|
h.register <- s
|
|
go s.writePump()
|
|
s.readPump()
|
|
}
|