Golang 实现 websocket 通讯（一个服务器，两个客户端）

运行客户端时首先要输入自己的名字，两个客户端输入的名字不能一样。

先上运行结果：

服务器端：

客户端1：

客户端2：

话不多说，直接上代码：

服务器端：

Go 复制代码

package main

import (
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// Client represents a connected client
type Client struct {
	conn *websocket.Conn
}

// Message represents a message from a client
type Message struct {
	Sender    string `json:"sender"`
	Recipient string `json:"recipient"`
	Content   string `json:"content"`
}

var clients = make(map[*Client]bool)
var broadcast = make(chan Message)

func handleConnections(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	client := &Client{conn: conn}
	clients[client] = true

	fmt.Println("Client connected")

	for {
		var msg Message
		err := conn.ReadJSON(&msg)
		if err != nil {
			log.Println(err)
			delete(clients, client)
			break
		}

		broadcast <- msg
	}
}

func handleMessages() {
	for {
		msg := <-broadcast

		for client := range clients {
			if client.conn != nil {
				err := client.conn.WriteJSON(msg)
				if err != nil {
					log.Println(err)
					client.conn.Close()
					delete(clients, client)
				}
			}
		}
	}
}

func main() {
	http.HandleFunc("/", handleConnections)
	go handleMessages()

	fmt.Println("WebSocket server listening on :8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

客户端1：

Go 复制代码

package main

import (
	"bufio"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"os"
)

type Message struct {
	Sender    string `json:"sender"`
	Recipient string `json:"recipient"`
	Content   string `json:"content"`
}

func main() {
	conn, _, err := websocket.DefaultDialer.Dial("ws://localhost:8080", nil)
	if err != nil {
		log.Println("Error connecting to server:", err)
		return
	}
	defer conn.Close()

	go func() {
		for {
			var msg Message
			err := conn.ReadJSON(&msg)
			if err != nil {
				log.Println("Error reading message:", err)
				return
			}
			fmt.Printf("Received from %s: %s\n", msg.Sender, msg.Content)
		}
	}()

	fmt.Print("Enter your name: ")
	var name string
	fmt.Scanln(&name)

	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		message := scanner.Text()
		msg := Message{Sender: name, Content: message}
		err := conn.WriteJSON(msg)
		if err != nil {
			log.Println("Error sending message:", err)
			return
		}
	}
}

客户端2：

Go 复制代码

package main

import (
	"bufio"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"os"
)

type Message struct {
	Sender    string `json:"sender"`
	Recipient string `json:"recipient"`
	Content   string `json:"content"`
}

func main() {
	conn, _, err := websocket.DefaultDialer.Dial("ws://localhost:8080", nil)
	if err != nil {
		log.Println("Error connecting to server:", err)
		return
	}
	defer conn.Close()

	go func() {
		for {
			var msg Message
			err := conn.ReadJSON(&msg)
			if err != nil {
				log.Println("Error reading message:", err)
				return
			}
			fmt.Printf("Received from %s: %s\n", msg.Sender, msg.Content)
		}
	}()

	fmt.Print("Enter your name: ")
	var name string
	fmt.Scanln(&name)

	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		message := scanner.Text()
		msg := Message{Sender: name, Content: message}
		err := conn.WriteJSON(msg)
		if err != nil {
			log.Println("Error sending message:", err)
			return
		}
	}
}