|
|
|
@ -1,10 +1,13 @@ |
|
|
|
package handlers |
|
|
|
package handlers |
|
|
|
|
|
|
|
|
|
|
|
import ( |
|
|
|
import ( |
|
|
|
|
|
|
|
"context" |
|
|
|
"fmt" |
|
|
|
"fmt" |
|
|
|
"html/template" |
|
|
|
"html/template" |
|
|
|
"net/http" |
|
|
|
"net/http" |
|
|
|
"sync" |
|
|
|
"sync" |
|
|
|
|
|
|
|
"sync/atomic" |
|
|
|
|
|
|
|
"time" |
|
|
|
|
|
|
|
|
|
|
|
"github.com/gorilla/websocket" |
|
|
|
"github.com/gorilla/websocket" |
|
|
|
"github.com/markbates/goth/gothic" |
|
|
|
"github.com/markbates/goth/gothic" |
|
|
|
@ -24,69 +27,178 @@ type PageData struct { |
|
|
|
Username string |
|
|
|
Username string |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
type Client struct { |
|
|
|
|
|
|
|
conn *websocket.Conn |
|
|
|
|
|
|
|
send chan []byte |
|
|
|
|
|
|
|
hub *hub |
|
|
|
|
|
|
|
ctx context.Context |
|
|
|
|
|
|
|
cancel context.CancelFunc |
|
|
|
|
|
|
|
id uint64 |
|
|
|
|
|
|
|
mu sync.Mutex // Protect concurrent operations
|
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
type hub struct { |
|
|
|
type hub struct { |
|
|
|
clients map[*websocket.Conn]bool |
|
|
|
clients sync.Map |
|
|
|
Broadcast chan []byte |
|
|
|
broadcast chan []byte |
|
|
|
register chan *websocket.Conn |
|
|
|
register chan *Client |
|
|
|
unregister chan *websocket.Conn |
|
|
|
unregister chan *Client |
|
|
|
mutex sync.RWMutex |
|
|
|
|
|
|
|
|
|
|
|
// Message pool for reusing byte slices
|
|
|
|
|
|
|
|
messagePool sync.Pool |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
var Hub = &hub{ |
|
|
|
var Hub = &hub{ |
|
|
|
clients: make(map[*websocket.Conn]bool), |
|
|
|
broadcast: make(chan []byte, 1024), |
|
|
|
Broadcast: make(chan []byte), |
|
|
|
register: make(chan *Client, 256), |
|
|
|
register: make(chan *websocket.Conn), |
|
|
|
unregister: make(chan *Client, 256), |
|
|
|
unregister: make(chan *websocket.Conn), |
|
|
|
|
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
var upgrader = websocket.Upgrader{ |
|
|
|
var upgrader = websocket.Upgrader{ |
|
|
|
|
|
|
|
ReadBufferSize: 4096, |
|
|
|
|
|
|
|
WriteBufferSize: 4096, |
|
|
|
CheckOrigin: func(r *http.Request) bool { |
|
|
|
CheckOrigin: func(r *http.Request) bool { |
|
|
|
return true // this should change
|
|
|
|
return true |
|
|
|
}, |
|
|
|
}, |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
var clientIDCounter uint64 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
const ( |
|
|
|
|
|
|
|
writeWait = 10 * time.Second |
|
|
|
|
|
|
|
pongWait = 60 * time.Second |
|
|
|
|
|
|
|
pingPeriod = 54 * time.Second |
|
|
|
|
|
|
|
maxMessageSize = 512 |
|
|
|
|
|
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
func init() { |
|
|
|
|
|
|
|
Hub.messagePool.New = func() interface{} { |
|
|
|
|
|
|
|
return make([]byte, 0, 1024) |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func (h *hub) Run() { |
|
|
|
func (h *hub) Run() { |
|
|
|
for { |
|
|
|
for { |
|
|
|
select { |
|
|
|
select { |
|
|
|
case conn := <-h.register: |
|
|
|
case client := <-h.register: |
|
|
|
h.mutex.Lock() |
|
|
|
h.clients.Store(client.id, client) |
|
|
|
h.clients[conn] = true |
|
|
|
fmt.Printf("Client connected. ID: %d\n", client.id) |
|
|
|
h.mutex.Unlock() |
|
|
|
|
|
|
|
fmt.Printf("Client connected. Total clients: %d\n", len(h.clients)) |
|
|
|
case client := <-h.unregister: |
|
|
|
|
|
|
|
if _, loaded := h.clients.LoadAndDelete(client.id); loaded { |
|
|
|
case conn := <-h.unregister: |
|
|
|
close(client.send) |
|
|
|
h.mutex.Lock() |
|
|
|
fmt.Printf("Client disconnected. ID: %d\n", client.id) |
|
|
|
if _, ok := h.clients[conn]; ok { |
|
|
|
|
|
|
|
delete(h.clients, conn) |
|
|
|
|
|
|
|
conn.Close() |
|
|
|
|
|
|
|
} |
|
|
|
} |
|
|
|
h.mutex.Unlock() |
|
|
|
|
|
|
|
fmt.Printf("Client disconnected. Total clients: %d\n", len(h.clients)) |
|
|
|
case message := <-h.broadcast: |
|
|
|
|
|
|
|
// Single-threaded broadcasting to avoid race conditions
|
|
|
|
case message := <-h.Broadcast: |
|
|
|
h.clients.Range(func(key, value interface{}) bool { |
|
|
|
h.mutex.RLock() |
|
|
|
client := value.(*Client) |
|
|
|
for conn := range h.clients { |
|
|
|
|
|
|
|
if err := conn.WriteMessage(websocket.TextMessage, message); err != nil { |
|
|
|
select { |
|
|
|
fmt.Printf("Error sending message to client: %v\n", err) |
|
|
|
case client.send <- message: |
|
|
|
// Remove failed connection
|
|
|
|
// Message sent successfully
|
|
|
|
delete(h.clients, conn) |
|
|
|
default: |
|
|
|
conn.Close() |
|
|
|
// Client is slow, remove it
|
|
|
|
|
|
|
|
h.clients.Delete(client.id) |
|
|
|
|
|
|
|
close(client.send) |
|
|
|
|
|
|
|
fmt.Printf("Slow client removed. ID: %d\n", client.id) |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
return true |
|
|
|
|
|
|
|
}) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
// Return message to pool after all clients processed
|
|
|
|
|
|
|
|
h.messagePool.Put(message) |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
func (c *Client) readPump() { |
|
|
|
|
|
|
|
defer func() { |
|
|
|
|
|
|
|
c.hub.unregister <- c |
|
|
|
|
|
|
|
c.conn.Close() |
|
|
|
|
|
|
|
c.cancel() |
|
|
|
|
|
|
|
}() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
c.conn.SetReadLimit(maxMessageSize) |
|
|
|
|
|
|
|
c.conn.SetReadDeadline(time.Now().Add(pongWait)) |
|
|
|
|
|
|
|
c.conn.SetPongHandler(func(string) error { |
|
|
|
|
|
|
|
c.conn.SetReadDeadline(time.Now().Add(pongWait)) |
|
|
|
|
|
|
|
return nil |
|
|
|
|
|
|
|
}) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
for { |
|
|
|
|
|
|
|
select { |
|
|
|
|
|
|
|
case <-c.ctx.Done(): |
|
|
|
|
|
|
|
return |
|
|
|
|
|
|
|
default: |
|
|
|
|
|
|
|
_, messageBytes, err := c.conn.ReadMessage() |
|
|
|
|
|
|
|
if err != nil { |
|
|
|
|
|
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { |
|
|
|
|
|
|
|
fmt.Printf("WebSocket error for client %d: %v\n", c.id, err) |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
return |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
// Get a message buffer from the pool
|
|
|
|
|
|
|
|
message := c.hub.messagePool.Get().([]byte) |
|
|
|
|
|
|
|
message = message[:len(messageBytes)] |
|
|
|
|
|
|
|
copy(message, messageBytes) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
// Non-blocking broadcast
|
|
|
|
|
|
|
|
select { |
|
|
|
|
|
|
|
case c.hub.broadcast <- message: |
|
|
|
|
|
|
|
// Message will be returned to pool by hub.Run()
|
|
|
|
|
|
|
|
default: |
|
|
|
|
|
|
|
// Broadcast buffer full, return message to pool and drop
|
|
|
|
|
|
|
|
c.hub.messagePool.Put(message) |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
func (c *Client) writePump() { |
|
|
|
|
|
|
|
ticker := time.NewTicker(pingPeriod) |
|
|
|
|
|
|
|
defer func() { |
|
|
|
|
|
|
|
ticker.Stop() |
|
|
|
|
|
|
|
c.conn.Close() |
|
|
|
|
|
|
|
}() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
for { |
|
|
|
|
|
|
|
select { |
|
|
|
|
|
|
|
case <-c.ctx.Done(): |
|
|
|
|
|
|
|
return |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
case message, ok := <-c.send: |
|
|
|
|
|
|
|
c.conn.SetWriteDeadline(time.Now().Add(writeWait)) |
|
|
|
|
|
|
|
if !ok { |
|
|
|
|
|
|
|
c.conn.WriteMessage(websocket.CloseMessage, []byte{}) |
|
|
|
|
|
|
|
return |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
// Write the message
|
|
|
|
|
|
|
|
if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil { |
|
|
|
|
|
|
|
fmt.Printf("Write error for client %d: %v\n", c.id, err) |
|
|
|
|
|
|
|
return |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
case <-ticker.C: |
|
|
|
|
|
|
|
c.conn.SetWriteDeadline(time.Now().Add(writeWait)) |
|
|
|
|
|
|
|
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { |
|
|
|
|
|
|
|
return |
|
|
|
} |
|
|
|
} |
|
|
|
h.mutex.RUnlock() |
|
|
|
|
|
|
|
} |
|
|
|
} |
|
|
|
} |
|
|
|
} |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func (h *Handler) Home(w http.ResponseWriter, r *http.Request) { |
|
|
|
func (h *Handler) Home(w http.ResponseWriter, r *http.Request) { |
|
|
|
session, err := gothic.Store.Get(r, "user-session") |
|
|
|
session, err := gothic.Store.Get(r, "user-session") |
|
|
|
|
|
|
|
|
|
|
|
if err != nil { |
|
|
|
if err != nil { |
|
|
|
http.Error(w, "Error retrieving session for welcome page", http.StatusInternalServerError) |
|
|
|
http.Error(w, "Error retrieving session for welcome page", http.StatusInternalServerError) |
|
|
|
|
|
|
|
return |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
username, ok := session.Values["user_name"].(string) |
|
|
|
username, ok := session.Values["user_name"].(string) |
|
|
|
var pagedata PageData |
|
|
|
var pagedata PageData |
|
|
|
|
|
|
|
|
|
|
|
if ok { |
|
|
|
if ok { |
|
|
|
pagedata.Username = username |
|
|
|
pagedata.Username = username |
|
|
|
} else { |
|
|
|
} else { |
|
|
|
@ -106,21 +218,22 @@ func (h *Handler) WsHandler(w http.ResponseWriter, r *http.Request) { |
|
|
|
return |
|
|
|
return |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
Hub.register <- conn |
|
|
|
ctx, cancel := context.WithCancel(context.Background()) |
|
|
|
|
|
|
|
client := &Client{ |
|
|
|
|
|
|
|
conn: conn, |
|
|
|
|
|
|
|
send: make(chan []byte, 256), |
|
|
|
|
|
|
|
hub: Hub, |
|
|
|
|
|
|
|
ctx: ctx, |
|
|
|
|
|
|
|
cancel: cancel, |
|
|
|
|
|
|
|
id: atomic.AddUint64(&clientIDCounter, 1), |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
defer conn.Close() |
|
|
|
Hub.register <- client |
|
|
|
|
|
|
|
|
|
|
|
for { |
|
|
|
// Start pumps in separate goroutines
|
|
|
|
_, messageBytes, err := conn.ReadMessage() |
|
|
|
go client.writePump() |
|
|
|
if err != nil && !websocket.IsCloseError( |
|
|
|
go client.readPump() |
|
|
|
err, |
|
|
|
|
|
|
|
websocket.CloseNormalClosure, |
|
|
|
|
|
|
|
websocket.CloseGoingAway, |
|
|
|
|
|
|
|
websocket.CloseNoStatusReceived, |
|
|
|
|
|
|
|
) { |
|
|
|
|
|
|
|
fmt.Printf("error reading message: %v", err) |
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
Hub.Broadcast <- messageBytes |
|
|
|
// Wait for completion
|
|
|
|
} |
|
|
|
<-ctx.Done() |
|
|
|
} |
|
|
|
} |
|
|
|
|