155 lines
2.6 KiB
Go
155 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"log"
|
|
"net/http"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
var upgrader = websocket.Upgrader{
|
|
ReadBufferSize: 1024,
|
|
WriteBufferSize: 1024,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
return true
|
|
},
|
|
}
|
|
|
|
type Hub struct {
|
|
clients map[*Client]bool
|
|
broadcast chan []byte
|
|
register chan *Client
|
|
unregister chan *Client
|
|
}
|
|
|
|
type Client struct {
|
|
hub *Hub
|
|
conn *websocket.Conn
|
|
send chan []byte
|
|
}
|
|
|
|
type Message struct {
|
|
Type string `json:"type"`
|
|
Content string `json:"content"`
|
|
Sender string `json:"sender"`
|
|
}
|
|
|
|
func newHub() *Hub {
|
|
return &Hub{
|
|
broadcast: make(chan []byte),
|
|
register: make(chan *Client),
|
|
unregister: make(chan *Client),
|
|
clients: make(map[*Client]bool),
|
|
}
|
|
}
|
|
|
|
func (h *Hub) run() {
|
|
for {
|
|
select {
|
|
case client := <-h.register:
|
|
h.clients[client] = true
|
|
case client := <-h.unregister:
|
|
if _, ok := h.clients[client]; ok {
|
|
delete(h.clients, client)
|
|
close(client.send)
|
|
}
|
|
case message := <-h.broadcast:
|
|
for client := range h.clients {
|
|
select {
|
|
case client.send <- message:
|
|
default:
|
|
close(client.send)
|
|
delete(h.clients, client)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Client) readPump() {
|
|
defer func() {
|
|
c.hub.unregister <- c
|
|
c.conn.Close()
|
|
}()
|
|
|
|
for {
|
|
_, message, err := c.conn.ReadMessage()
|
|
if err != nil {
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
|
log.Printf("error: %v", err)
|
|
}
|
|
break
|
|
}
|
|
c.hub.broadcast <- message
|
|
}
|
|
}
|
|
|
|
func (c *Client) writePump() {
|
|
defer func() {
|
|
c.conn.Close()
|
|
}()
|
|
|
|
for {
|
|
message, ok := <-c.send
|
|
if !ok {
|
|
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
|
|
return
|
|
}
|
|
|
|
w, err := c.conn.NextWriter(websocket.TextMessage)
|
|
if err != nil {
|
|
return
|
|
}
|
|
w.Write(message)
|
|
|
|
n := len(c.send)
|
|
for i := 0; i < n; i++ {
|
|
w.Write([]byte("\n"))
|
|
w.Write(<-c.send)
|
|
}
|
|
|
|
if err := w.Close(); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func handleWebSocket(hub *Hub) gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
|
if err != nil {
|
|
log.Println(err)
|
|
return
|
|
}
|
|
|
|
client := &Client{
|
|
hub: hub,
|
|
conn: conn,
|
|
send: make(chan []byte, 256),
|
|
}
|
|
client.hub.register <- client
|
|
|
|
go client.writePump()
|
|
go client.readPump()
|
|
}
|
|
}
|
|
|
|
func main() {
|
|
hub := newHub()
|
|
go hub.run()
|
|
|
|
r := gin.Default()
|
|
|
|
r.GET("/ws", handleWebSocket(hub))
|
|
r.GET("/health", func(c *gin.Context) {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"status": "ok",
|
|
})
|
|
})
|
|
|
|
log.Println("Gateway server starting on port 8080...")
|
|
if err := r.Run(":8080"); err != nil {
|
|
log.Fatalf("Failed to start server: %v", err)
|
|
}
|
|
} |