galene/diskwriter/diskwriter.go

package diskwriter

import (
	crand "crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/at-wat/ebml-go/webm"
	"github.com/pion/rtp"
	"github.com/pion/rtp/codecs"
	"github.com/pion/webrtc/v3"
	"github.com/pion/webrtc/v3/pkg/media/samplebuilder"

	"sfu/conn"
	"sfu/group"
)

var Directory string

type Client struct {
	group *group.Group
	id    string

	mu     sync.Mutex
	down   map[string]*diskConn
	closed bool
}

func newId() string {
	b := make([]byte, 16)
	crand.Read(b)
	return hex.EncodeToString(b)
}

func New(g *group.Group) *Client {
	return &Client{group: g, id: newId()}
}

func (client *Client) Group() *group.Group {
	return client.group
}

func (client *Client) Id() string {
	return client.id
}

func (client *Client) Credentials() group.ClientCredentials {
	return group.ClientCredentials{"RECORDING", ""}
}

func (client *Client) OverridePermissions(g *group.Group) bool {
	return true
}

func (client *Client) SetPermissions(perms group.ClientPermissions) {
	return
}

func (client *Client) PushClient(id, username string, add bool) error {
	return nil
}

func (client *Client) Close() error {
	client.mu.Lock()
	defer client.mu.Unlock()

	for _, down := range client.down {
		down.Close()
	}
	client.down = nil
	client.closed = true
	return nil
}

func (client *Client) Kick(message string) error {
	err := client.Close()
	group.DelClient(client)
	return err
}

func (client *Client) PushConn(id string, up conn.Up, tracks []conn.UpTrack, label string) error {
	client.mu.Lock()
	defer client.mu.Unlock()

	if client.closed {
		return errors.New("disk client is closed")
	}

	old := client.down[id]
	if old != nil {
		old.Close()
		delete(client.down, id)
	}

	if up == nil {
		return nil
	}

	directory := filepath.Join(Directory, client.group.Name())
	err := os.MkdirAll(directory, 0700)
	if err != nil {
		return err
	}

	if client.down == nil {
		client.down = make(map[string]*diskConn)
	}

	down, err := newDiskConn(directory, label, up, tracks)
	if err != nil {
		return err
	}

	client.down[up.Id()] = down
	return nil
}

type diskConn struct {
	directory string
	label     string
	hasVideo  bool

	mu            sync.Mutex
	file          *os.File
	remote        conn.Up
	tracks        []*diskTrack
	width, height uint32
}

// called locked
func (conn *diskConn) reopen() error {
	for _, t := range conn.tracks {
		if t.writer != nil {
			t.writer.Close()
			t.writer = nil
		}
	}
	conn.file = nil

	file, err := openDiskFile(conn.directory, conn.label)
	if err != nil {
		return err
	}

	conn.file = file
	return nil
}

func (conn *diskConn) Close() error {
	conn.remote.DelLocal(conn)

	conn.mu.Lock()
	tracks := make([]*diskTrack, 0, len(conn.tracks))
	for _, t := range conn.tracks {
		if t.writer != nil {
			t.writer.Close()
			t.writer = nil
		}
		tracks = append(tracks, t)
	}
	conn.mu.Unlock()

	for _, t := range tracks {
		t.remote.DelLocal(t)
	}
	return nil
}

func openDiskFile(directory, label string) (*os.File, error) {
	filename := time.Now().Format("2006-01-02T15:04:05.000")
	if label != "" {
		filename = filename + "-" + label
	}
	for counter := 0; counter < 100; counter++ {
		var fn string
		if counter == 0 {
			fn = fmt.Sprintf("%v.webm", filename)
		} else {
			fn = fmt.Sprintf("%v-%02d.webm", filename, counter)
		}

		fn = filepath.Join(directory, fn)
		f, err := os.OpenFile(
			fn, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600,
		)
		if err == nil {
			return f, nil
		} else if !os.IsExist(err) {
			return nil, err
		}
	}
	return nil, errors.New("couldn't create file")
}

type diskTrack struct {
	remote conn.UpTrack
	conn   *diskConn

	writer  webm.BlockWriteCloser
	builder *samplebuilder.SampleBuilder

	// bit 32 is a boolean indicating that the origin is valid
	origin uint64

	lastKf uint32
}

func newDiskConn(directory, label string, up conn.Up, remoteTracks []conn.UpTrack) (*diskConn, error) {
	conn := diskConn{
		directory: directory,
		label:     label,
		tracks:    make([]*diskTrack, 0, len(remoteTracks)),
		remote:    up,
	}
	for _, remote := range remoteTracks {
		var builder *samplebuilder.SampleBuilder
		switch remote.Codec().Name {
		case webrtc.Opus:
			builder = samplebuilder.New(
				16, &codecs.OpusPacket{},
				samplebuilder.WithPartitionHeadChecker(
					&codecs.OpusPartitionHeadChecker{},
				),
			)
		case webrtc.VP8:
			if conn.hasVideo {
				return nil, errors.New("multiple video tracks not supported")
			}
			builder = samplebuilder.New(
				128, &codecs.VP8Packet{},
				samplebuilder.WithPartitionHeadChecker(
					&codecs.VP8PartitionHeadChecker{},
				),
			)
			conn.hasVideo = true
		}
		track := &diskTrack{
			remote:  remote,
			builder: builder,
			conn:    &conn,
		}
		conn.tracks = append(conn.tracks, track)
		remote.AddLocal(track)
	}

	err := up.AddLocal(&conn)
	if err != nil {
		return nil, err
	}

	return &conn, nil
}

