moq-rs/server/internal/warp/session.go

280 lines
6.2 KiB
Go
Raw Normal View History

2022-06-29 16:17:02 +00:00
package warp
import (
"context"
"encoding/binary"
"encoding/json"
"errors"
"fmt"
"io"
2022-11-18 23:13:35 +00:00
"log"
2022-06-29 16:17:02 +00:00
"math"
"time"
"github.com/kixelated/invoker"
2022-12-06 00:26:35 +00:00
"github.com/kixelated/quic-go"
"github.com/kixelated/webtransport-go"
2022-06-29 16:17:02 +00:00
)
// A single WebTransport session
type Session struct {
2022-11-18 23:13:35 +00:00
conn quic.Connection
inner *webtransport.Session
2022-11-18 23:13:35 +00:00
media *Media
2022-11-18 23:13:35 +00:00
inits map[string]*MediaInit
audio *MediaStream
video *MediaStream
2022-06-29 16:17:02 +00:00
streams invoker.Tasks
}
2022-11-18 23:13:35 +00:00
func NewSession(connection quic.Connection, session *webtransport.Session, media *Media) (s *Session, err error) {
2022-06-29 16:17:02 +00:00
s = new(Session)
2022-11-18 23:13:35 +00:00
s.conn = connection
2022-06-29 16:17:02 +00:00
s.inner = session
s.media = media
return s, nil
}
func (s *Session) Run(ctx context.Context) (err error) {
2022-11-18 23:13:35 +00:00
s.inits, s.audio, s.video, err = s.media.Start(s.conn.GetMaxBandwidth)
2022-06-29 16:17:02 +00:00
if err != nil {
return fmt.Errorf("failed to start media: %w", err)
}
// Once we've validated the session, now we can start accessing the streams
2022-11-18 23:13:35 +00:00
return invoker.Run(ctx, s.runAccept, s.runAcceptUni, s.runInit, s.runAudio, s.runVideo, s.streams.Repeat)
2022-06-29 16:17:02 +00:00
}
func (s *Session) runAccept(ctx context.Context) (err error) {
for {
stream, err := s.inner.AcceptStream(ctx)
2022-06-29 16:17:02 +00:00
if err != nil {
return fmt.Errorf("failed to accept bidirectional stream: %w", err)
}
// Warp doesn't utilize bidirectional streams so just close them immediately.
// We might use them in the future so don't close the connection with an error.
stream.CancelRead(1)
}
}
func (s *Session) runAcceptUni(ctx context.Context) (err error) {
for {
stream, err := s.inner.AcceptUniStream(ctx)
if err != nil {
return fmt.Errorf("failed to accept unidirectional stream: %w", err)
}
s.streams.Add(func(ctx context.Context) (err error) {
return s.handleStream(ctx, stream)
2022-06-29 16:17:02 +00:00
})
}
}
func (s *Session) handleStream(ctx context.Context, stream webtransport.ReceiveStream) (err error) {
2022-06-29 16:17:02 +00:00
defer func() {
if err != nil {
stream.CancelRead(1)
}
}()
var header [8]byte
for {
_, err = io.ReadFull(stream, header[:])
if errors.Is(io.EOF, err) {
return nil
} else if err != nil {
return fmt.Errorf("failed to read atom header: %w", err)
}
size := binary.BigEndian.Uint32(header[0:4])
name := string(header[4:8])
if size < 8 {
return fmt.Errorf("atom size is too small")
} else if size > 42069 { // arbitrary limit
return fmt.Errorf("atom size is too large")
} else if name != "warp" {
return fmt.Errorf("only warp atoms are supported")
}
payload := make([]byte, size-8)
_, err = io.ReadFull(stream, payload)
if err != nil {
return fmt.Errorf("failed to read atom payload: %w", err)
}
2022-11-18 23:13:35 +00:00
log.Println("received message:", string(payload))
2022-06-29 16:17:02 +00:00
msg := Message{}
err = json.Unmarshal(payload, &msg)
if err != nil {
return fmt.Errorf("failed to decode json payload: %w", err)
}
2022-11-18 23:13:35 +00:00
if msg.Debug != nil {
s.setDebug(msg.Debug)
2022-06-29 16:17:02 +00:00
}
}
}
2022-11-18 23:13:35 +00:00
func (s *Session) runInit(ctx context.Context) (err error) {
for _, init := range s.inits {
err = s.writeInit(ctx, init)
if err != nil {
return fmt.Errorf("failed to write init stream: %w", err)
}
2022-06-29 16:17:02 +00:00
}
2022-11-18 23:13:35 +00:00
return nil
}
2022-06-29 16:17:02 +00:00
2022-11-18 23:13:35 +00:00
func (s *Session) runAudio(ctx context.Context) (err error) {
2022-06-29 16:17:02 +00:00
for {
2022-11-18 23:13:35 +00:00
segment, err := s.audio.Next(ctx)
2022-06-29 16:17:02 +00:00
if err != nil {
return fmt.Errorf("failed to get next segment: %w", err)
}
if segment == nil {
return nil
}
2022-11-18 23:13:35 +00:00
err = s.writeSegment(ctx, segment)
2022-06-29 16:17:02 +00:00
if err != nil {
return fmt.Errorf("failed to write segment stream: %w", err)
}
}
}
func (s *Session) runVideo(ctx context.Context) (err error) {
for {
2022-11-18 23:13:35 +00:00
segment, err := s.video.Next(ctx)
2022-06-29 16:17:02 +00:00
if err != nil {
return fmt.Errorf("failed to get next segment: %w", err)
}
if segment == nil {
return nil
}
2022-11-18 23:13:35 +00:00
err = s.writeSegment(ctx, segment)
2022-06-29 16:17:02 +00:00
if err != nil {
return fmt.Errorf("failed to write segment stream: %w", err)
}
}
}
// Create a stream for an INIT segment and write the container.
2022-11-18 23:13:35 +00:00
func (s *Session) writeInit(ctx context.Context, init *MediaInit) (err error) {
2022-06-29 16:17:02 +00:00
temp, err := s.inner.OpenUniStreamSync(ctx)
if err != nil {
return fmt.Errorf("failed to create stream: %w", err)
}
2023-03-28 05:57:18 +00:00
if temp == nil {
// Not sure when this happens, perhaps when closing a connection?
return fmt.Errorf("received a nil stream from quic-go")
}
2022-06-29 16:17:02 +00:00
// Wrap the stream in an object that buffers writes instead of blocking.
stream := NewStream(temp)
s.streams.Add(stream.Run)
defer func() {
if err != nil {
stream.WriteCancel(1)
}
}()
stream.SetPriority(math.MaxInt)
err = stream.WriteMessage(Message{
2022-11-18 23:13:35 +00:00
Init: &MessageInit{Id: init.ID},
2022-06-29 16:17:02 +00:00
})
if err != nil {
return fmt.Errorf("failed to write init header: %w", err)
}
_, err = stream.Write(init.Raw)
if err != nil {
return fmt.Errorf("failed to write init data: %w", err)
}
2023-03-27 22:54:41 +00:00
err = stream.Close()
if err != nil {
return fmt.Errorf("failed to close init stream: %w", err)
}
2022-06-29 16:17:02 +00:00
return nil
}
// Create a stream for a segment and write the contents, chunk by chunk.
2022-11-18 23:13:35 +00:00
func (s *Session) writeSegment(ctx context.Context, segment *MediaSegment) (err error) {
2022-06-29 16:17:02 +00:00
temp, err := s.inner.OpenUniStreamSync(ctx)
if err != nil {
return fmt.Errorf("failed to create stream: %w", err)
}
2023-03-27 22:54:41 +00:00
if temp == nil {
// Not sure when this happens, perhaps when closing a connection?
2023-03-28 05:57:18 +00:00
return fmt.Errorf("received a nil stream from quic-go")
2023-03-27 22:54:41 +00:00
}
2022-06-29 16:17:02 +00:00
// Wrap the stream in an object that buffers writes instead of blocking.
stream := NewStream(temp)
s.streams.Add(stream.Run)
defer func() {
if err != nil {
stream.WriteCancel(1)
}
}()
ms := int(segment.timestamp / time.Millisecond)
// newer segments take priority
stream.SetPriority(ms)
err = stream.WriteMessage(Message{
Segment: &MessageSegment{
2022-11-18 23:13:35 +00:00
Init: segment.Init.ID,
2022-06-29 16:17:02 +00:00
Timestamp: ms,
},
})
if err != nil {
return fmt.Errorf("failed to write segment header: %w", err)
}
for {
// Get the next fragment
buf, err := segment.Read(ctx)
if errors.Is(err, io.EOF) {
break
} else if err != nil {
return fmt.Errorf("failed to read segment data: %w", err)
}
// NOTE: This won't block because of our wrapper
_, err = stream.Write(buf)
if err != nil {
return fmt.Errorf("failed to write segment data: %w", err)
}
}
err = stream.Close()
if err != nil {
return fmt.Errorf("failed to close segemnt stream: %w", err)
}
return nil
}
2022-11-18 23:13:35 +00:00
func (s *Session) setDebug(msg *MessageDebug) {
s.conn.SetMaxBandwidth(uint64(msg.MaxBitrate))
2022-06-29 16:17:02 +00:00
}