internal/app/livestream/rtmp.go

Rama main · réplica de lectura del repositorio original.

package livestream

import (
	"context"
	"io"
	"net"

	rtmp "github.com/yutopp/go-rtmp"
	rtmpmsg "github.com/yutopp/go-rtmp/message"
	"go.uber.org/zap"
)

// rtmpHandler is the callback for a single RTMP connection: on publish it authenticates
// with the stream key and binds the channel; media tags pass straight through FLV → ffmpeg.
// A connection allows only one publish.
type rtmpHandler struct {
	rtmp.DefaultHandler
	svc    *Service
	conn   net.Conn
	stream *stream
}

func (h *rtmpHandler) OnPublish(_ *rtmp.StreamContext, _ uint32, cmd *rtmpmsg.NetStreamPublish) error {
	if h.stream != nil {
		return errPublishRejected
	}
	st, err := h.svc.startPublish(context.Background(), cmd.PublishingName, h.conn)
	if err != nil {
		h.svc.log.Warn("rtmp publish rejected", zap.Error(err))
		return errPublishRejected
	}
	h.stream = st
	return nil
}

func (h *rtmpHandler) OnSetDataFrame(timestamp uint32, data *rtmpmsg.NetStreamSetDataFrame) error {
	if h.stream == nil {
		return nil
	}
	// onMetaData is passed through to ffmpeg verbatim (optional info; a write failure does not break the stream).
	_ = h.stream.writeTag(18, timestamp, data.Payload)
	return nil
}

func (h *rtmpHandler) OnAudio(timestamp uint32, payload io.Reader) error {
	return h.writeMedia(8, timestamp, payload)
}

func (h *rtmpHandler) OnVideo(timestamp uint32, payload io.Reader) error {
	return h.writeMedia(9, timestamp, payload)
}

func (h *rtmpHandler) writeMedia(tagType byte, timestamp uint32, payload io.Reader) error {
	if h.stream == nil {
		return nil
	}
	body, err := io.ReadAll(payload)
	if err != nil {
		return err
	}
	return h.stream.writeTag(tagType, timestamp, body)
}

func (h *rtmpHandler) OnClose() {
	if h.stream != nil {
		h.svc.endPublish(h.stream)
		h.stream = nil
	}
}