neko/server/internal/webrtc/webrtc.go

package webrtc

import (
	"encoding/json"
	"fmt"
	"io"
	"strings"

	"github.com/pion/interceptor"
	"github.com/pion/rtcp"
	"github.com/pion/webrtc/v3"
	"github.com/pion/webrtc/v3/pkg/media"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"

	"n.eko.moe/neko/internal/types"
	"n.eko.moe/neko/internal/types/config"
)

func New(sessions types.SessionManager, remote types.RemoteManager, config *config.WebRTC) *WebRTCManager {
	return &WebRTCManager{
		logger:   log.With().Str("module", "webrtc").Logger(),
		remote:   remote,
		sessions: sessions,
		config:   config,
	}
}

type WebRTCManager struct {
	logger     zerolog.Logger
	videoTrack *webrtc.TrackLocalStaticSample
	audioTrack *webrtc.TrackLocalStaticSample
	videoCodec webrtc.RTPCodecParameters
	audioCodec webrtc.RTPCodecParameters
	sessions   types.SessionManager
	remote     types.RemoteManager
	config     *config.WebRTC
}

func (manager *WebRTCManager) Start() {
	var err error
	manager.audioTrack, manager.audioCodec, err = manager.createTrack(manager.remote.AudioCodec())
	if err != nil {
		manager.logger.Panic().Err(err).Msg("unable to create audio track")
	}

	manager.remote.OnAudioFrame(func(sample types.Sample) {
		if err := manager.audioTrack.WriteSample(media.Sample(sample)); err != nil && err != io.ErrClosedPipe {
			manager.logger.Warn().Err(err).Msg("audio pipeline failed to write")
		}
	})

	manager.videoTrack, manager.videoCodec, err = manager.createTrack(manager.remote.VideoCodec())
	if err != nil {
		manager.logger.Panic().Err(err).Msg("unable to create video track")
	}

	manager.remote.OnVideoFrame(func(sample types.Sample) {
		if err := manager.videoTrack.WriteSample(media.Sample(sample)); err != nil && err != io.ErrClosedPipe {
			manager.logger.Warn().Err(err).Msg("video pipeline failed to write")
		}
	})

	manager.logger.Info().
		Str("ice_lite", fmt.Sprintf("%t", manager.config.ICELite)).
		Str("ice_servers", strings.Join(manager.config.ICEServers, ",")).
		Str("ephemeral_port_range", fmt.Sprintf("%d-%d", manager.config.EphemeralMin, manager.config.EphemeralMax)).
		Str("nat_ips", strings.Join(manager.config.NAT1To1IPs, ",")).
		Msgf("webrtc starting")
}

func (manager *WebRTCManager) Shutdown() error {
	manager.logger.Info().Msgf("webrtc shutting down")
	return nil
}

