1+ package socket
2+
3+ import (
4+ "bytes"
5+ "log"
6+ "net/http"
7+ "time"
8+
9+ "github.com/gorilla/websocket"
10+ )
11+
12+ const (
13+ // Time allowed to write a message to the peer.
14+ writeWait = 10 * time .Second
15+
16+ // Time allowed to read the next pong message from the peer.
17+ pongWait = 60 * time .Second
18+
19+ // Send pings to peer with this period. Must be less than pongWait.
20+ pingPeriod = (pongWait * 9 ) / 10
21+
22+ // Maximum message size allowed from peer.
23+ maxMessageSize = 512
24+ )
25+
26+ var (
27+ newline = []byte {'\n' }
28+ space = []byte {' ' }
29+ )
30+
31+ var upgrader = websocket.Upgrader {
32+ ReadBufferSize : 1024 ,
33+ WriteBufferSize : 1024 ,
34+ }
35+
36+ // Client is a middleman between the websocket connection and the hub.
37+ type Client struct {
38+ hub * Hub
39+
40+ // The websocket connection.
41+ conn * websocket.Conn
42+
43+ // Buffered channel of outbound messages.
44+ send chan []byte
45+ }
46+
47+ // readPump pumps messages from the websocket connection to the hub.
48+ //
49+ // The application runs readPump in a per-connection goroutine. The application
50+ // ensures that there is at most one reader on a connection by executing all
51+ // reads from this goroutine.
52+ func (c * Client ) readPump () {
53+ defer func () {
54+ c .hub .unregister <- c
55+ c .conn .Close ()
56+ }()
57+ c .conn .SetReadLimit (maxMessageSize )
58+ c .conn .SetReadDeadline (time .Now ().Add (pongWait ))
59+ c .conn .SetPongHandler (func (string ) error { c .conn .SetReadDeadline (time .Now ().Add (pongWait )); return nil })
60+ for {
61+ _ , message , err := c .conn .ReadMessage ()
62+ if err != nil {
63+ if websocket .IsUnexpectedCloseError (err , websocket .CloseGoingAway , websocket .CloseAbnormalClosure ) {
64+ log .Printf ("error: %v" , err )
65+ }
66+ break
67+ }
68+ message = bytes .TrimSpace (bytes .Replace (message , newline , space , - 1 ))
69+ c .hub .Broadcast <- message
70+ }
71+ }
72+
73+ // writePump pumps messages from the hub to the websocket connection.
74+ //
75+ // A goroutine running writePump is started for each connection. The
76+ // application ensures that there is at most one writer to a connection by
77+ // executing all writes from this goroutine.
78+ func (c * Client ) writePump () {
79+ ticker := time .NewTicker (pingPeriod )
80+ defer func () {
81+ ticker .Stop ()
82+ c .conn .Close ()
83+ }()
84+ for {
85+ select {
86+ case message , ok := <- c .send :
87+ c .conn .SetWriteDeadline (time .Now ().Add (writeWait ))
88+ if ! ok {
89+ // The hub closed the channel.
90+ c .conn .WriteMessage (websocket .CloseMessage , []byte {})
91+ return
92+ }
93+
94+ w , err := c .conn .NextWriter (websocket .TextMessage )
95+ if err != nil {
96+ return
97+ }
98+ w .Write (message )
99+
100+ // Add queued chat messages to the current websocket message.
101+ n := len (c .send )
102+ for i := 0 ; i < n ; i ++ {
103+ w .Write (newline )
104+ w .Write (<- c .send )
105+ }
106+
107+ if err := w .Close (); err != nil {
108+ return
109+ }
110+ case <- ticker .C :
111+ c .conn .SetWriteDeadline (time .Now ().Add (writeWait ))
112+ if err := c .conn .WriteMessage (websocket .PingMessage , nil ); err != nil {
113+ return
114+ }
115+ }
116+ }
117+ }
118+
119+ // ServeSocket handles a websocket request from a client and registers the client with the socket hub
120+ func ServeSocket (hub * Hub , w http.ResponseWriter , r * http.Request ) {
121+ conn , err := upgrader .Upgrade (w , r , nil )
122+ if err != nil {
123+ log .Println (err )
124+ return
125+ }
126+ client := & Client {hub : hub , conn : conn , send : make (chan []byte , 256 )}
127+ client .hub .register <- client
128+
129+ // Allow collection of memory referenced by the caller by doing all work in
130+ // new goroutines.
131+ go client .writePump ()
132+ go client .readPump ()
133+ }
0 commit comments