telegabber/calls/audio/xmpp_to_tg.go

package audio

import (
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pion/webrtc/v4"
	log "github.com/sirupsen/logrus"
)

type XmppToTgOptions struct {
	Ntg          NtgClient
	ChatID       int64  // P2P: equals the Telegram peer's user_id
	PrimePackets uint16 // buffered before playout starts; default 2 (~40ms)
	Logger       *log.Entry
}

// xmpp->tg half: pion RemoteTrack RTP -> opus decode -> ntgcalls
// SendMicrophonePCM. One reader goroutine per HandleTrack, plus one
// playout ticker spawned by Start.
type XmppToTg struct {
	opts XmppToTgOptions
	log  *log.Entry

	dec     *Decoder
	playout *Playout

	stop    chan struct{}
	wg      sync.WaitGroup
	pcmBuf  []int16 // decoder output (worst-case 120ms stereo)
	silence []byte  // 10ms silence for underrun

	closed  atomic.Bool
	started atomic.Bool
	// false until Start runs. Audio between pion OnTrack (during
	// SetRemoteDescription) and bridge OnEstablished is drained from
	// pion's per-track queue but NOT enqueued - otherwise every packet
	// from the setup window would back up and the ticker would replay
	// it at real time, manifesting as call-long lag.
	livePush         atomic.Bool
	preLiveDiscards  atomic.Int64

	capture *captureWriter
}

func NewXmppToTg(opts XmppToTgOptions) (*XmppToTg, error) {
	if opts.Ntg == nil {
		return nil, errors.New("audio: Ntg required")
	}
	if opts.ChatID == 0 {
		return nil, errors.New("audio: ChatID required")
	}
	logger := opts.Logger
	if logger == nil {
		logger = log.WithField("module", "audio")
	}
	logger = logger.WithFields(log.Fields{"dir": "xmpp->tg", "chat_id": opts.ChatID})

	dec, err := NewDecoder()
	if err != nil {
		return nil, err
	}

	h := &XmppToTg{
		opts:    opts,
		log:     logger,
		dec:     dec,
		playout: NewPlayout(opts.PrimePackets),
		pcmBuf:  make([]int16, MaxDecodeSamplesPerChannel*Channels),
		silence: make([]byte, NtgFrameBytes),
		stop:    make(chan struct{}),
	}
	if cap, err := newCaptureWriter(
		fmt.Sprintf("xmpp-to-tg-%d-%d.bin", opts.ChatID, time.Now().Unix()),
		magicRTP,
	); err != nil {
		logger.WithError(err).Warn("capture: xmpp->tg writer disabled")
	} else if cap != nil {
		logger.Info("capture: xmpp->tg RTP capture enabled")
		h.capture = cap
	}
	return h, nil
}

// flips on the ntgcalls Capture source, the playout ticker, and livePush
func (h *XmppToTg) Start() error {
	if !h.started.CompareAndSwap(false, true) {
		return nil
	}
	if err := h.opts.Ntg.SetExternalMicrophone(h.opts.ChatID, true, SampleRate, Channels); err != nil {
		return fmt.Errorf("SetExternalMicrophone(capture): %w", err)
	}
	h.livePush.Store(true)
	h.wg.Add(1)
	go h.tickerLoop()
	h.log.WithField("pre_live_discards", h.preLiveDiscards.Load()).
		Info("xmpp->tg half started")
	return nil
}

// spawns a reader for each pc.OnTrack
func (h *XmppToTg) HandleTrack(t *webrtc.TrackRemote) {
	if t == nil || t.Kind() != webrtc.RTPCodecTypeAudio {
		return
	}
	h.log.WithFields(log.Fields{
		"codec": t.Codec().MimeType,
		"ssrc":  t.SSRC(),
	}).Info("xmpp->tg: remote audio track received from pion")
	h.wg.Add(1)
	go h.reader(t)
}

