github.com/oarkflow/sio@v0.0.6/socket.go (about) 1 package sio 2 3 import ( 4 "encoding/json" 5 "fmt" 6 "log/slog" 7 "sync" 8 "time" 9 10 "github.com/oarkflow/frame" 11 "github.com/oarkflow/frame/pkg/common/xid" 12 "github.com/oarkflow/frame/pkg/websocket" 13 14 "github.com/oarkflow/sio/internal/bpool" 15 "github.com/oarkflow/sio/internal/maps" 16 ) 17 18 // Socket represents a websocket connection 19 type Socket struct { 20 l *sync.RWMutex 21 id string 22 ws *websocket.Conn 23 closed bool 24 serv *Server 25 roomsl *sync.RWMutex 26 context *maps.Map[string, any] 27 rooms map[string]bool 28 pingTicker *time.Ticker 29 tickerDone chan bool 30 pingInterval time.Duration 31 Ctx *frame.Context 32 } 33 34 const ( 35 typeJSON string = "J" 36 typeBin = "B" 37 typeStr = "S" 38 ) 39 40 func newSocket(serv *Server, ctx *frame.Context, ws *websocket.Conn) *Socket { 41 s := &Socket{ 42 l: &sync.RWMutex{}, 43 id: xid.New().String(), 44 ws: ws, 45 closed: false, 46 serv: serv, 47 roomsl: &sync.RWMutex{}, 48 rooms: make(map[string]bool), 49 context: maps.New[string, any](100000), 50 pingTicker: time.NewTicker(5 * time.Second), 51 tickerDone: make(chan bool), 52 Ctx: ctx, 53 } 54 serv.hub.addSocket(s) 55 go s.Ping() 56 return s 57 } 58 59 func (s *Socket) receive() ([]byte, error) { 60 _, data, err := s.ws.ReadMessage() 61 return data, err 62 } 63 64 func (s *Socket) send(msgType int, data []byte) error { 65 s.l.Lock() 66 defer s.l.Unlock() 67 return s.ws.WriteMessage(msgType, data) 68 } 69 70 func (s *Socket) Ping() error { 71 for { 72 select { 73 case <-s.tickerDone: 74 return nil 75 case <-s.pingTicker.C: 76 buf := bpool.Get() 77 defer bpool.Put(buf) 78 buf.WriteString(fmt.Sprintf("%d", websocket.PongMessage)) 79 s.ws.WriteMessage(websocket.TextMessage, buf.Bytes()) 80 } 81 } 82 } 83 84 // InRoom returns true if s is currently a member of roomName 85 func (s *Socket) InRoom(roomName string) bool { 86 s.roomsl.RLock() 87 defer s.roomsl.RUnlock() 88 inRoom := s.rooms[roomName] 89 return inRoom 90 } 91 92 // Set get request 93 func (s *Socket) Set(key string, val any) { 94 s.context.Put(key, val) 95 } 96 97 // Get gets value 98 func (s *Socket) Get(key string) (any, bool) { 99 return s.context.Get(key) 100 } 101 102 // Context gets value 103 func (s *Socket) Context() *maps.Map[string, any] { 104 return s.context 105 } 106 107 // GetRooms returns a list of rooms that s is a member of 108 func (s *Socket) GetRooms() []string { 109 s.roomsl.RLock() 110 defer s.roomsl.RUnlock() 111 112 var roomList []string 113 for room := range s.rooms { 114 roomList = append(roomList, room) 115 } 116 return roomList 117 } 118 119 // Join adds s to the specified room. If the room does 120 // not exist, it will be created 121 func (s *Socket) Join(roomName string) { 122 s.roomsl.Lock() 123 defer s.roomsl.Unlock() 124 s.serv.hub.joinRoom(&joinRequest{roomName, s}) 125 s.rooms[roomName] = true 126 } 127 128 // Leave removes s from the specified room. If s 129 // is not a member of the room, nothing will happen. If the room is 130 // empty upon removal of s, the room will be closed 131 func (s *Socket) Leave(roomName string) { 132 s.roomsl.Lock() 133 defer s.roomsl.Unlock() 134 s.serv.hub.leaveRoom(&leaveRequest{roomName, s}) 135 delete(s.rooms, roomName) 136 } 137 138 // LeaveAll removes s from the specified room. If s 139 // is not a member of the room, nothing will happen. If the room is 140 // empty upon removal of s, the room will be closed 141 func (s *Socket) LeaveAll() { 142 for roomName := range s.rooms { 143 s.Leave(roomName) 144 } 145 } 146 147 // ToRoom dispatches an event to all Sockets in the specified room. 148 func (s *Socket) ToRoom(roomName, eventName string, data any) { 149 s.serv.hub.toRoom(&RoomMsg{RoomName: roomName, EventName: eventName, Data: data}) 150 } 151 152 // ToRoomExcept dispatches an event to all Sockets in the specified room. 153 func (s *Socket) ToRoomExcept(roomName string, except []string, eventName string, data any) { 154 s.serv.hub.toRoom(&RoomMsg{RoomName: roomName, EventName: eventName, Data: data, Except: except}) 155 } 156 157 // Broadcast dispatches an event to all Sockets on the Server. 158 func (s *Socket) Broadcast(eventName string, data any) { 159 s.serv.hub.broadcast(&BroadcastMsg{EventName: eventName, Data: data}) 160 } 161 162 // BroadcastExcept dispatches an event to all Sockets on the Server. 163 func (s *Socket) BroadcastExcept(except []string, eventName string, data any) { 164 s.serv.hub.broadcast(&BroadcastMsg{EventName: eventName, Data: data, Except: except}) 165 } 166 167 // ToSocket dispatches an event to the specified socket ID. 168 func (s *Socket) ToSocket(socketID, eventName string, data any) { 169 s.serv.ToRoom("__socket_id:"+socketID, eventName, data) 170 } 171 172 // Emit dispatches an event to s. 173 func (s *Socket) Emit(eventName string, data any) error { 174 return s.send(emitData(eventName, data)) 175 } 176 177 // ID returns the unique ID of s 178 func (s *Socket) ID() string { 179 return s.id 180 } 181 182 // emitData combines the eventName and data into a payload that is understood 183 // by the sac-sock protocol. 184 func emitData(eventName string, data any) (int, []byte) { 185 buf := bpool.Get() 186 defer bpool.Put(buf) 187 buf.WriteString(eventName) 188 buf.WriteByte(startOfHeaderByte) 189 190 switch d := data.(type) { 191 case string: 192 buf.WriteString(typeStr) 193 buf.WriteByte(startOfDataByte) 194 buf.WriteString(d) 195 return websocket.TextMessage, buf.Bytes() 196 197 case []byte: 198 buf.WriteString(typeBin) 199 buf.WriteByte(startOfDataByte) 200 buf.Write(d) 201 return websocket.BinaryMessage, buf.Bytes() 202 203 default: 204 buf.WriteString(typeJSON) 205 buf.WriteByte(startOfDataByte) 206 jsonData, err := json.Marshal(d) 207 if err != nil { 208 slog.Error(err.Error()) 209 } else { 210 buf.Write(jsonData) 211 } 212 return websocket.TextMessage, buf.Bytes() 213 } 214 } 215 216 // Close closes the Socket connection and removes the Socket 217 // from any rooms that it was a member of 218 func (s *Socket) Close() error { 219 s.l.Lock() 220 isAlreadyClosed := s.closed 221 s.closed = true 222 s.l.Unlock() 223 224 if isAlreadyClosed { // can't reclose the socket 225 return nil 226 } 227 228 defer slog.Debug(s.ID(), "disconnected") 229 230 err := s.ws.Close() 231 if err != nil { 232 return err 233 } 234 235 rooms := s.GetRooms() 236 237 for _, room := range rooms { 238 s.Leave(room) 239 } 240 241 s.serv.l.RLock() 242 event := s.serv.onDisconnectFunc 243 s.serv.l.RUnlock() 244 245 if event != nil { 246 if err := event(s); err != nil { 247 return err 248 } 249 } 250 251 s.serv.hub.removeSocket(s) 252 if s.pingTicker != nil { 253 s.pingTicker.Stop() 254 s.tickerDone <- true 255 } 256 return nil 257 }