This commit is contained in:
Yvan
2026-09-23 17:42:55 +08:00
parent 53e8599bdd
commit 0a8ecbe09b
75 changed files with 2299 additions and 276 deletions
+122 -45
View File
@@ -34,6 +34,7 @@ type Manager struct {
peers map[*melody.Session]*peer
observers map[*melody.Session]struct{}
nextSub atomic.Uint64
nextGuest atomic.Uint64
stop chan struct{}
stopOnce sync.Once
}
@@ -53,16 +54,19 @@ type publishedTrack struct {
}
type peer struct {
session *melody.Session
manager *Manager
room *room
streamID string
userID uint
role string
subscriberID string
pc *webrtc.PeerConnection
lastBeat atomic.Int64
cleanup sync.Once
session *melody.Session
manager *Manager
room *room
streamID string
userID uint
role string
authenticated bool
subscriberID string
pcMu sync.RWMutex
pc *webrtc.PeerConnection
closing bool
lastBeat atomic.Int64
cleanup sync.Once
}
type message struct {
@@ -210,20 +214,28 @@ func (m *Manager) join(session *melody.Session, payload joinPayload) {
m.send(session, "JoinRejected", map[string]any{"code": "INVALID_JOIN", "message": "role and streamId are required"})
return
}
userID, ok := sessionUserID(session)
if !ok {
userID, authenticated := sessionUserID(session)
if role == "publisher" && !authenticated {
m.send(session, "JoinRejected", map[string]any{"code": "AUTH_REQUIRED", "message": "authenticated user is required"})
return
}
if role == "subscriber" && !authenticated {
userID = m.anonymousViewerID()
}
if existing := m.peer(session); existing != nil {
m.closePeer(existing, "rejoin")
}
r := m.getOrCreateRoom(streamID)
if role == "publisher" {
if err := m.authorizeRoomPublisher(streamID, userID); err != nil {
m.send(session, "JoinRejected", map[string]any{"code": "PUBLISH_FORBIDDEN", "message": err.Error()})
return
}
} else if err := m.authorizeRoomSubscriber(streamID); err != nil {
m.send(session, "JoinRejected", map[string]any{"code": "ROOM_UNAVAILABLE", "message": err.Error()})
return
}
r := m.getOrCreateRoom(streamID)
if role == "publisher" {
r.mu.RLock()
occupied := r.publisher != nil
r.mu.RUnlock()
@@ -235,7 +247,10 @@ func (m *Manager) join(session *melody.Session, payload joinPayload) {
m.send(session, "JoinRejected", map[string]any{"code": errorCode(err), "message": err.Error()})
return
}
p := &peer{session: session, manager: m, room: r, streamID: streamID, userID: userID, role: role}
if roomID, ok := roomIDFromStream(streamID); ok {
_ = m.chat.Clear(context.Background(), roomID)
}
p := &peer{session: session, manager: m, room: r, streamID: streamID, userID: userID, role: role, authenticated: true}
p.lastBeat.Store(time.Now().UnixMilli())
r.mu.Lock()
r.publisher = p
@@ -244,7 +259,6 @@ func (m *Manager) join(session *melody.Session, payload joinPayload) {
m.peers[session] = p
m.mu.Unlock()
m.send(session, "Joined", map[string]any{"role": role, "streamId": streamID, "mode": "pion-sfu"})
m.sendChatHistory(session, streamID)
m.sendPublisherStatus(streamID, false, "publisher_waiting_for_offer")
m.broadcastStreamList()
return
@@ -269,7 +283,7 @@ func (m *Manager) join(session *melody.Session, payload joinPayload) {
m.send(session, "JoinRejected", map[string]any{"code": errorCode(err), "message": err.Error()})
return
}
p := &peer{session: session, manager: m, room: r, streamID: streamID, userID: userID, role: role, subscriberID: "sub_" + strconv.FormatUint(m.nextSub.Add(1), 10)}
p := &peer{session: session, manager: m, room: r, streamID: streamID, userID: userID, role: role, authenticated: authenticated, subscriberID: "sub_" + strconv.FormatUint(m.nextSub.Add(1), 10)}
p.lastBeat.Store(time.Now().UnixMilli())
r.mu.Lock()
r.viewers[p] = struct{}{}
@@ -278,14 +292,17 @@ func (m *Manager) join(session *melody.Session, payload joinPayload) {
m.mu.Lock()
m.peers[session] = p
m.mu.Unlock()
m.send(session, "Joined", map[string]any{"role": role, "streamId": streamID, "subscriberId": p.subscriberID, "userId": userID, "mode": "pion-sfu"})
m.sendChatHistory(session, streamID)
publicUserID := userID
if !authenticated {
publicUserID = 0
}
m.send(session, "Joined", map[string]any{"role": role, "streamId": streamID, "subscriberId": p.subscriberID, "userId": publicUserID, "mode": "pion-sfu"})
m.sendPublisherStatusTo(session, streamID, ready, "publisher_offline")
if ready {
m.send(session, "PublisherStatus", map[string]any{"streamId": streamID, "ready": true})
}
if publisher := m.publisher(r); publisher != nil {
m.send(publisher.session, "NewSubscriber", map[string]any{"streamId": streamID, "subscriberId": p.subscriberID, "userId": userID})
m.send(publisher.session, "NewSubscriber", map[string]any{"streamId": streamID, "subscriberId": p.subscriberID, "userId": publicUserID})
}
m.broadcastStreamList()
}
@@ -296,6 +313,10 @@ func (m *Manager) handleChat(session *melody.Session, payload chatSendPayload) {
m.send(session, "Error", map[string]any{"code": "NOT_JOINED", "message": "join a room before chatting"})
return
}
if !p.authenticated {
m.send(session, "Error", map[string]any{"code": "AUTH_REQUIRED", "message": "authenticated user is required for chat"})
return
}
roomID, ok := roomIDFromStream(p.streamID)
if !ok {
m.send(session, "Error", map[string]any{"code": "CHAT_UNAVAILABLE", "message": "chat is only available in live rooms"})
@@ -309,23 +330,6 @@ func (m *Manager) handleChat(session *melody.Session, payload chatSendPayload) {
m.broadcastRoom(p.streamID, "ChatMessage", chatMessagePayload(item))
}
func (m *Manager) sendChatHistory(session *melody.Session, streamID string) {
roomID, ok := roomIDFromStream(streamID)
if !ok {
return
}
items, err := m.chat.ListRecent(context.Background(), roomID, 100)
if err != nil {
m.log("load live chat history failed", "room_id", roomID, "error", err)
return
}
messages := make([]map[string]any, 0, len(items))
for _, item := range items {
messages = append(messages, chatMessagePayload(item))
}
m.send(session, "ChatHistory", map[string]any{"roomId": roomID, "messages": messages})
}
// GiftSent implements gift.Notifier. Gift persistence and wallet debit have
// already committed; chat persistence is therefore best-effort and must never
// turn a successful debit into a failed gift response.
@@ -388,6 +392,23 @@ func (m *Manager) authorizeRoomPublisher(streamID string, userID uint) error {
return nil
}
func (m *Manager) authorizeRoomSubscriber(streamID string) error {
if strings.HasPrefix(streamID, "room-") {
id, err := strconv.ParseUint(strings.TrimPrefix(streamID, "room-"), 10, 64)
if err != nil || id == 0 {
return errors.New("invalid live room stream id")
}
if _, err = m.roomCatalog.GetPublic(context.Background(), uint(id)); err != nil {
return errors.New("live room is unavailable")
}
}
stream, err := m.streams.Get(context.Background(), streamID)
if err != nil || stream == nil || !stream.PublisherOnline {
return errors.New("publisher is offline")
}
return nil
}
func (m *Manager) handleOffer(session *melody.Session, payload sessionDescriptionPayload) {
p := m.peer(session)
if p == nil || (p.role != "publisher" && p.role != "subscriber") {
@@ -398,20 +419,22 @@ func (m *Manager) handleOffer(session *melody.Session, payload sessionDescriptio
m.send(session, "Error", map[string]any{"code": "INVALID_SDP", "message": "sdp is required"})
return
}
if p.pc != nil {
_ = p.pc.Close()
}
pc, err := m.api.NewPeerConnection(webrtc.Configuration{ICEServers: liveICEServers()})
if err != nil {
m.send(session, "Error", map[string]any{"code": "PEER_CONNECTION", "message": err.Error()})
return
}
p.pc = pc
oldPC, installed := p.replacePeerConnection(pc)
if !installed {
_ = pc.Close()
return
}
pc.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
if state == webrtc.ICEConnectionStateFailed || state == webrtc.ICEConnectionStateClosed {
m.closePeer(p, "ice_"+state.String())
}
m.handleICEConnectionStateChange(p, pc, state)
})
if oldPC != nil {
_ = oldPC.Close()
}
if p.role == "publisher" {
pc.OnTrack(func(track *webrtc.TrackRemote, _ *webrtc.RTPReceiver) {
m.onPublisherTrack(p, track)
@@ -448,6 +471,13 @@ func (m *Manager) handleOffer(session *melody.Session, payload sessionDescriptio
}
}
func (m *Manager) handleICEConnectionStateChange(p *peer, pc *webrtc.PeerConnection, state webrtc.ICEConnectionState) {
if state != webrtc.ICEConnectionStateFailed && state != webrtc.ICEConnectionStateClosed {
return
}
m.closePeerConnection(p, pc, "ice_"+state.String())
}
func (m *Manager) onPublisherTrack(p *peer, remote *webrtc.TrackRemote) {
codec := remote.Codec().RTPCodecCapability
local, err := webrtc.NewTrackLocalStaticRTP(codec, remote.ID(), p.streamID)
@@ -514,15 +544,26 @@ func (m *Manager) heartbeatLoop() {
}
func (m *Manager) closePeer(p *peer, reason string) {
m.closePeerConnection(p, nil, reason)
}
func (m *Manager) closePeerConnection(p *peer, candidate *webrtc.PeerConnection, reason string) {
if p == nil {
return
}
pc, ok := p.beginClose(candidate)
if !ok {
return
}
p.cleanup.Do(func() {
if p.pc != nil {
_ = p.pc.Close()
if pc != nil {
_ = pc.Close()
}
if p.role == "publisher" {
_ = m.streams.UnregisterPublisher(context.Background(), p.streamID, p.userID)
if roomID, ok := roomIDFromStream(p.streamID); ok {
_ = m.chat.Clear(context.Background(), roomID)
}
p.room.mu.Lock()
tracks := make([]*publishedTrack, 0, len(p.room.tracks))
for _, track := range p.room.tracks {
@@ -566,6 +607,36 @@ func (m *Manager) closePeer(p *peer, reason string) {
})
}
func (p *peer) replacePeerConnection(next *webrtc.PeerConnection) (*webrtc.PeerConnection, bool) {
p.pcMu.Lock()
defer p.pcMu.Unlock()
if p.closing {
return nil, false
}
previous := p.pc
p.pc = next
return previous, true
}
func (p *peer) hasPeerConnection(candidate *webrtc.PeerConnection) bool {
p.pcMu.RLock()
current := p.pc
p.pcMu.RUnlock()
return candidate != nil && current == candidate
}
func (p *peer) beginClose(candidate *webrtc.PeerConnection) (*webrtc.PeerConnection, bool) {
p.pcMu.Lock()
defer p.pcMu.Unlock()
if p.closing || candidate != nil && p.pc != candidate {
return nil, false
}
p.closing = true
current := p.pc
p.pc = nil
return current, true
}
func (m *Manager) getOrCreateRoom(streamID string) *room {
m.mu.Lock()
defer m.mu.Unlock()
@@ -685,6 +756,12 @@ func sessionUserID(session *melody.Session) (uint, bool) {
}
}
func (m *Manager) anonymousViewerID() uint {
// Allocate downward from the maximum uint so anonymous connections never
// share the zero ID and cannot collide with ordinary auto-increment users.
return ^uint(0) - uint(m.nextGuest.Add(1))
}
func peerStreamID(m *Manager, session *melody.Session) string {
if p := m.peer(session); p != nil {
return p.streamID
+201
View File
@@ -0,0 +1,201 @@
package live
import (
"context"
"sync"
"testing"
"github.com/olahol/melody"
"github.com/pion/webrtc/v4"
bizlive "kra/internal/biz/live"
platformws "kra/pkg/websocket"
)
type liveHubStub struct {
mu sync.Mutex
messages map[*melody.Session][][]byte
}
func (h *liveHubStub) Enabled() bool { return true }
func (h *liveHubStub) Broadcast([]byte) error { return nil }
func (h *liveHubStub) BroadcastBinary([]byte) error { return nil }
func (h *liveHubStub) Sessions() ([]*platformws.Session, error) { return nil, nil }
func (h *liveHubStub) Len() int { return 0 }
func (h *liveHubStub) SendBinary(*platformws.Session, []byte) error { return nil }
func (h *liveHubStub) OnMessage(func(*platformws.Session, []byte)) {}
func (h *liveHubStub) OnBinaryMessage(func(*platformws.Session, []byte)) {}
func (h *liveHubStub) OnConnect(func(*platformws.Session)) {}
func (h *liveHubStub) OnDisconnect(func(*platformws.Session)) {}
func (h *liveHubStub) Send(session *platformws.Session, message []byte) error {
h.mu.Lock()
defer h.mu.Unlock()
if h.messages == nil {
h.messages = make(map[*melody.Session][][]byte)
}
h.messages[session] = append(h.messages[session], append([]byte(nil), message...))
return nil
}
type liveCatalogStub struct {
bizlive.CatalogRepo
room *bizlive.Room
err error
}
func (s *liveCatalogStub) GetRoom(context.Context, uint) (*bizlive.Room, error) {
if s.err != nil {
return nil, s.err
}
copy := *s.room
return &copy, nil
}
func (s *liveCatalogStub) GetCategory(context.Context, uint) (*bizlive.Category, error) {
return &bizlive.Category{ID: s.room.CategoryID, Enabled: true}, nil
}
type liveStreamStub struct {
bizlive.StreamRepo
mu sync.Mutex
streams map[string]*bizlive.Stream
users map[uint]struct{}
registerCalls int
}
func (s *liveStreamStub) Get(_ context.Context, streamID string) (*bizlive.Stream, error) {
s.mu.Lock()
defer s.mu.Unlock()
stream := s.streams[streamID]
if stream == nil {
return &bizlive.Stream{StreamID: streamID}, nil
}
copy := *stream
return &copy, nil
}
func (s *liveStreamStub) List(context.Context) ([]*bizlive.Stream, error) {
s.mu.Lock()
defer s.mu.Unlock()
items := make([]*bizlive.Stream, 0, len(s.streams))
for _, stream := range s.streams {
if stream.PublisherOnline {
copy := *stream
items = append(items, &copy)
}
}
return items, nil
}
func (s *liveStreamStub) RegisterSubscriber(_ context.Context, streamID string, userID uint) (*bizlive.Stream, error) {
s.mu.Lock()
defer s.mu.Unlock()
s.registerCalls++
if s.users == nil {
s.users = make(map[uint]struct{})
}
if _, exists := s.users[userID]; exists {
return nil, bizlive.ErrSubscriberExists
}
s.users[userID] = struct{}{}
return &bizlive.Stream{StreamID: streamID, SubscriberCount: len(s.users)}, nil
}
func newAnonymousTestManager(catalog bizlive.CatalogRepo, streams bizlive.StreamRepo) *Manager {
streamUC := bizlive.NewStreamUsecase(streams)
return &Manager{
hub: &liveHubStub{},
streams: streamUC,
roomCatalog: bizlive.NewRoomUsecase(catalog, streamUC),
rooms: make(map[string]*room),
peers: make(map[*melody.Session]*peer),
observers: make(map[*melody.Session]struct{}),
stop: make(chan struct{}),
}
}
func TestAnonymousSubscribersReceiveDistinctRuntimeIdentities(t *testing.T) {
streams := &liveStreamStub{streams: map[string]*bizlive.Stream{"room-7": {StreamID: "room-7", PublisherOnline: true}}}
manager := newAnonymousTestManager(&liveCatalogStub{room: &bizlive.Room{ID: 7, CategoryID: 1, Enabled: true}}, streams)
first := &melody.Session{Keys: map[string]any{}}
second := &melody.Session{Keys: map[string]any{}}
manager.join(first, joinPayload{Role: "subscriber", StreamID: "room-7"})
manager.join(second, joinPayload{Role: "subscriber", StreamID: "room-7"})
p1, p2 := manager.peer(first), manager.peer(second)
if p1 == nil || p2 == nil {
t.Fatalf("anonymous subscribers did not join: first=%+v second=%+v", p1, p2)
}
if p1.authenticated || p2.authenticated || p1.userID == 0 || p2.userID == 0 || p1.userID == p2.userID {
t.Fatalf("anonymous identities are not unique: first=%d second=%d", p1.userID, p2.userID)
}
if len(streams.users) != 2 {
t.Fatalf("subscriber count = %d, want 2", len(streams.users))
}
}
func TestSubscriberCannotCreateUnknownCatalogRoom(t *testing.T) {
manager := newAnonymousTestManager(&liveCatalogStub{err: bizlive.ErrRoomNotFound}, &liveStreamStub{})
session := &melody.Session{Keys: map[string]any{}}
manager.join(session, joinPayload{Role: "subscriber", StreamID: "room-999999"})
if manager.peer(session) != nil || len(manager.rooms) != 0 {
t.Fatalf("unknown room created runtime state: peers=%d rooms=%d", len(manager.peers), len(manager.rooms))
}
}
func TestSubscriberCannotCreateOfflineRuntimeState(t *testing.T) {
streams := &liveStreamStub{}
manager := newAnonymousTestManager(&liveCatalogStub{room: &bizlive.Room{ID: 7, CategoryID: 1, Enabled: true}}, streams)
session := &melody.Session{Keys: map[string]any{}}
manager.join(session, joinPayload{Role: "subscriber", StreamID: "room-7"})
if manager.peer(session) != nil || len(manager.rooms) != 0 || len(streams.users) != 0 || streams.registerCalls != 0 {
t.Fatalf("offline room created runtime state: peers=%d rooms=%d users=%d registrations=%d", len(manager.peers), len(manager.rooms), len(streams.users), streams.registerCalls)
}
}
func TestAnonymousPublisherIsRejected(t *testing.T) {
manager := newAnonymousTestManager(&liveCatalogStub{room: &bizlive.Room{ID: 7, Enabled: true}}, &liveStreamStub{})
session := &melody.Session{Keys: map[string]any{}}
manager.join(session, joinPayload{Role: "publisher", StreamID: "room-7"})
if manager.peer(session) != nil {
t.Fatal("anonymous publisher joined")
}
}
func TestClosedSupersededPeerConnectionDoesNotRemovePeer(t *testing.T) {
oldPC, err := webrtc.NewPeerConnection(webrtc.Configuration{})
if err != nil {
t.Fatal(err)
}
defer oldPC.Close()
newPC, err := webrtc.NewPeerConnection(webrtc.Configuration{})
if err != nil {
t.Fatal(err)
}
defer newPC.Close()
p := &peer{pc: oldPC}
replaced, installed := p.replacePeerConnection(newPC)
if !installed || replaced != oldPC {
t.Fatalf("replaced peer connection = %p, want %p", replaced, oldPC)
}
manager := &Manager{}
manager.handleICEConnectionStateChange(p, oldPC, webrtc.ICEConnectionStateClosed)
cleanupAvailable := false
p.cleanup.Do(func() { cleanupAvailable = true })
if !cleanupAvailable {
t.Fatal("superseded peer connection consumed peer cleanup")
}
if !p.hasPeerConnection(newPC) {
t.Fatal("replacement peer connection is no longer active")
}
}
var _ platformws.Hub = (*liveHubStub)(nil)