// stops both pumps; ntgcalls Capture+Playback are released by Call.Close
func (h *XmppToTg) Close() error {
	if !h.closed.CompareAndSwap(false, true) {
		return nil
	}
	close(h.stop)
	h.wg.Wait()
	h.capture.close()
	h.log.Info("xmpp->tg half closed")
	return nil
}

func (h *XmppToTg) reader(t *webrtc.TrackRemote) {
	defer h.wg.Done()
	// negotiated PT for this track; peer is free to pick anything in 96-127,
	// hardcoding 111 would silently drop opus from peers that picked another
	expectedPT := uint8(t.PayloadType())
	for {
		select {
		case <-h.stop:
			return
		default:
		}
		pkt, _, err := t.ReadRTP()
		if err != nil {
			return
		}
		if pkt.PayloadType != expectedPT {
			// CN/DTMF/RED can share an SSRC; libopus rejects them
			continue
		}
		if !h.livePush.Load() {
			// pre-Start: drain pion's per-track queue but don't enqueue;
			// see livePush docstring for the lag rationale
			h.preLiveDiscards.Add(1)
			continue
		}
		h.capture.writeRTP(pkt.SSRC, pkt.SequenceNumber, pkt.Timestamp, pkt.PayloadType, pkt.Payload)
		h.playout.Push(pkt)
	}
}

// cross-tick playout state; separated so tests can drive processTick
// step by step without a time.Ticker
type tickState struct {
	lastDecodedSamples int      // samples/ch from last Decode; DecodePLC output matches
	pending10ms        [][]byte // 10ms PCM chunks awaiting shipment, one per tick
}

func newTickState() *tickState {
	return &tickState{lastDecodedSamples: SamplesPerOpusFrame}
}

func (h *XmppToTg) tickerLoop() {
	defer h.wg.Done()
	ticker := time.NewTicker(NtgFrameMs * time.Millisecond)
	defer ticker.Stop()

	st := newTickState()
	for {
		select {
		case <-h.stop:
			return
		case <-ticker.C:
		}
		h.processTick(st)
	}
}

// one iteration of the playout loop
func (h *XmppToTg) processTick(st *tickState) {
	if len(st.pending10ms) == 0 {
		pkt, plcCount, err := h.playout.PopOrSkip()
		if err != nil {
			// empty or still priming - feed silence
			h.sendToNtg(h.silence)
			return
		}
		for i := 0; i < plcCount; i++ {
			plcOut := h.pcmBuf[:st.lastDecodedSamples*Channels]
			if err := h.dec.DecodePLC(plcOut); err != nil {
				h.log.WithError(err).Warn("DecodePLC failed")
				break
			}
			st.pending10ms = appendChunks10ms(st.pending10ms, plcOut)
		}
		n, err := h.dec.Decode(pkt.Payload, h.pcmBuf)
		if err != nil {
			h.log.WithError(err).Warn("opus decode failed")
			h.sendToNtg(h.silence)
			return
		}
		st.lastDecodedSamples = n
		st.pending10ms = appendChunks10ms(st.pending10ms, h.pcmBuf[:n*Channels])
	}
	chunk := st.pending10ms[0]
	st.pending10ms = st.pending10ms[1:]
	h.sendToNtg(chunk)
}

// slice interleaved stereo PCM into 10ms chunks; partial trailing chunks
// are dropped (libopus's 2.5/5ms frames aren't multiples of 10ms)
func appendChunks10ms(dst [][]byte, pcm []int16) [][]byte {
	samplesPerChunk := SamplesPerNtgFrame * Channels
	for off := 0; off+samplesPerChunk <= len(pcm); off += samplesPerChunk {
		buf := make([]byte, NtgFrameBytes)
		PCMInt16ToBytes(pcm[off:off+samplesPerChunk], buf)
		dst = append(dst, buf)
	}
	return dst
}

func (h *XmppToTg) sendToNtg(pcm []byte) {
	if err := h.opts.Ntg.SendMicrophonePCM(h.opts.ChatID, pcm); err != nil {
		h.log.WithError(err).Debug("SendMicrophonePCM failed")
	}
}