-
Notifications
You must be signed in to change notification settings - Fork 52
/
main.go
106 lines (87 loc) · 1.83 KB
/
main.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
package main
import (
"errors"
"io"
"log"
"net"
"net/http"
"strconv"
"sync"
"github.com/gin-gonic/gin"
"github.com/gobwas/ws"
"github.com/gobwas/ws/wsutil"
)
const (
addr = ":9000"
)
type wsService struct {
connections map[int]net.Conn
mu sync.Mutex
index int
}
func (s *wsService) addConnection(conn net.Conn) int {
s.mu.Lock()
defer s.mu.Unlock()
i := s.index
s.connections[s.index] = conn
s.index += 1
return i
}
func (s *wsService) chat(w http.ResponseWriter, r *http.Request) {
conn, _, _, err := ws.UpgradeHTTP(r, w)
if err != nil {
log.Printf("can't upgrade connection: %s\n", err.Error())
return
}
connID := s.addConnection(conn)
id := strconv.Itoa(connID)
ch := make(chan []byte)
go func() {
defer func() {
conn.Close()
close(ch)
}()
for {
msg, _, err := wsutil.ReadClientData(conn)
if err != nil {
if !errors.Is(err, io.EOF) {
log.Printf("can't read message from connection: %s\n", err.Error())
}
break
}
ch <- msg
}
}()
go func() {
for msg := range ch {
msg = append([]byte(id+": "), msg...)
s.mu.Lock()
for key, connection := range s.connections {
if err := wsutil.WriteServerMessage(connection, ws.OpText, msg); err != nil {
log.Printf("can't write message: %s\n", err.Error())
delete(s.connections, key)
}
}
s.mu.Unlock()
}
log.Println("go func stop")
s.mu.Lock()
delete(s.connections, connID)
s.mu.Unlock()
}()
}
func main() {
service := wsService{
connections: make(map[int]net.Conn, 0),
mu: sync.Mutex{},
}
router := gin.Default()
router.GET("/chat", gin.WrapF(service.chat))
server := http.Server{
Addr: addr,
Handler: router,
}
if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
log.Fatalf("can't listen and serve server: %s", err.Error())
}
}