main.go
package main
import (
"fmt"
"net/http"
"sync"
"github.com/gorilla/websocket"
)
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
fmt.Println("upgrade")
return true
},
}
type Client struct {
conn *websocket.Conn
send chan []byte
}
type Hub struct {
clients map[*Client]bool
broadcast chan []byte
mutex sync.Mutex
}
var hub = Hub{
clients: make(map[*Client]bool),
broadcast: make(chan []byte),
}
func (h *Hub) run() {
for {
msg := <-h.broadcast
fmt.Println("hub get msg")
h.mutex.Lock()
for client := range h.clients {
select {
case client.send <- msg:
fmt.Println("send msg to client channel")
default:
close(client.send)
delete(h.clients, client)
}
}
h.mutex.Unlock()
}
}
func (c *Client) read() {
defer func() {
hub.mutex.Lock()
delete(hub.clients, c)
hub.mutex.Unlock()
c.conn.Close()
}()
for {
_, msg, err := c.conn.ReadMessage()
fmt.Println("read client msg", msg)
if err != nil {
break
}
hub.broadcast <- msg
}
}
func (c *Client) write() {
defer c.conn.Close()
// send是channel，for循环会阻塞等待消息
for msg := range c.send {
fmt.Println("send msg to client")
if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
break
}
}
}
func handleConnection(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
fmt.Println(err)
return
}
client := &Client{conn: conn, send: make(chan []byte)}
hub.mutex.Lock()
hub.clients[client] = true
hub.mutex.Unlock()
go client.write()
client.read()
}
func main() {
go hub.run()
http.HandleFunc("/ws", handleConnection)
http.Handle("/", http.FileServer(http.Dir(".")))
fmt.Println("Server started at :6060")
if err := http.ListenAndServe(":6060", nil); err != nil {
fmt.Println("Error while starting the server:", err)
}
}
消息通过read函数读取到内存中，并通过channel转发给hub.broadcast。
在hub.run函数中，从broadcast通道中获取消息，遍历所有client，将消息依次转发给所有client的sendchannel
write函数，读取channel中的消息，并将消息发送给client，for循环会阻塞等待send 通道的数据。
index.html
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Chat Room</title>
<link rel="stylesheet" href="./style.css">
</head>
<body>
<div id="chat-container">
<div id="messages"></div>
<input id="message-input" type="text" placeholder="Type a message..." autocomplete="off" />
<button id="send-button">Send</button>
</div>
<script>
const messages = document.getElementById('messages');
const input = document.getElementById('message-input');
const sendButton = document.getElementById('send-button');
const connection = new WebSocket('ws://' + window.location.host + '/ws');
connection.onmessage = function(event) {
const messageElement = document.createElement('div');
messageElement.textContent = event.data;
messages.appendChild(messageElement);
};
sendButton.onclick = function() {
const message = input.value;
if (message) {
connection.send(message);
input.value = '';
}
};
</script>
</body>
</html>
button按钮发送消息
窗口显示发送出去的消息
style.css
body {
font-family: Arial, sans-serif;
}
#chat-container {
width: 600px;
margin: 0 auto;
border: 1px solid #ccc;
padding: 10px;
}
#messages {
height: 400px;
overflow-y: scroll;
border: 1px solid #ccc;
margin-bottom: 10px;
}
#message-input {
width: 80%;
padding: 10px;
}
#send-button {
padding: 10px;
}
在输入框输入，点击发送，消息会经过go编写的后台服务器，回到消息框
