package plugin_room import ( "context" _ "embed" "encoding/json" "fmt" "net" "net/http" "strings" "sync" "github.com/gobwas/ws" "github.com/gobwas/ws/wsutil" "github.com/google/uuid" . "m7s.live/v5" "m7s.live/v5/pkg" "m7s.live/v5/pkg/task" "m7s.live/v5/pkg/util" ) // User 结构体代表一个房间用户 type User struct { task.Task Subscriber *Subscriber ID string StreamPath string Token string `json:"-" yaml:"-"` Room *Room `json:"-" yaml:"-"` net.Conn `json:"-" yaml:"-"` writeLock sync.Mutex } func (u *User) GetKey() string { return u.ID } func (u *User) Start() (err error) { u.Subscriber, err = u.Room.Plugin.Subscribe(u, u.StreamPath) if err != nil { return err } u.Subscriber.DataChannel = make(chan pkg.IDataFrame, 10) return nil } func (u *User) Go() error { for data := range u.Subscriber.DataChannel { u.writeLock.Lock() wsutil.WriteServerText(u.Conn, data.([]byte)) u.writeLock.Unlock() } return nil } func (u *User) Dispose() { close(u.Subscriber.DataChannel) } func (u *User) Send(event string, data any) { if u.Conn != nil { u.writeLock.Lock() defer u.writeLock.Unlock() j, err := json.Marshal(map[string]any{"event": event, "data": data}) if err == nil { wsutil.WriteServerText(u.Conn, j) } } } type Room struct { *Publisher ID string Users task.Manager[string, *User] } //go:embed default.yaml var defaultYaml DefaultYaml type RoomPlugin struct { Plugin AppName string `default:"room" desc:"用于订阅房间消息的应用名(streamPath第一段)"` Size int `default:"20" desc:"房间大小"` Private map[string]string `desc:"私密房间" key:"房间号" value:"密码"` Verify struct { URL string `desc:"验证用户身份的URL"` Method string `desc:"验证用户身份的HTTP方法"` Header map[string]string `desc:"验证用户身份的HTTP头" key:"名称" value:"值"` } `desc:"验证用户身份"` Ping string `default:"ping" desc:"用于客户端与服务器保持心跳时客户端发送的特殊字符串"` Pong string `default:"pong" desc:"用于客户端与服务器保持心跳时服务器响应的特殊字符串"` lock sync.RWMutex rooms util.Collection[string, *Room] } var _ = InstallPlugin[RoomPlugin](defaultYaml) func (rc *RoomPlugin) OnPublish(p *Publisher) { args := p.Args token := args.Get("token") ss := strings.Split(token, ":") if len(ss) != 3 { return } roomId := ss[0] userId := ss[1] if roomId != "" && rc.rooms.Has(roomId) { room, _ := rc.rooms.Get(roomId) if user, ok := room.Users.Get(userId); ok { if user.Token == token { user.StreamPath = p.StreamPath data, _ := json.Marshal(map[string]any{ "event": "publish", "data": p.StreamPath, "userId": user.ID, }) room.WriteData(data) } } } } func (rc *RoomPlugin) ServeHTTP(w http.ResponseWriter, r *http.Request) { ss := strings.Split(r.URL.Path, "/")[1:] var roomId, userId, token string if len(ss) == 2 { roomId = ss[0] userId = ss[1] } else { http.Error(w, "invalid url", http.StatusBadRequest) return } if rc.Verify.URL != "" { req, _ := http.NewRequest(rc.Verify.Method, rc.Verify.URL, nil) req.Header = r.Header res, _ := http.DefaultClient.Do(req) if res.StatusCode != 200 { http.Error(w, "verify failed", http.StatusForbidden) } } if rc.Private != nil { rc.lock.RLock() pass, ok := rc.Private[roomId] rc.lock.Unlock() if ok { if pass != r.URL.Query().Get("password") { http.Error(w, "password wrong", http.StatusForbidden) return } } } var room *Room var ok bool if room, ok = rc.rooms.Get(roomId); !ok { room = &Room{ID: roomId} if _, err := rc.Publish(context.Background(), rc.AppName+"/"+roomId); err == nil { rc.rooms.Add(room) } else { http.Error(w, "room already exist", http.StatusBadRequest) return } } if room.Users.Has(userId) { http.Error(w, "user exist", http.StatusBadRequest) return } conn, _, _, err := ws.UpgradeHTTP(r, w) if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } defer conn.Close() token = fmt.Sprintf("%s:%s:%s", roomId, userId, uuid.NewString()) user := &User{Room: room, Conn: conn, Token: token, ID: userId} data, _ := json.Marshal(map[string]any{"event": "userjoin", "data": user}) room.WriteData(data) room.Users.Add(user) user.Send("joined", map[string]any{"token": token, "userList": room.Users.Items}) defer func() { user.Stop(err) if room.Users.Length == 0 { room.Stop(err) rc.rooms.RemoveByKey(roomId) } }() var msg []byte var op ws.OpCode for { msg, op, err = wsutil.ReadClientData(conn) if string(msg) == rc.Ping { wsutil.WriteServerText(conn, []byte(rc.Pong)) } else { data, _ := json.Marshal(map[string]any{"event": "msg", "data": string(msg), "userId": userId}) room.WriteData(data) } if op == ws.OpClose || err != nil { data, _ := json.Marshal(map[string]any{"event": "userleave", "userId": userId, "data": user}) room.WriteData(data) return } } }