Go实现简单的IM系统

Posted on In

客户端，服务器。net库，flag库。

1 main.go

1
2
3
4
5
6
package main

func main() {
	server := NewServer("127.0.0.1", 9999)
	server.Start()
}

2 server.go

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
package main

import (
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

type Server struct {
	IP string
	Port int
	C chan string
	OnlineMap map[string]*User
	Locker sync.RWMutex
}

func NewServer(ip string, port int) *Server {
	server := &Server{
		IP: ip,
		Port: port,
		C: make(chan string),
		OnlineMap: make(map[string]*User),
	}
	return server
}

func (s *Server) ListenMsg() {
	for {
		msg := <- s.C
		s.Locker.Lock()
		for _, usr := range s.OnlineMap {
			usr.C <- msg
		}
		s.Locker.Unlock()
	}
}

func (s *Server) Start() {
	addr := fmt.Sprintf("%s:%d", s.IP, s.Port)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer listener.Close()

	go s.ListenMsg()

	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println(err)
			continue
		}
		go s.handler(conn)
	}
}

func (s *Server) handler(conn net.Conn) {
	user := NewUser(s, conn)
	user.Online()
	isLive := make(chan bool)

	go func() {
		buf := make([]byte, 4096)
		for {
			n, err := conn.Read(buf)
			if n == 0 {
				user.Offline()
				return
			}
			if err != nil && err != io.EOF {
				fmt.Println(err)
				return
			}

			msg := string(buf)
			user.DoMessage(msg)
			isLive <- true
		}
	}()

	for {
		select {
		case <-isLive:

		case <-time.After(time.Second * 300):
			user.Offline()
			close(user.C)
			conn.Close()
			return
		}
	}
}

func (s *Server) BroadCast(u *User, msg string) {
	sendMsg := "[" + u.Addr + "]" +u.Name + ":" + msg
	s.Locker.RLock()
	for _, user := range s.OnlineMap {
		user.C <- sendMsg
	}
	s.Locker.RUnlock()
}

3 client.go

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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
package main

import (
	"flag"
	"fmt"
	"io"
	"net"
	"os"
)

type Client struct {
	ServerIP   string
	ServerPort int
	Name       string
	conn       net.Conn
	flag       int
}

var ip string
var port int

func init() {
	flag.StringVar(&ip, "ip", "127.0.0.1", "sever ip address")
	flag.IntVar(&port, "port", 9999, "server port")
}

func main() {
	flag.Parse()
	client := NewClient()
	if client == nil {
		fmt.Println("link server err")
		return
	}
	go client.DealResponse()
	fmt.Println("success")
	client.Run()
}

func NewClient() *Client {
	cli := &Client{
		ServerIP: ip,
		ServerPort: port,
		flag: 999,
	}
	addr := fmt.Sprintf("%s:%d", ip, port)
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		fmt.Println(err)
		return nil
	}
	cli.conn = conn
	cli.Name = conn.RemoteAddr().String()
	return cli
}

func (cli *Client) Run() {
	for cli.flag != 0 {
		for cli.menu() != true {}
		switch cli.flag {
		case 1:
			cli.PublicChat()
		case 2:
			cli.PrivateChat()
		case 3:
			cli.UpdateName()
		}
	}
}

func (cli *Client) DealResponse() {
	io.Copy(os.Stdout, cli.conn)
}

func (cli *Client) menu() bool {
	var flag int

	fmt.Println("1.公聊模式")
	fmt.Println("2.私聊模式")
	fmt.Println("3.更新用户名")
	fmt.Println("0.退出")

	fmt.Scanln(&flag)

	if flag >= 0 && flag <= 3 {
		cli.flag = flag
		return true
	} else {
		fmt.Println(">>>>请输入合法范围内的数字<<<<")
		return false
	}
}

func (cli *Client) SelectUsers() {
	sendMsg := "who\n"
	_, err := cli.conn.Write([]byte(sendMsg))
	if err != nil {
		fmt.Println("conn Write err:", err)
		return
	}
}

func (cli *Client) PrivateChat() {
	var remoteName string
	var chatMsg string

	cli.SelectUsers()
	fmt.Println(">>>>请输入聊天对象[用户名], exit退出:")
	fmt.Scanln(&remoteName)

	for remoteName != "exit" {
		fmt.Println(">>>>请输入消息内容, exit退出:")
		fmt.Scanln(&chatMsg)
		for chatMsg != "exit" {
			if len(chatMsg) != 0 {
				sendMsg := "to|" + remoteName + "|" + chatMsg + "\n\n"
				_, err := cli.conn.Write([]byte(sendMsg))
				if err != nil {
					fmt.Println("conn Write err:", err)
					break
				}
			}
			chatMsg = ""
			fmt.Println(">>>>请输入消息内容, exit退出:")
			fmt.Scanln(&chatMsg)
		}
		cli.SelectUsers()
		fmt.Println(">>>>请输入聊天对象[用户名], exit退出:")
		fmt.Scanln(&remoteName)
	}
}

