first commit
This commit is contained in:
378
main.go
Normal file
378
main.go
Normal file
@ -0,0 +1,378 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user