IT博客汇 | 像Redis作者那样，使用Go实现一个聊天服务器，不到100行代码

像Redis作者那样，使用Go实现一个聊天服务器，不到100行代码

smallnest发表于 2023-10-30 13:33:08

昨天Redis的作者 antirez (Salvatore Sanfilippo) 昨天创建一个新的演示项目：smallchat，用了200行C语言代码实现了一个聊天室。我看了一下，觉得很有意思，于是就用Go语言实现了一下，代码不到100行，功能和antirez的实现一样。

antirez 三年前停止写代码，专心写他的科幻小说《Wohpe》，今天看起来他有回到编程的状态了。关于这个小项目的背景是：

昨天我正在与几个前端开发者朋友闲聊,他们距离系统编程有些远。我们回忆起了过去的IRC时光。不可避免地,我说:编写一个非常简单的IRC服务器每个人都应该做一次。这样程序中有非常有趣的部分。一个进程进行多路复用,维护客户端状态,可以用不同的方式实现等等。

然后讨论继续,我想,我会给你们展示一个极简的C语言例子。但是你能编写出啥样的最小聊天服务器呢?要真正做到极简,我们不应该需要任何特殊的客户端，即使不是很完美,它应该可以用telnet或nc(netcat)作为客户端连接。服务器的主要功能只是接收一些聊天信息并发送给所有其他客户端,这有时称为扇出操作。这还需要一个合适的readline()函数,然后是缓冲等等。我们想要更简单的:利用内核缓冲区,假装我们每次都从客户端收到一个完整的行(这个假设在实际中通常是正确的,所以这个假设没啥问题)。

好吧,有了这些技巧,我们可以用只有200行代码实现一个聊天室,用户甚至可以设置昵称(当然,不计空格和注释)。由于我将这个小程序作为示例编写给我的朋友,我决定也把它推到Github上。

嗯，挺有趣的事情，我也很羡慕 antirez 有时间聊一聊编程中的一些趣事和想法。

这也不免让我想起上大学的时候，大家还沉迷于在终端中使用telnet连接BBS服务器，或者玩mud的游戏，窗外还飘着《Yesterday Once More》的旋律。那时候的互联网刚刚开始。

嗯，然后这个无聊的下午，我就想使用Go实现antirez的这个程序，我也不知道目的是啥，就纯粹想玩一玩，练一练手，最终用了不到100行代码实现了一个聊天服务器，功能和antirez的实现一样。

这个代码我也放到了github上: smallnest/smallchat。

我们不妨看看代码：

package main
import (
	"flag"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"
)
const (
	maxClients = 1000
	maxNickLen = 32
)
var (
	serverPort = flag.Int("p", 8972, "server port")
)
type Client struct {
	conn net.Conn
	nick string
}
type ChatState struct {
	listener net.Listener
	clientsLock sync.RWMutex
	clients     map[net.Conn]*Client
	numClients  int
}
var chatState = &ChatState{
	clients: make(map[net.Conn]*Client),
}
func initChat() {
	var err error
	chatState.listener, err = net.Listen("tcp", fmt.Sprintf(":%d", *serverPort))
	if err != nil {
		fmt.Println("listen error:", err)
		os.Exit(1)
	}
}
func handleClient(client *Client) {
	// 发送欢迎信息
	welcomeMsg := "Welcome Simple Chat! Use /nick to change nick name.\n"
	client.conn.Write([]byte(welcomeMsg))
	buf := make([]byte, 256)
	for {
		n, err := client.conn.Read(buf)
		if err != nil {
			fmt.Printf("client left: %s\n", client.conn.RemoteAddr())
			chatState.clientsLock.Lock()
			delete(chatState.clients, client.conn)
			chatState.numClients--
			chatState.clientsLock.Unlock()
			return
		}
		msg := string(buf[:n])
		msg = strings.TrimSpace(msg)
		if msg[0] == '/' {
			// 处理命令
			parts := strings.SplitN(msg, " ", 2)
			cmd := parts[0]
			if cmd == "/nick" && len(parts) > 1 {
				client.nick = parts[1]
			}
			continue
		}
		fmt.Printf("%s: %s\n", client.nick, msg)
		// 将消息转发给其他客户端
		chatState.clientsLock.RLock()
		for conn, cl := range chatState.clients {
			if cl != client {
				conn.Write([]byte(client.nick + ": " + msg))
			}
		}
		chatState.clientsLock.RUnlock()
	}
}
func main() {
	flag.Parse()
	initChat()
	for {
		conn, err := chatState.listener.Accept()
		if err != nil {
			fmt.Println("accept error:", err)
			continue
		}
		client := &Client{conn: conn}
		client.nick = fmt.Sprintf("user%d", conn.RemoteAddr().(*net.TCPAddr).Port)
		chatState.clientsLock.Lock()
		chatState.clients[conn] = client
		chatState.numClients++
		chatState.clientsLock.Unlock()
		go handleClient(client)
		fmt.Printf("new client: %s\n", conn.RemoteAddr())
	}
}

首先我们从main函数说起。

main函数中我们首先调用initChat函数，这个函数中我们使用net.Listen创建了一个net.Listener，然后使用Accept方法接收客户端的连接。Accept方法返回一个net.Conn，这个net.Conn代表了一个客户端的连接，我们可以使用Read和Write方法读写数据。

为了跟踪每一个用户，我们定义了一个Client结构体，其中包含了一个net.Conn和一个nick字段，nick字段代表了用户的昵称。

我们使用一个ChatState结构体来保存聊天室的状态，其中包含了一个net.Listener和一个clients字段，clients字段是一个map[net.Conn]*Client，用来保存所有的客户端连接。ChatState还包含了一个clientsLock字段，这个字段是一个sync.RWMutex，用来保护clients字段，因为clients字段会被多个goroutine访问。

main函数中我们使用一个for循环来接收客户端的连接，然后调用handleClient函数来处理客户端的连接。

接下来就是handleClient函数了，这个函数中我们首先发送一个欢迎信息给客户端，然后使用一个for循环来读取客户端发送的消息，如果客户端断开连接，我们就从clients中删除这个客户端，然后退出循环。

我们假定用户的输入不超过256字节，然后我们使用strings.TrimSpace函数去掉消息前后的空格，然后判断消息是否以/开头。

如果客户端发送的消息以/开头，我们就认为这是一个命令，我们只处理/nick命令，这个命令用来设置客户端的昵称。

如果客户端发送的消息不是以/开头，我们就认为这是一个聊天消息，我们将这个消息转发给所有的客户端。

handleClient函数中我们使用了chatState.clientsLock来保护clients字段，因为clients字段会被多个goroutine访问。

这就是一个可工作的聊天服务器了，我们可以使用telnet或nc来连接这个服务器，然后就可以跟其他用户聊天了。登录进去后你可以使用/nick命令进行改名。

当然这只是一个玩具，没有任何的安全性检查，也没有任何的错误处理，但是它可以工作，而且代码很简单。