func (t *diskTrack) SetTimeOffset(ntp uint64, rtp uint32) {
}

func (t *diskTrack) SetCname(string) {
}

func clonePacket(packet *rtp.Packet) *rtp.Packet {
	buf, err := packet.Marshal()
	if err != nil {
		return nil
	}
	var p rtp.Packet
	err = p.Unmarshal(buf)
	if err != nil {
		return nil
	}
	return &p
}

func (t *diskTrack) WriteRTP(packet *rtp.Packet) error {
	// since we call initWriter, we take the connection lock for simplicity.
	t.conn.mu.Lock()
	defer t.conn.mu.Unlock()

	if t.builder == nil {
		return nil
	}

	p := clonePacket(packet)
	if p == nil {
		return nil
	}

	kfNeeded := false

	t.builder.Push(p)

	for {
		sample, ts := t.builder.PopWithTimestamp()
		if sample == nil {
			if kfNeeded {
				return conn.ErrKeyframeNeeded
			}
			return nil
		}

		keyframe := true

		switch t.remote.Codec().Name {
		case webrtc.VP8:
			if len(sample.Data) < 1 {
				continue
			}
			keyframe = (sample.Data[0]&0x1 == 0)
			if keyframe {
				err := t.initWriter(sample.Data)
				if err != nil {
					return err
				}
				t.lastKf = ts
			} else if t.writer != nil {
				// Request a keyframe every 10s
				delta := ts - t.lastKf
				if (delta&0x80000000) == 0 &&
					delta > 10*90000 {
					kfNeeded = true
				}
			}
		default:
			if t.writer == nil {
				if !t.conn.hasVideo {
					err := t.conn.initWriter(0, 0)
					if err != nil {
						return err
					}
				}
			}
		}

		if t.writer == nil {
			if !keyframe {
				return conn.ErrKeyframeNeeded
			}
			return nil
		}

		if t.origin == 0 {
			t.origin = uint64(ts) | (1 << 32)
		}
		ts -= uint32(t.origin)

		tm := ts / (t.remote.Codec().ClockRate / 1000)
		_, err := t.writer.Write(keyframe, int64(tm), sample.Data)
		if err != nil {
			return err
		}
	}
}

// called locked
func (t *diskTrack) initWriter(data []byte) error {
	switch t.remote.Codec().Name {
	case webrtc.VP8:
		if len(data) < 10 {
			return nil
		}
		keyframe := (data[0]&0x1 == 0)
		if !keyframe {
			return nil
		}
		raw := uint32(data[6]) | uint32(data[7])<<8 |
			uint32(data[8])<<16 | uint32(data[9])<<24
		width := raw & 0x3FFF
		height := (raw >> 16) & 0x3FFF
		return t.conn.initWriter(width, height)
	}
	return nil
}

// called locked
func (conn *diskConn) initWriter(width, height uint32) error {
	if conn.file != nil && width == conn.width && height == conn.height {
		return nil
	}
	var entries []webm.TrackEntry
	for i, t := range conn.tracks {
		codec := t.remote.Codec()
		var entry webm.TrackEntry
		switch t.remote.Codec().Name {
		case webrtc.Opus:
			entry = webm.TrackEntry{
				Name:        "Audio",
				TrackNumber: uint64(i + 1),
				CodecID:     "A_OPUS",
				TrackType:   2,
				Audio: &webm.Audio{
					SamplingFrequency: float64(codec.ClockRate),
					Channels:          uint64(codec.Channels),
				},
			}
		case webrtc.VP8:
			entry = webm.TrackEntry{
				Name:        "Video",
				TrackNumber: uint64(i + 1),
				CodecID:     "V_VP8",
				TrackType:   1,
				Video: &webm.Video{
					PixelWidth:  uint64(width),
					PixelHeight: uint64(height),
				},
			}
		default:
			return errors.New("unknown track type")
		}
		entries = append(entries, entry)
	}

	err := conn.reopen()
	if err != nil {
		return err
	}

	writers, err := webm.NewSimpleBlockWriter(conn.file, entries)
	if err != nil {
		conn.file.Close()
		conn.file = nil
		return err
	}

	if len(writers) != len(conn.tracks) {
		conn.file.Close()
		conn.file = nil
		return errors.New("unexpected number of writers")
	}

	conn.width = width
	conn.height = height

	for i, t := range conn.tracks {
		t.writer = writers[i]
	}
	return nil
}

func (down *diskConn) GetMaxBitrate(now uint64) uint64 {
	return ^uint64(0)
}

func (t *diskTrack) Accumulate(bytes uint32) {
	return
}