Files
2026-05-24 00:46:14 -04:00

1026 lines
20 KiB
Go

package diskwriter
import (
crand "crypto/rand"
"encoding/binary"
"encoding/hex"
"errors"
"fmt"
"io"
"log"
"math"
"net"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"sync"
"time"
"github.com/pion/rtp"
"git.stormux.org/storm/skald/conn"
"git.stormux.org/storm/skald/hall"
)
const (
sampleRate = 48000
channels = 2
frameDuration = 20 * time.Millisecond
samplesPerFrame = sampleRate / 1000 * 20
bytesPerFrame = samplesPerFrame * channels * 2
sourceBufferSize = 16
encoderBufferSize = 50
sourceStaleAfter = 2 * time.Second
activeSampleLevel = 512
limiterThreshold = 28000
limiterRatio = 4
)
var Directory string
var activeRecordings sync.Map
type Client struct {
hall *hall.Hall
id string
mu sync.Mutex
down map[string]*diskConn
closed bool
recorder *recorder
}
func newId() string {
b := make([]byte, 16)
crand.Read(b)
return hex.EncodeToString(b)
}
func New(g *hall.Hall) *Client {
return &Client{hall: g, id: newId()}
}
func (client *Client) Start() error {
client.mu.Lock()
defer client.mu.Unlock()
if client.closed {
return errors.New("disk client is closed")
}
if client.recorder != nil {
return nil
}
directory := filepath.Join(Directory, client.hall.Name())
err := os.MkdirAll(directory, 0700)
if err != nil {
return err
}
r, err := newRecorder(client, directory)
if err != nil {
return err
}
client.recorder = r
return nil
}
func (client *Client) Hall() *hall.Hall {
return client.hall
}
func (client *Client) Id() string {
return client.id
}
func (client *Client) Username() string {
return "RECORDING"
}
func (client *Client) SetUsername(string) {
}
func (client *Client) SetPermissions(perms []string) {
}
func (client *Client) Permissions() []string {
return []string{"system"}
}
func (client *Client) Data() map[string]interface{} {
return nil
}
func (client *Client) PushClient(hall, kind, id, username string, perms []string, data map[string]interface{}) error {
return nil
}
func (client *Client) RequestConns(target hall.Client, g *hall.Hall, id string) 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
if client.recorder != nil {
err := client.recorder.Close()
client.recorder = nil
return err
}
return nil
}
func (client *Client) Kick(id string, user *string, message string) error {
err := client.Close()
hall.DelClient(client)
return err
}
func (client *Client) Addr() net.Addr {
return nil
}
func (client *Client) Joined(hall, kind string) error {
return nil
}
func (client *Client) PushConn(g *hall.Hall, id string, up conn.Up, tracks []conn.UpTrack, replace string) error {
if client.hall != g {
return nil
}
client.mu.Lock()
defer client.mu.Unlock()
if client.closed {
return errors.New("disk client is closed")
}
if client.recorder == nil {
return errors.New("recorder is not started")
}
if replace != "" {
rp := client.down[replace]
if rp != nil {
rp.Close()
delete(client.down, replace)
} else {
log.Printf("Disk writer: replacing unknown connection")
}
}
old := client.down[id]
if old != nil {
old.Close()
delete(client.down, id)
}
if up == nil {
return nil
}
if client.down == nil {
client.down = make(map[string]*diskConn)
}
down, err := newDiskConn(client, up, tracks)
if err != nil {
g.WallOps("Write to disk: " + err.Error())
return err
}
if down == nil {
return nil
}
client.down[up.Id()] = down
return nil
}
type recorder struct {
client *Client
path string
cmd *exec.Cmd
stdin io.WriteCloser
done chan error
writeCh chan []byte
writerDone chan error
mu sync.Mutex
sources map[*diskTrack]*recordingSource
metrics RecorderMetrics
closed bool
}
type recordingSource struct {
ch chan []byte
lastFrame time.Time
decoderStarted bool
decoderStopped bool
staleLogged bool
}
type RecorderMetrics struct {
Started time.Time
Duration time.Duration
ActiveSources int
DecoderStarts uint64
DecoderStops uint64
DroppedSourceFrames uint64
StaleSourceTicks uint64
DroppedMixedFrames uint64
EncoderWriteErrors uint64
MixerFrames uint64
}
type SyntheticRecordingResult struct {
Path string
Metrics RecorderMetrics
}
func newRecorder(client *Client, directory string) (*recorder, error) {
file, err := openDiskFile(directory, "ogg")
if err != nil {
return nil, err
}
path := file.Name()
if err := file.Close(); err != nil {
return nil, err
}
cmd := exec.Command(
"ffmpeg",
"-y",
"-hide_banner",
"-loglevel", "error",
"-f", "s16le",
"-ar", fmt.Sprint(sampleRate),
"-ac", fmt.Sprint(channels),
"-i", "pipe:0",
"-c:a", "libopus",
"-b:a", "128k",
path,
)
stdin, err := cmd.StdinPipe()
if err != nil {
return nil, err
}
stderr, err := cmd.StderrPipe()
if err != nil {
return nil, err
}
if err := cmd.Start(); err != nil {
return nil, fmt.Errorf("could not start ffmpeg: %w", err)
}
go drainStderr("ffmpeg encoder", stderr)
r := &recorder{
client: client,
path: path,
cmd: cmd,
stdin: stdin,
done: make(chan error, 1),
writeCh: make(chan []byte, encoderBufferSize),
writerDone: make(chan error, 1),
sources: make(map[*diskTrack]*recordingSource),
metrics: RecorderMetrics{Started: time.Now()},
}
activeRecordings.Store(path, struct{}{})
go r.writeEncoder()
go r.mix()
return r, nil
}
func (r *recorder) addSource(t *diskTrack) *recordingSource {
r.mu.Lock()
defer r.mu.Unlock()
if r.closed {
return nil
}
source := &recordingSource{
ch: make(chan []byte, sourceBufferSize),
lastFrame: time.Now(),
}
r.sources[t] = source
return source
}
func (r *recorder) noteDecoderStarted(t *diskTrack) {
r.mu.Lock()
defer r.mu.Unlock()
source := r.sources[t]
if source == nil || source.decoderStarted {
return
}
source.decoderStarted = true
r.metrics.DecoderStarts++
log.Printf("Recording %s: decoder started, active sources %d",
filepath.Base(r.path), len(r.sources))
}
func (r *recorder) removeSource(t *diskTrack) {
r.mu.Lock()
defer r.mu.Unlock()
source := r.sources[t]
if source == nil {
return
}
if source.decoderStarted && !source.decoderStopped {
source.decoderStopped = true
r.metrics.DecoderStops++
}
delete(r.sources, t)
log.Printf("Recording %s: source removed, active sources %d",
filepath.Base(r.path), len(r.sources))
}
func (r *recorder) noteDecoderStopped(t *diskTrack, err error) {
r.mu.Lock()
defer r.mu.Unlock()
source := r.sources[t]
if source == nil || source.decoderStopped {
return
}
source.decoderStopped = true
r.metrics.DecoderStops++
delete(r.sources, t)
if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) {
log.Printf("Recording %s: decoder stopped: %v",
filepath.Base(r.path), err)
} else {
log.Printf("Recording %s: decoder stopped, active sources %d",
filepath.Base(r.path), len(r.sources))
}
}
func (r *recorder) pushSourceFrame(t *diskTrack, frame []byte) bool {
r.mu.Lock()
defer r.mu.Unlock()
source := r.sources[t]
if source == nil {
return false
}
now := time.Now()
source.lastFrame = now
source.staleLogged = false
select {
case source.ch <- frame:
default:
<-source.ch
source.ch <- frame
r.metrics.DroppedSourceFrames++
}
return true
}
func (r *recorder) mix() {
ticker := time.NewTicker(frameDuration)
defer ticker.Stop()
silence := make([]byte, bytesPerFrame)
for {
select {
case err := <-r.writerDone:
if err != nil {
r.mu.Lock()
r.metrics.EncoderWriteErrors++
r.mu.Unlock()
}
r.done <- err
return
case <-ticker.C:
if r.isClosed() {
close(r.writeCh)
err := <-r.writerDone
if err != nil {
r.mu.Lock()
r.metrics.EncoderWriteErrors++
r.mu.Unlock()
}
r.done <- err
return
}
frame := r.nextMixedFrame(time.Now())
if frame == nil {
frame = silence
}
r.queueMixedFrame(frame)
}
}
}
func (r *recorder) isClosed() bool {
r.mu.Lock()
defer r.mu.Unlock()
return r.closed
}
func (r *recorder) nextMixedFrame(now time.Time) []byte {
r.mu.Lock()
frames := make([][]byte, 0, len(r.sources))
for _, source := range r.sources {
select {
case frame := <-source.ch:
if len(frame) >= bytesPerFrame {
frames = append(frames, frame)
}
default:
if now.Sub(source.lastFrame) >= sourceStaleAfter && !source.staleLogged {
source.staleLogged = true
r.metrics.StaleSourceTicks++
log.Printf("Recording %s: source has no decoded frames for %v",
filepath.Base(r.path), sourceStaleAfter)
}
}
}
r.metrics.MixerFrames++
r.mu.Unlock()
return mixPCMFrames(frames)
}
func (r *recorder) queueMixedFrame(frame []byte) {
select {
case r.writeCh <- frame:
default:
r.mu.Lock()
r.metrics.DroppedMixedFrames++
r.mu.Unlock()
log.Printf("Recording %s: encoder queue full, dropping mixed frame",
filepath.Base(r.path))
}
}
func (r *recorder) writeEncoder() {
for frame := range r.writeCh {
if _, err := r.stdin.Write(frame); err != nil {
r.writerDone <- err
return
}
}
r.writerDone <- nil
}
func mixFrame(sources []chan []byte) []byte {
frames := make([][]byte, 0, len(sources))
for _, ch := range sources {
var frame []byte
select {
case frame = <-ch:
default:
}
if len(frame) >= bytesPerFrame {
frames = append(frames, frame)
}
}
return mixPCMFrames(frames)
}
func mixPCMFrames(frames [][]byte) []byte {
if len(frames) == 0 {
return nil
}
var mixed []int32
activeFrames := 0
for _, frame := range frames {
if mixed == nil {
mixed = make([]int32, samplesPerFrame*channels)
}
active := false
for i := range mixed {
sample := int16(binary.LittleEndian.Uint16(frame[i*2:]))
if !active && abs16(sample) >= activeSampleLevel {
active = true
}
mixed[i] += int32(sample)
}
if active {
activeFrames++
}
}
gain := 1.0
if activeFrames > 1 {
gain = 1 / math.Sqrt(float64(activeFrames))
}
out := make([]byte, bytesPerFrame)
for i, sample := range mixed {
sample = int32(math.Round(float64(sample) * gain))
sample = limitSample(sample)
binary.LittleEndian.PutUint16(out[i*2:], uint16(int16(sample)))
}
return out
}
func abs16(sample int16) int16 {
if sample < 0 {
return -sample
}
return sample
}
func limitSample(sample int32) int32 {
if sample > limiterThreshold {
sample = limiterThreshold + (sample-limiterThreshold)/limiterRatio
} else if sample < -limiterThreshold {
sample = -limiterThreshold + (sample+limiterThreshold)/limiterRatio
}
if sample > 32767 {
return 32767
}
if sample < -32768 {
return -32768
}
return sample
}
func (r *recorder) Close() error {
r.mu.Lock()
if r.closed {
r.mu.Unlock()
return nil
}
r.closed = true
r.mu.Unlock()
err := <-r.done
stdinErr := r.stdin.Close()
waitErr := r.cmd.Wait()
activeRecordings.Delete(r.path)
metrics := r.snapshotMetrics()
log.Printf("Recording %s: closed after %v, sources=%d, decoders=%d/%d, source drops=%d, stale ticks=%d, mixed drops=%d, encoder errors=%d",
filepath.Base(r.path),
metrics.Duration,
metrics.ActiveSources,
metrics.DecoderStarts,
metrics.DecoderStops,
metrics.DroppedSourceFrames,
metrics.StaleSourceTicks,
metrics.DroppedMixedFrames,
metrics.EncoderWriteErrors,
)
if err != nil {
return err
}
if stdinErr != nil {
return stdinErr
}
if waitErr != nil {
return waitErr
}
now := time.Now()
if err := os.Chtimes(r.path, now, now); err != nil {
return err
}
return nil
}
func (r *recorder) snapshotMetrics() RecorderMetrics {
r.mu.Lock()
defer r.mu.Unlock()
metrics := r.metrics
metrics.Duration = time.Since(metrics.Started)
metrics.ActiveSources = len(r.sources)
return metrics
}
func RunSyntheticRecording(directory string, sources int, duration time.Duration) (SyntheticRecordingResult, error) {
if sources <= 0 {
return SyntheticRecordingResult{}, errors.New("sources must be greater than zero")
}
if duration <= 0 {
return SyntheticRecordingResult{}, errors.New("duration must be greater than zero")
}
if err := os.MkdirAll(directory, 0700); err != nil {
return SyntheticRecordingResult{}, err
}
r, err := newRecorder(nil, directory)
if err != nil {
return SyntheticRecordingResult{}, err
}
var wg sync.WaitGroup
for i := 0; i < sources; i++ {
track := &diskTrack{}
if r.addSource(track) == nil {
r.Close()
return SyntheticRecordingResult{}, errors.New("recorder is closed")
}
r.noteDecoderStarted(track)
wg.Add(1)
go func(index int, t *diskTrack) {
defer wg.Done()
defer r.removeSource(t)
ticker := time.NewTicker(frameDuration)
defer ticker.Stop()
timer := time.NewTimer(duration)
defer timer.Stop()
frequency := 220 + float64(index)*37
phase := 0.0
for {
select {
case <-timer.C:
return
case <-ticker.C:
r.pushSourceFrame(t, sinePCMFrame(frequency, &phase))
}
}
}(i, track)
}
wg.Wait()
if err := r.Close(); err != nil {
return SyntheticRecordingResult{}, err
}
return SyntheticRecordingResult{
Path: r.path,
Metrics: r.snapshotMetrics(),
}, nil
}
func sinePCMFrame(frequency float64, phase *float64) []byte {
frame := make([]byte, bytesPerFrame)
step := 2 * math.Pi * frequency / sampleRate
for i := 0; i < samplesPerFrame; i++ {
sample := int16(math.Sin(*phase) * 6000)
*phase += step
if *phase >= 2*math.Pi {
*phase -= 2 * math.Pi
}
offset := i * channels * 2
for ch := 0; ch < channels; ch++ {
binary.LittleEndian.PutUint16(
frame[offset+ch*2:], uint16(sample),
)
}
}
return frame
}
type diskConn struct {
client *Client
remote conn.Up
tracks []*diskTrack
}
func (conn *diskConn) Close() error {
conn.remote.DelLocal(conn)
for _, t := range conn.tracks {
t.Close()
}
return nil
}
type diskTrack struct {
remote conn.UpTrack
conn *diskConn
mu sync.Mutex
started bool
closed bool
payload uint8
rtpConn *net.UDPConn
cmd *exec.Cmd
stdin io.WriteCloser
ch chan []byte
}
func newDiskConn(client *Client, up conn.Up, remoteTracks []conn.UpTrack) (*diskConn, error) {
var tracks []conn.UpTrack
for _, remote := range remoteTracks {
codec := remote.Codec().MimeType
if strings.EqualFold(codec, "audio/opus") {
tracks = append(tracks, remote)
} else {
client.hall.WallOps("Unknown codec, " + codec + ", not recording")
}
}
if len(tracks) == 0 {
return nil, nil
}
down := &diskConn{
client: client,
remote: up,
tracks: make([]*diskTrack, 0, len(tracks)),
}
for _, remote := range tracks {
track := &diskTrack{
remote: remote,
conn: down,
}
down.tracks = append(down.tracks, track)
}
for _, t := range down.tracks {
err := t.remote.AddLocal(t)
if err != nil {
log.Printf("Couldn't add disk track: %v", err)
client.hall.WallOps("Couldn't add disk track: " + err.Error())
}
}
err := up.AddLocal(down)
if err != nil {
return nil, err
}
return down, nil
}
func (t *diskTrack) SetCname(string) {
}
func (t *diskTrack) Write(buf []byte) (int, error) {
var p rtp.Packet
if err := p.Unmarshal(buf); err != nil {
log.Printf("Diskwriter: %v", err)
return 0, nil
}
t.mu.Lock()
defer t.mu.Unlock()
if t.closed {
return 0, conn.ErrConnectionClosed
}
if !t.started {
if err := t.start(p.PayloadType); err != nil {
t.conn.client.hall.WallOps("Write to disk: " + err.Error())
return 0, err
}
}
if t.payload != p.PayloadType {
return 0, nil
}
if _, err := t.rtpConn.Write(buf); err != nil {
return 0, err
}
return len(buf), nil
}
func (t *diskTrack) start(payload uint8) error {
source := t.conn.client.recorder.addSource(t)
if source == nil {
return errors.New("recorder is closed")
}
port, err := freeUDPPort()
if err != nil {
t.conn.client.recorder.removeSource(t)
return err
}
cmd := exec.Command(
"ffmpeg",
"-hide_banner",
"-loglevel", "error",
"-protocol_whitelist", "file,pipe,udp,rtp",
"-f", "sdp",
"-i", "pipe:0",
"-f", "s16le",
"-ar", fmt.Sprint(sampleRate),
"-ac", fmt.Sprint(channels),
"pipe:1",
)
stdin, err := cmd.StdinPipe()
if err != nil {
t.conn.client.recorder.removeSource(t)
return err
}
stdout, err := cmd.StdoutPipe()
if err != nil {
t.conn.client.recorder.removeSource(t)
return err
}
stderr, err := cmd.StderrPipe()
if err != nil {
t.conn.client.recorder.removeSource(t)
return err
}
if err := cmd.Start(); err != nil {
t.conn.client.recorder.removeSource(t)
return fmt.Errorf("could not start ffmpeg: %w", err)
}
sdp := fmt.Sprintf(
"v=0\r\n"+
"o=- 0 0 IN IP4 127.0.0.1\r\n"+
"s=Skald recording\r\n"+
"c=IN IP4 127.0.0.1\r\n"+
"t=0 0\r\n"+
"m=audio %d RTP/AVP %d\r\n"+
"a=rtpmap:%d opus/48000/2\r\n"+
"a=recvonly\r\n",
port, payload, payload,
)
if _, err := io.WriteString(stdin, sdp); err != nil {
cmd.Process.Kill()
t.conn.client.recorder.removeSource(t)
return err
}
stdin.Close()
addr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: port}
rtpConn, err := net.DialUDP("udp", nil, addr)
if err != nil {
cmd.Process.Kill()
t.conn.client.recorder.removeSource(t)
return err
}
t.started = true
t.payload = payload
t.rtpConn = rtpConn
t.cmd = cmd
t.stdin = stdin
t.ch = source.ch
t.conn.client.recorder.noteDecoderStarted(t)
go drainStderr("ffmpeg decoder", stderr)
go t.readPCM(stdout)
return nil
}
func (t *diskTrack) readPCM(r io.Reader) {
for {
buf := make([]byte, bytesPerFrame)
_, err := io.ReadFull(r, buf)
if err != nil {
if t.conn.client.recorder != nil {
t.conn.client.recorder.noteDecoderStopped(t, err)
}
return
}
if t.conn.client.recorder != nil {
t.conn.client.recorder.pushSourceFrame(t, buf)
}
}
}
func (t *diskTrack) Close() {
t.mu.Lock()
defer t.mu.Unlock()
if t.closed {
return
}
t.closed = true
t.remote.DelLocal(t)
if t.rtpConn != nil {
t.rtpConn.Close()
t.rtpConn = nil
}
if t.cmd != nil && t.cmd.Process != nil {
t.cmd.Process.Kill()
t.cmd.Wait()
}
if t.conn.client.recorder != nil {
t.conn.client.recorder.removeSource(t)
}
}
func (t *diskTrack) SetTimeOffset(ntp uint64, rtp uint32) {
}
func (t *diskTrack) GetMaxBitrate() (uint64, int, int) {
return ^uint64(0), -1, -1
}
func freeUDPPort() (int, error) {
addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
if err != nil {
return 0, err
}
conn, err := net.ListenUDP("udp", addr)
if err != nil {
return 0, err
}
defer conn.Close()
return conn.LocalAddr().(*net.UDPAddr).Port, nil
}
var replacer = strings.NewReplacer(
"/", "-slash-",
"\\", "-backslash-",
)
// sanitise sanitises a string so it can be safely used in a filename.
func sanitise(s string) string {
return replacer.Replace(s)
}
func openDiskFile(directory, extension string) (*os.File, error) {
filenameFormat := "20060102-150405"
if runtime.GOOS == "windows" {
filenameFormat = "20060102-150405"
}
filename := "skald-recording-" + time.Now().Format(filenameFormat)
for counter := 0; counter < 100; counter++ {
var fn string
if counter == 0 {
fn = fmt.Sprintf("%v.%v", filename, extension)
} else {
fn = fmt.Sprintf("%v-%02d.%v",
filename, counter, extension,
)
}
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 !errors.Is(err, os.ErrExist) {
return nil, err
}
}
return nil, errors.New("couldn't create file")
}
func CleanupOldRecordings(retention time.Duration) error {
if retention <= 0 {
return nil
}
cutoff := time.Now().Add(-retention)
err := filepath.WalkDir(
Directory,
func(path string, d os.DirEntry, err error) error {
if err != nil {
log.Printf("Recording cleanup %v: %v", path, err)
return nil
}
if d.IsDir() {
return nil
}
name := d.Name()
if !strings.HasPrefix(name, "skald-recording-") ||
filepath.Ext(name) != ".ogg" {
return nil
}
if _, ok := activeRecordings.Load(path); ok {
return nil
}
info, err := d.Info()
if err != nil {
log.Printf("Recording cleanup %v: %v", path, err)
return nil
}
if info.ModTime().After(cutoff) {
return nil
}
if err := os.Remove(path); err != nil {
log.Printf("Recording cleanup %v: %v", path, err)
}
return nil
},
)
if errors.Is(err, os.ErrNotExist) {
return nil
}
return err
}
func drainStderr(prefix string, r io.Reader) {
buf := make([]byte, 1024)
for {
n, err := r.Read(buf)
if n > 0 {
msg := strings.TrimSpace(string(buf[:n]))
if msg != "" {
log.Printf("%s: %s", prefix, msg)
}
}
if err != nil {
return
}
}
}