BareRTC/pkg/websocket.go

319 lines
7.6 KiB
Go

package barertc
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"time"
"git.kirsle.net/apps/barertc/pkg/jwt"
"git.kirsle.net/apps/barertc/pkg/log"
"git.kirsle.net/apps/barertc/pkg/util"
"nhooyr.io/websocket"
)
// Subscriber represents a connected WebSocket session.
type Subscriber struct {
// User properties
ID int // ID assigned by server
Username string
VideoActive bool
JWTClaims *jwt.Claims
authenticated bool // has passed the login step
conn *websocket.Conn
ctx context.Context
cancel context.CancelFunc
messages chan []byte
closeSlow func()
}
// ReadLoop spawns a goroutine that reads from the websocket connection.
func (sub *Subscriber) ReadLoop(s *Server) {
go func() {
for {
msgType, data, err := sub.conn.Read(sub.ctx)
if err != nil {
log.Error("ReadLoop error(%d=%s): %+v", sub.ID, sub.Username, err)
s.DeleteSubscriber(sub)
// Notify if this user was auth'd
if sub.authenticated {
s.Broadcast(Message{
Action: ActionPresence,
Username: sub.Username,
Message: "has exited the room!",
})
s.SendWhoList()
}
return
}
if msgType != websocket.MessageText {
log.Error("Unexpected MessageType")
continue
}
// Read the user's posted message.
var msg Message
log.Debug("Read(%d=%s): %s", sub.ID, sub.Username, data)
if err := json.Unmarshal(data, &msg); err != nil {
log.Error("Read(%d=%s) Message error: %s", sub.ID, sub.Username, err)
continue
}
// What action are they performing?
switch msg.Action {
case ActionLogin:
s.OnLogin(sub, msg)
case ActionMessage:
s.OnMessage(sub, msg)
case ActionMe:
s.OnMe(sub, msg)
case ActionOpen:
s.OnOpen(sub, msg)
case ActionCandidate:
s.OnCandidate(sub, msg)
case ActionSDP:
s.OnSDP(sub, msg)
case ActionWatch:
s.OnWatch(sub, msg)
case ActionUnwatch:
s.OnUnwatch(sub, msg)
default:
sub.ChatServer("Unsupported message type.")
}
}
}()
}
// SendJSON sends a JSON message to the websocket client.
func (sub *Subscriber) SendJSON(v interface{}) error {
data, err := json.Marshal(v)
if err != nil {
return err
}
log.Debug("SendJSON(%d=%s): %s", sub.ID, sub.Username, data)
return sub.conn.Write(sub.ctx, websocket.MessageText, data)
}
// SendMe sends the current user state to the client.
func (sub *Subscriber) SendMe() {
sub.SendJSON(Message{
Action: ActionMe,
Username: sub.Username,
VideoActive: sub.VideoActive,
})
}
// ChatServer is a convenience function to deliver a ChatServer error to the client.
func (sub *Subscriber) ChatServer(message string, v ...interface{}) {
sub.SendJSON(Message{
Action: ActionError,
Username: "ChatServer",
Message: fmt.Sprintf(message, v...),
})
}
// WebSocket handles the /ws websocket connection.
func (s *Server) WebSocket() http.HandlerFunc {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ip := util.IPAddress(r)
log.Info("WebSocket connection from %s - %s", ip, r.Header.Get("User-Agent"))
log.Debug("Headers: %+v", r.Header)
c, err := websocket.Accept(w, r, nil)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
fmt.Fprintf(w, "Could not accept websocket connection: %s", err)
return
}
defer c.Close(websocket.StatusInternalError, "the sky is falling")
log.Debug("WebSocket: %s has connected", ip)
// CloseRead starts a goroutine that will read from the connection
// until it is closed.
// ctx := c.CloseRead(r.Context())
ctx, cancel := context.WithCancel(r.Context())
sub := &Subscriber{
conn: c,
ctx: ctx,
cancel: cancel,
messages: make(chan []byte, s.subscriberMessageBuffer),
closeSlow: func() {
c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages")
},
}
s.AddSubscriber(sub)
defer s.DeleteSubscriber(sub)
go sub.ReadLoop(s)
pinger := time.NewTicker(PingInterval)
for {
select {
case msg := <-sub.messages:
err = writeTimeout(ctx, time.Second*5, c, msg)
if err != nil {
return
}
case timestamp := <-pinger.C:
sub.SendJSON(Message{
Action: ActionPing,
Message: timestamp.Format(time.RFC3339),
})
case <-ctx.Done():
pinger.Stop()
return
}
}
})
}
// Auto incrementing Subscriber ID, assigned in AddSubscriber.
var SubscriberID int
// AddSubscriber adds a WebSocket subscriber to the server.
func (s *Server) AddSubscriber(sub *Subscriber) {
// Assign a unique ID.
SubscriberID++
sub.ID = SubscriberID
log.Debug("AddSubscriber: ID #%d", sub.ID)
s.subscribersMu.Lock()
s.subscribers[sub] = struct{}{}
s.subscribersMu.Unlock()
}
// GetSubscriber by username.
func (s *Server) GetSubscriber(username string) (*Subscriber, error) {
s.subscribersMu.RLock()
defer s.subscribersMu.RUnlock()
for _, sub := range s.IterSubscribers(true) {
if sub.Username == username {
return sub, nil
}
}
return nil, errors.New("not found")
}
// DeleteSubscriber removes a subscriber from the server.
func (s *Server) DeleteSubscriber(sub *Subscriber) {
log.Error("DeleteSubscriber: %s", sub.Username)
// Cancel its context to clean up the for-loop goroutine.
if sub.cancel != nil {
sub.cancel()
}
s.subscribersMu.Lock()
delete(s.subscribers, sub)
s.subscribersMu.Unlock()
}
// IterSubscribers loops over the subscriber list with a read lock. If the
// caller already holds a lock, pass the optional `true` parameter for isLocked.
func (s *Server) IterSubscribers(isLocked ...bool) []*Subscriber {
var result = []*Subscriber{}
// Has the caller already taken the read lock or do we get it?
if locked := len(isLocked) > 0 && isLocked[0]; !locked {
s.subscribersMu.RLock()
defer s.subscribersMu.RUnlock()
}
for sub := range s.subscribers {
result = append(result, sub)
}
return result
}
// Broadcast a message to the chat room.
func (s *Server) Broadcast(msg Message) {
log.Debug("Broadcast: %+v", msg)
s.subscribersMu.RLock()
defer s.subscribersMu.RUnlock()
for _, sub := range s.IterSubscribers(true) {
if !sub.authenticated {
continue
}
sub.SendJSON(msg)
}
}
// SendTo sends a message to a given username.
func (s *Server) SendTo(username string, msg Message) error {
log.Debug("SendTo(%s): %+v", username, msg)
username = strings.TrimPrefix(username, "@")
s.subscribersMu.RLock()
defer s.subscribersMu.RUnlock()
// If no timestamp, add it.
if msg.Timestamp.IsZero() {
msg.Timestamp = time.Now()
}
var found bool
for _, sub := range s.IterSubscribers(true) {
if sub.Username == username {
found = true
sub.SendJSON(Message{
Action: msg.Action,
Channel: msg.Channel,
Username: msg.Username,
Message: msg.Message,
})
}
}
if !found {
return fmt.Errorf("%s is not online", username)
}
return nil
}
// SendWhoList broadcasts the connected members to everybody in the room.
func (s *Server) SendWhoList() {
var (
users = []WhoList{}
subscribers = s.IterSubscribers()
)
for _, sub := range subscribers {
if !sub.authenticated {
continue
}
who := WhoList{
Username: sub.Username,
VideoActive: sub.VideoActive,
}
if sub.JWTClaims != nil {
who.Operator = sub.JWTClaims.IsAdmin
who.Avatar = sub.JWTClaims.Avatar
who.ProfileURL = sub.JWTClaims.ProfileURL
}
users = append(users, who)
}
for _, sub := range subscribers {
sub.SendJSON(Message{
Action: ActionWhoList,
WhoList: users,
Timestamp: time.Now(),
})
}
}
func writeTimeout(ctx context.Context, timeout time.Duration, c *websocket.Conn, msg []byte) error {
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
return c.Write(ctx, websocket.MessageText, msg)
}