379 lines
9.4 KiB
Go
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)
|
|
}
|
|
}
|