Files
2025-08-19 10:38:15 -04:00

379 lines
9.4 KiB
Go

package main
import (
"bufio"
"encoding/binary"
"fmt"
"log/slog"
"os"
"sync"
"github.com/bluenviron/gortsplib/v4"
"github.com/bluenviron/gortsplib/v4/pkg/base"
"github.com/bluenviron/gortsplib/v4/pkg/description"
"github.com/bluenviron/gortsplib/v4/pkg/format"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/g711"
"github.com/pion/rtp"
"gopkg.in/yaml.v3"
)
// Config holds RTSP URL and output file path
type Config struct {
RTSPURL string `yaml:"rtsp_url"`
OutputFile string `yaml:"output_file"`
}
// Event represents lifecycle events (e.g., source ready, recording started)
type Event struct {
Type string
Data interface{}
}
// Path manages a stream, distributing packets to readers/recorders
type Path struct {
audioMedia *description.Media
audioFormat format.Format
readers map[string]chan *Data
readerMu sync.RWMutex
eventChan chan Event
}
// Data holds an RTP packet and its media type
type Data struct {
Media *description.Media
Pkt *rtp.Packet
}
// Recorder saves packets to a file
type Recorder struct {
name string
path *Path
file *os.File
writer *bufio.Writer
isMuLaw bool // true for μ-law, false for A-law
dataChan chan *Data
terminate chan struct{}
sampleRate int
channels int
}
// NewPath creates a stream hub
func NewPath(eventChan chan Event) *Path {
return &Path{
readers: make(map[string]chan *Data),
eventChan: eventChan,
}
}
// AddReader adds a reader (e.g., recorder) with its own channel
func (p *Path) AddReader(name string, ch chan *Data) {
p.readerMu.Lock()
defer p.readerMu.Unlock()
p.readers[name] = ch
}
// RemoveReader removes a reader
func (p *Path) RemoveReader(name string) {
p.readerMu.Lock()
defer p.readerMu.Unlock()
if ch, ok := p.readers[name]; ok {
close(ch)
delete(p.readers, name)
}
}
// DistributeData sends packets to all readers
func (p *Path) DistributeData(data *Data) {
p.readerMu.RLock()
defer p.readerMu.RUnlock()
for _, ch := range p.readers {
select {
case ch <- data:
default:
slog.Warn("Reader channel full, dropping packet")
}
}
}
// SourceReady sets up the path with media info
func (p *Path) SourceReady(audioMedia *description.Media, audioFormat format.Format) {
p.audioMedia = audioMedia
p.audioFormat = audioFormat
p.eventChan <- Event{Type: "source_ready", Data: audioMedia}
}
// NewRecorder creates a recorder that subscribes to the path
func NewRecorder(path *Path, outputFile string, isMuLaw bool) (*Recorder, error) {
file, err := os.Create(outputFile)
if err != nil {
return nil, fmt.Errorf("create file: %w", err)
}
r := &Recorder{
name: "recorder",
path: path,
file: file,
writer: bufio.NewWriter(file),
isMuLaw: isMuLaw,
dataChan: make(chan *Data, 100), // Buffered channel
terminate: make(chan struct{}),
sampleRate: 8000, // Fixed for G.711
channels: 1, // Mono
}
path.AddReader(r.name, r.dataChan)
return r, nil
}
// writeWAVHeader writes a basic WAV header for PCM audio
func (r *Recorder) writeWAVHeader(dataSize uint32) error {
header := make([]byte, 44)
copy(header[0:4], "RIFF")
binary.LittleEndian.PutUint32(header[4:8], 36+dataSize) // ChunkSize
copy(header[8:12], "WAVE")
copy(header[12:16], "fmt ")
binary.LittleEndian.PutUint32(header[16:20], 16) // Subchunk1Size (16 for PCM)
binary.LittleEndian.PutUint16(header[20:22], 1) // AudioFormat (1 for PCM)
binary.LittleEndian.PutUint16(header[22:24], uint16(r.channels))
binary.LittleEndian.PutUint32(header[24:28], uint32(r.sampleRate))
binary.LittleEndian.PutUint32(header[28:32], uint32(r.sampleRate*r.channels*2)) // ByteRate
binary.LittleEndian.PutUint16(header[32:34], uint16(r.channels*2)) // BlockAlign
binary.LittleEndian.PutUint16(header[34:36], 16) // BitsPerSample
copy(header[36:40], "data")
binary.LittleEndian.PutUint32(header[40:44], dataSize)
_, err := r.writer.Write(header)
return err
}
// muLawToPCM converts μ-law samples to 16-bit PCM
func muLawToPCM(sample byte) int16 {
const muLawBias = 33
const muLawMax = 32767
var sign int
if sample&0x80 != 0 {
sign = -1
sample = ^sample
} else {
sign = 1
}
sample = sample & 0x7F
exponent := int(sample>>4) & 0x07
mantissa := int(sample & 0x0F)
value := (mantissa << (exponent + 3)) + muLawBias
if exponent > 0 {
value += (1 << (exponent + 2))
}
value -= muLawBias
return int16(sign * value * muLawMax / 8159)
}
// aLawToPCM converts A-law samples to 16-bit PCM
func aLawToPCM(sample byte) int16 {
const aLawMax = 32767
var sign int
if sample&0x80 != 0 {
sign = -1
sample = ^sample
} else {
sign = 1
}
sample = sample & 0x7F
exponent := int(sample>>4) & 0x07
mantissa := int(sample & 0x0F)
value := 0
if exponent == 0 {
value = mantissa
} else {
value = (mantissa << 4) + (1 << (exponent + 2))
}
return int16(sign * value * aLawMax / 4096)
}
// Start runs the recorder's async writing loop
func (r *Recorder) Start() {
// Placeholder for data size; update at end
if err := r.writeWAVHeader(0); err != nil {
slog.Error("Write header error", "error", err)
}
var totalDataSize uint32
go func() {
defer r.file.Close()
defer r.writer.Flush()
for {
select {
case data := <-r.dataChan:
if data.Media.Type == "audio" {
// Unmarshal G.711 payload
var samples []byte
var err error
if r.isMuLaw {
var mu g711.Mulaw
mu.Unmarshal(data.Pkt.Payload)
samples = []byte(mu)
} else {
var al g711.Alaw
al.Unmarshal(data.Pkt.Payload)
samples = []byte(al)
}
if err != nil {
slog.Error("Unmarshal audio error", "error", err)
continue
}
// Convert to 16-bit PCM and write
for _, sample := range samples {
var pcmSample int16
if r.isMuLaw {
pcmSample = muLawToPCM(sample)
} else {
pcmSample = aLawToPCM(sample)
}
var buf [2]byte
binary.LittleEndian.PutUint16(buf[:], uint16(pcmSample))
_, err := r.writer.Write(buf[:])
if err != nil {
slog.Error("Write error", "error", err)
}
totalDataSize += 2
}
}
case <-r.terminate:
// Update WAV header with actual data size
if err := r.writer.Flush(); err != nil {
slog.Error("Flush error", "error", err)
}
if _, err := r.file.Seek(4, 0); err != nil {
slog.Error("Seek error", "error", err)
return
}
var buf [4]byte
binary.LittleEndian.PutUint32(buf[:], 36+totalDataSize)
r.file.Write(buf[:])
if _, err := r.file.Seek(40, 0); err != nil {
slog.Error("Seek error", "error", err)
return
}
binary.LittleEndian.PutUint32(buf[:], totalDataSize)
r.file.Write(buf[:])
return
}
}
}()
}
// Stop terminates the recorder
func (r *Recorder) Stop() {
r.path.RemoveReader(r.name)
close(r.terminate)
}
func main() {
// Initialize slog (default handler for structured logging)
slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo})))
// Step 1: Parse config
configFile, err := os.ReadFile("config.yaml")
if err != nil {
slog.Error("Read config", "error", err)
os.Exit(1)
}
var conf Config
if err := yaml.Unmarshal(configFile, &conf); err != nil {
slog.Error("Unmarshal config", "error", err)
os.Exit(1)
}
if conf.RTSPURL == "" || conf.OutputFile == "" {
slog.Error("RTSP URL or output file missing")
os.Exit(1)
}
// Step 2: Set up event bus
eventChan := make(chan Event, 10)
go func() {
for ev := range eventChan {
slog.Info("Event", "type", ev.Type, "data", ev.Data)
}
}()
// Step 3: Create path (stream hub)
path := NewPath(eventChan)
// Step 4: Set up RTSP client (using provided connection logic)
u, err := base.ParseURL(conf.RTSPURL)
if err != nil {
slog.Error("Failed to parse RTSP URL", "error", err, "url", conf.RTSPURL)
os.Exit(1)
}
c := gortsplib.Client{
Scheme: u.Scheme,
Host: u.Host,
}
// Connect to the server
err = c.Start2()
if err != nil {
slog.Error("Failed to connect to RTSP server", "error", err)
os.Exit(1)
}
defer c.Close()
// Find available medias
desc, _, err := c.Describe(u)
if err != nil {
slog.Error("Failed to describe RTSP stream", "error", err)
os.Exit(1)
}
// Find the G711 media and format
var audioFormat *format.G711
audioMedia := desc.FindFormat(&audioFormat)
if audioMedia == nil {
slog.Error("G711 media not found in RTSP stream")
os.Exit(1)
}
if audioFormat.SampleRate != 8000 || audioFormat.ChannelCount != 1 {
slog.Error("Audio format is not mono 8000Hz G711")
os.Exit(1)
}
_, err = c.Setup(desc.BaseURL, audioMedia, 0, 0)
if err != nil {
slog.Error("Failed to setup RTSP media", "error", err)
os.Exit(1)
}
// Step 5: Set up recorder (choose μ-law or A-law based on format)
isMuLaw := audioFormat.MULaw
recorder, err := NewRecorder(path, conf.OutputFile, isMuLaw)
if err != nil {
slog.Error("Create recorder", "error", err)
os.Exit(1)
}
recorder.Start()
// Step 6: Signal source ready
path.SourceReady(audioMedia, audioFormat)
// Step 7: Read RTP packets and distribute
go func() {
c.OnPacketRTP(audioMedia, audioFormat, func(pkt *rtp.Packet) {
path.DistributeData(&Data{Media: audioMedia, Pkt: pkt})
})
_, err := c.Play(nil)
if err != nil {
slog.Error("Play failed", "error", err)
eventChan <- Event{Type: "source_error", Data: err}
return
}
}()
// Step 8: Wait for errors or interruption
err = c.Wait()
if err != nil {
slog.Error("RTSP client error", "error", err)
recorder.Stop()
c.Close()
close(eventChan)
os.Exit(1)
}
}