FreePad/lib/socketmanager/socketmanager.go

198 lines
4.8 KiB
Go

package socketmanager
import (
"encoding/json"
"errors"
"fmt"
"net/http"
"github.com/JustKato/FreePad/lib/objects"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"github.com/gorilla/websocket"
)
var wsUpgrader = websocket.Upgrader{
ReadBufferSize: 1024, // TODO: Make it configurable via the .env file
WriteBufferSize: 1024, // TODO: Make it configurable via the .env file
}
// The pad socket map caches all of the existing sockets
var padSocketMap map[string]map[string]*websocket.Conn = make(map[string]map[string]*websocket.Conn)
// TODO: Use generics so that we can take string messages, that'd be nice!
type SocketMessage struct {
EventType string `json:"eventType"`
PadName string `json:"padName"`
Message map[string]interface{} `json:"message"`
}
// Bind the websockets to the gin router
func BindSocket(router *gin.RouterGroup) {
router.GET("/get/:pad", func(ctx *gin.Context) {
// Get the name of the pad to assign to this socket
padName := ctx.Param("pad")
// Upgrade the socket connection
webSocketUpgrade(ctx.Writer, ctx.Request, padName)
})
}
func webSocketUpgrade(w http.ResponseWriter, r *http.Request, padName string) {
conn, err := wsUpgrader.Upgrade(w, r, nil)
if err != nil {
fmt.Printf("Failed to set websocket upgrade: %v\n", err)
return
}
// Check if we have any sockets in this padName
if _, ok := padSocketMap[padName]; !ok {
// Initialize a new map of sockets
padSocketMap[padName] = make(map[string]*websocket.Conn)
}
// Give this socket a token
socketToken := uuid.NewString()
// Set the current connection at the socket Token position
padSocketMap[padName][socketToken] = conn
// Somone just connected
UpdatePadStatus(padName)
// Start listening to this socket
for {
// Try Read the JSON input from the socket
_, msg, err := conn.ReadMessage()
// Check if anything but a read limit was created
if err != nil && !errors.Is(err, websocket.ErrReadLimit) {
// Remove self from the cache
delete(padSocketMap[padName], socketToken)
// Somone just disconnected
UpdatePadStatus(padName)
break
}
if err != nil {
// There has been an error reading the message
fmt.Println("Failed to read from the socket but probably still connected")
// Skip this cycle
continue
}
// Init the variable
var p SocketMessage
// Try and parse the json
err = json.Unmarshal([]byte(msg), &p)
if err != nil {
// There has been an error reading the message
fmt.Println("Failed to parse the JSON", err)
// Skip this cycle
continue
}
// Pass the message to the proper handlers
handleSocketMessage(p, socketToken, padName)
}
}
// Handle the socket's message
func handleSocketMessage(msg SocketMessage, socketToken string, padName string) {
// Check if this is a pad Update
if msg.EventType == `padUpdate` {
handlePadUpdate(msg, socketToken, padName)
// Serialize the message
serialized, err := json.Marshal(msg)
// Check if there was an error
if err != nil {
fmt.Println(`Failed to broadcast the padUpdate`, err)
// Stop the execution
return
}
// Alert all the other pads other than this one.
for k, pad := range padSocketMap[padName] {
// Check if this is the same socket.
if k == socketToken {
// Skip self
continue
}
// Send the message to the others.
pad.WriteMessage(websocket.TextMessage, serialized)
}
}
}
func handlePadUpdate(msg SocketMessage, socketToken string, padName string) {
// Check if the msg content is valid
if _, ok := msg.Message[`content`]; !ok {
fmt.Printf("Failed to update pad %s, invalid message\n", padName)
return
}
// Check that the content is string
newPadContent, ok := msg.Message[`content`].(string)
if !ok {
fmt.Printf("Type assertion failed for %s, invalid message\n", padName)
return
}
// Get the pad
pad := objects.GetPost(padName, false)
// Update the pad contents
pad.Content = newPadContent
// Save to file
objects.WritePost(pad)
}
// Update the current users of the pad about the amount of live viewers.
func UpdatePadStatus(padName string) {
// Grab info about the map's key
sockets, ok := padSocketMap[padName]
// Check if the pad is set and has sockets connected.
if !ok || len(sockets) < 1 {
// Quit
return
}
// Generate the message
msg := SocketMessage{
EventType: `statusUpdate`,
PadName: padName,
Message: gin.H{
// Send the current amount of live viewers
"currentViewers": len(sockets),
},
}
BroadcastMessage(padName, msg)
}
func BroadcastMessage(padName string, msg SocketMessage) {
// Grab info about the map's key
sockets, ok := padSocketMap[padName]
// Check if the pad is set and has sockets connected.
if !ok || len(sockets) < 1 {
// Quit
return
}
// Get all the participants of the pad group
for _, s := range sockets {
// Send the message to the socket
s.WriteJSON(msg)
}
}