udemy-go-web-2/cmd/web/ws-handlers.go

110 lines
2.3 KiB
Go

package main
import (
"fmt"
"net/http"
"github.com/gorilla/websocket"
)
type WebSocketConnection struct {
*websocket.Conn
}
type WsPayload struct {
Action string `json:"action"`
Message string `json:"message"`
UserName string `json:"user_name"`
MessageType string `json:"message_type"`
UserID int `json:"user_id"`
Conn WebSocketConnection `json:"-"`
}
type WsJSONResponse struct {
Action string `json:"action"`
Message string `json:"message"`
UserID int `json:"user_id"`
}
var upgradeConnection = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
// since we don't expect any communication from the client side
CheckOrigin: func(r *http.Request) bool { return true },
}
var clients = make(map[WebSocketConnection]string)
var wsChan = make(chan WsPayload)
func (app *application) WsEndPoint(w http.ResponseWriter, r *http.Request) {
ws, err := upgradeConnection.Upgrade(w, r, nil)
if err != nil {
app.errorLog.Println(err)
return
}
app.infoLog.Printf("Client connected from %s\n", r.RemoteAddr)
var response WsJSONResponse
response.Message = "Connected to server"
err = ws.WriteJSON(response)
if err != nil {
app.errorLog.Println(err)
return
}
conn := WebSocketConnection{Conn: ws}
clients[conn] = ""
go app.ListenForWS(&conn)
}
func (app *application) ListenForWS(conn *WebSocketConnection) {
defer func() {
if r := recover(); r != nil {
app.errorLog.Println("ERROR:", fmt.Sprintf("%v", r))
}
}()
var payload WsPayload
for {
err := conn.ReadJSON(&payload)
if err != nil {
// do nothing
app.errorLog.Println(err)
break
}
payload.Conn = *conn
wsChan <- payload
}
}
func (app *application) ListenToWsChannel() {
var response WsJSONResponse
for {
e := <-wsChan
switch e.Action {
case "deleteUser":
response.Action = "logout"
response.Message = "Your account has ben deleted"
response.UserID = e.UserID
app.broadcastToAll(response)
default:
}
}
}
func (app *application) broadcastToAll(response WsJSONResponse) {
for client := range clients {
// broadcast to every connected client
err := client.WriteJSON(response)
if err != nil {
app.errorLog.Printf("Websocket err on %s: %s", response.Action, err)
_ = client.Close()
delete(clients, client)
}
}
}