func (cli *Client) PublicChat() {
	var chatMsg string
	fmt.Println(">>>>请输入聊天内容，exit退出.")
	fmt.Scanln(&chatMsg)
	for chatMsg != "exit" {
		if len(chatMsg) != 0 {
			sendMsg := chatMsg + "\n"
			_, err := cli.conn.Write([]byte(sendMsg))
			if err != nil {
				fmt.Println("conn Write err:", err)
				break
			}
		}
		chatMsg = ""
		fmt.Println(">>>>请输入聊天内容，exit退出.")
		fmt.Scanln(&chatMsg)
	}

}

func (cli *Client) UpdateName() bool {
	fmt.Println(">>>>请输入用户名:")
	fmt.Scanln(&cli.Name)
	sendMsg := "rename|" + cli.Name + "\n"
	_, err := cli.conn.Write([]byte(sendMsg))
	if err != nil {
		fmt.Println("conn.Write err:", err)
		return false
	}
	return true
}


4 user.go

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
114
115
116
117
package main

import (
	"fmt"
	"net"
	"strings"
)

type User struct {
	Addr   string
	Name   string
	C      chan string
	conn   net.Conn
	server *Server
}

func NewUser(server *Server, conn net.Conn) *User {
	user := &User{
		Addr: conn.RemoteAddr().String(),
		Name: conn.RemoteAddr().String(),
		C: make(chan string),
		conn: conn,
		server: server,
	}
	go user.receiveMsg()
	return user
}

func (u *User) receiveMsg() {
	for {
		msg := <- u.C
		u.conn.Write([]byte(msg))
	}
}

func (u *User) Online() {
	//addr := fmt.Sprintf("[%s:%d]SERVER: ", u.server.IP, u.server.Port)
	//msg := addr + u.Name + " online\n"
	//u.server.BroadCast(msg)
	u.server.Locker.Lock()
	u.server.OnlineMap[u.Name] = u
	u.server.Locker.Unlock()
	u.server.BroadCast(u, "online\n")
}

func (u *User) Offline() {
	//addr := fmt.Sprintf("[%s:%d]SERVER: ", u.server.IP, u.server.Port)
	//msg := addr + u.Name + " offline\n"
	//u.server.BroadCast(msg)
	u.server.Locker.Lock()
	delete(u.server.OnlineMap, u.Name)
	u.server.Locker.Unlock()
	u.server.BroadCast(u, "offline\n")
}



func (u *User) SendMsg(msg string) {
	u.C <- msg
}

func (u *User) DoMessage(msg string) {
	// rename|name
	if len(msg) > 7 && msg[:7] == "rename|" {
		newName := msg[7:]
		u.rename(newName)
	// to|name|msg
	} else if len(msg) > 3 && msg[:3] == "to|" {
		nameAndMsg := msg[3:]
		u.privateChat(nameAndMsg)
	} else if msg == "who" {
		u.who()
	} else {
		u.server.BroadCast(u, msg)
	}
}

func (u *User) who() {
	addr := fmt.Sprintf("[%s:%d]SERVER:", u.server.IP, u.server.Port)
	u.SendMsg(addr + ">>>>>>>>online map<<<<<<<<\n")
	u.server.Locker.RLock()
	for _, user := range u.server.OnlineMap {
		msg :=  addr + user.Name + " online\n"
		u.SendMsg(msg)
	}
	u.server.Locker.RUnlock()
	u.SendMsg(addr + ">>>>>>>>>end map<<<<<<<<<\n")
}

func (u *User) rename(newName string) {
	if _, ok := u.server.OnlineMap[newName]; ok {
		u.SendMsg("name exist\n")
	} else {
		u.server.Locker.Lock()
		delete(u.server.OnlineMap, u.Name)
		u.server.OnlineMap[newName] = u
		u.server.Locker.Unlock()
		u.Name = newName
		u.SendMsg("rename success\n")
	}
}

// to|name|msg
func (u *User) privateChat(nameAndMsg string) {
	name := strings.Split(nameAndMsg, "|")[1]
	msg := strings.Split(nameAndMsg, "|")[2]
	if _, ok := u.server.OnlineMap[name]; !ok {
		u.SendMsg("user not exist\n")
		return
	}
	if len(msg) == 0 {
		return
	}
	remoteUser := u.server.OnlineMap[name]
	remoteUser.SendMsg(u.Name + "say:" + msg)
}