func (manager *WebRTCManager) CreatePeer(id string, session types.Session) (string, bool, []string, error) {
	configuration := &webrtc.Configuration{
		ICEServers: []webrtc.ICEServer{
			{
				URLs: manager.config.ICEServers,
			},
		},
		SDPSemantics: webrtc.SDPSemanticsUnifiedPlanWithFallback,
	}

	settings := webrtc.SettingEngine{
		LoggerFactory: loggerFactory{
			logger: manager.logger,
		},
	}

	if manager.config.ICELite {
		configuration = &webrtc.Configuration{
			SDPSemantics: webrtc.SDPSemanticsUnifiedPlanWithFallback,
		}
		settings.SetLite(true)
	}

	settings.SetEphemeralUDPPortRange(manager.config.EphemeralMin, manager.config.EphemeralMax)
	settings.SetNAT1To1IPs(manager.config.NAT1To1IPs, webrtc.ICECandidateTypeHost)
	settings.SetSRTPReplayProtectionWindow(512)

	// Create MediaEngine based off sdp
	engine := webrtc.MediaEngine{}

	engine.RegisterCodec(manager.audioCodec, webrtc.RTPCodecTypeAudio)
	engine.RegisterCodec(manager.videoCodec, webrtc.RTPCodecTypeVideo)

	i := &interceptor.Registry{}
	if err := webrtc.RegisterDefaultInterceptors(&engine, i); err != nil {
		return "", manager.config.ICELite, manager.config.ICEServers, err
	}

	// Create API with MediaEngine and SettingEngine
	api := webrtc.NewAPI(webrtc.WithMediaEngine(&engine), webrtc.WithSettingEngine(settings), webrtc.WithInterceptorRegistry(i))

	// Create new peer connection
	connection, err := api.NewPeerConnection(*configuration)
	if err != nil {
		return "", manager.config.ICELite, manager.config.ICEServers, err
	}
	negotiated := true
	connection.CreateDataChannel("data", &webrtc.DataChannelInit{
		Negotiated: &negotiated,
	})
	connection.OnDataChannel(func(d *webrtc.DataChannel) {
		d.OnMessage(func(msg webrtc.DataChannelMessage) {
			if err = manager.handle(id, msg); err != nil {
				manager.logger.Warn().Err(err).Msg("data handle failed")
			}
		})
	})

	// Set the handler for ICE connection state
	// This will notify you when the peer has connected/disconnected
	connection.OnICEConnectionStateChange(func(connectionState webrtc.ICEConnectionState) {
		fmt.Printf("Connection State has changed %s \n", connectionState.String())
	})

	rtpSender, viderr := connection.AddTrack(manager.videoTrack)
	if viderr != nil {
		return "", manager.config.ICELite, manager.config.ICEServers, viderr
	}

	if _, err = connection.AddTrack(manager.audioTrack); err != nil {
		return "", manager.config.ICELite, manager.config.ICEServers, err
	}

	description, err := connection.CreateOffer(nil)
	if err != nil {
		return "", manager.config.ICELite, manager.config.ICEServers, err
	}

	err = connection.SetLocalDescription(description)
	if err != nil {
		panic(err)
	}

	connection.OnConnectionStateChange(func(state webrtc.PeerConnectionState) {
		switch state {
		case webrtc.PeerConnectionStateDisconnected:
		case webrtc.PeerConnectionStateFailed:
			manager.logger.Info().Str("id", id).Msg("peer disconnected")
			manager.sessions.Destroy(id)
			break
		case webrtc.PeerConnectionStateConnected:
			manager.logger.Info().Str("id", id).Msg("peer connected")
			if err = session.SetConnected(true); err != nil {
				manager.logger.Warn().Err(err).Msg("unable to set connected on peer")
				manager.sessions.Destroy(id)
			}
			break
		}
	})

	connection.OnICECandidate(func(i *webrtc.ICECandidate) {
		if i != nil {
			candidateString, err := json.Marshal(i.ToJSON())
			if err != nil {
				manager.logger.Info().Msg("error")
				return
			}

			if err = session.SignalCandidate(string(candidateString));err != nil {
				manager.logger.Info().Msg("err")
				return
			}
		}
	})


	// Read incoming RTCP packets
	// Before these packets are retuned they are processed by interceptors. For things
	// like NACK this needs to be called.
	go func() {
		rtcpBuf := make([]byte, 1500)
		for {
			n, _, rtcpErr := rtpSender.Read(rtcpBuf)
			if rtcpErr != nil {
				return
			}
			ps, err := rtcp.Unmarshal(rtcpBuf[:n])
				if err != nil {
				log.Printf("Unmarshal RTCP: %v", err)
				continue
			}
			for _, p := range ps {
				switch p.(type) {
				case *rtcp.TransportLayerNack:
					manager.logger.Info().Msg("got a nack")
				}
			}
		}
	}()


	if err := session.SetPeer(&Peer{
		id:            id,
		api:           api,
		engine:        &engine,
		manager:       manager,
		settings:      &settings,
		connection:    connection,
		configuration: configuration,
	}); err != nil {
		return "", manager.config.ICELite, manager.config.ICEServers, err
	}

	return description.SDP, manager.config.ICELite, manager.config.ICEServers, nil
}

func (m *WebRTCManager) createTrack(codecName string) (*webrtc.TrackLocalStaticSample, webrtc.RTPCodecParameters, error) {
	var codec webrtc.RTPCodecParameters
	var fb  []webrtc.RTCPFeedback
	var fba []webrtc.RTCPFeedback
	fb = []webrtc.RTCPFeedback{
		{"goog-remb", ""},
		{"nack", ""},
		{"nack", "pli"},
		{"ccm", "fir"},
	}
	fba = []webrtc.RTCPFeedback{}

	switch codecName {
	case "VP8":
		codec = webrtc.RTPCodecParameters{ RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: "video/VP8", ClockRate: 90000, Channels: 0, SDPFmtpLine: "", RTCPFeedback: fb}, PayloadType: 96,}
	case "VP9":
		codec = webrtc.RTPCodecParameters{ RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: "video/VP9", ClockRate: 90000, Channels: 0, SDPFmtpLine: "", RTCPFeedback: fb}, PayloadType: 98,}
	case "H264":
		codec = webrtc.RTPCodecParameters{ RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: "video/H264", ClockRate: 90000, Channels: 0, SDPFmtpLine: "level-asymmetry-allowed=1;packetization-mode=1;profile-level-id=42001f", RTCPFeedback: fb}, PayloadType: 102,}
	case "Opus":
		codec = webrtc.RTPCodecParameters{ RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: "audio/opus", ClockRate: 48000, Channels: 2, SDPFmtpLine: "", RTCPFeedback: fba}, PayloadType: 111,}
	case "G722":
		codec = webrtc.RTPCodecParameters{ RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: "audio/G722", ClockRate: 8000, Channels: 0, SDPFmtpLine: "", RTCPFeedback: fba}, PayloadType: 9,}
	case "PCMU":
		codec = webrtc.RTPCodecParameters{ RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: "audio/PCMU", ClockRate: 8000, Channels: 0, SDPFmtpLine: "", RTCPFeedback: fba}, PayloadType: 0,}
	case "PCMA":
		codec = webrtc.RTPCodecParameters{ RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: "audio/PCMA", ClockRate: 8000, Channels: 0, SDPFmtpLine: "", RTCPFeedback: fba}, PayloadType: 8,}
	default:
		return nil, codec, fmt.Errorf("unknown codec %s", codecName)
	}

	track, err := webrtc.NewTrackLocalStaticSample(codec.RTPCodecCapability, "stream", "stream")
	if err != nil {
		return nil, codec, err
	}

	return track, codec, nil
}