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

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

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

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

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

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

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

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

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

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

我们不妨看看代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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代表了一个客户端的连接,我们可以使用ReadWrite方法读写数据。

为了跟踪每一个用户,我们定义了一个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访问。

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

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