rkKVM/http/hw/rtc/webrtc.go

package rtc

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"

	log "github.com/sirupsen/logrus"

	"github.com/pion/webrtc/v4"
)

// https://github.com/pion/example-webrtc-applications/blob/master/sfu-ws/main.go

var rtc *RTC

func Get() *RTC {
	return rtc
}

var ErrPeerClosedConn = errors.New("webrtc: peer closed conn")

type RTC struct {
	l            *net.UDPConn
	peer         *webrtc.PeerConnection
	track        *webrtc.TrackLocalStaticRTP
	sender       *webrtc.RTPSender
	localSession string
}

func (r *RTC) Close() error {
	return r.l.Close()
}

// Read incoming RTCP packets
// Before these packets are returned they are processed by interceptors. For things
// like NACK this needs to be called.
func (r *RTC) Read() {
	rtcpBuf := make([]byte, 1500)
	for {
		if _, _, rtcpErr := r.sender.Read(rtcpBuf); rtcpErr != nil {
			log.Errorf("failed to read RTCP packet: %v", rtcpErr)
			return
		}
	}
}

func Init(host string, port int) (*RTC, error) {
	peerConnection, err := webrtc.NewPeerConnection(webrtc.Configuration{
		ICEServers: []webrtc.ICEServer{
			{
				URLs: []string{"stun:stun.l.google.com:19302"},
			},
		},
	})
	if err != nil {
		return nil, fmt.Errorf("failed to create peer connection: %v", err)
	}

	// Open a UDP Listener for RTP Packets on port 5004
	l, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP(host), Port: port})
	if err != nil {
		return nil, fmt.Errorf("failed to init webrtc listener: %v", err)
	}

	// Increase the UDP receive buffer size
	// Default UDP buffer sizes vary on different operating systems
	bufferSize := 300000 // 300KB
	err = l.SetReadBuffer(bufferSize)
	if err != nil {
		return nil, fmt.Errorf("failed to set read buffer: %v", err)
	}

	track, err := webrtc.NewTrackLocalStaticRTP(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeH264}, "video", "pion")
	if err != nil { // it should never happens
		panic(fmt.Sprintf("failed to create video track: %v", err))
	}

	rtpSender, err := peerConnection.AddTrack(track)
	if err != nil {
		return nil, fmt.Errorf("failed to add track to peer connection: %v", err)
	}

	r := &RTC{
		peer:   peerConnection,
		sender: rtpSender,
		track:  track,
		l:      l,
	}
	rtc = r

	return r, nil
}

func (r *RTC) Handshake(clientSession string) (string, error) {
	// Set the handler for ICE connection state
	// This will notify you when the peer has connected/disconnected
	r.peer.OnICEConnectionStateChange(func(connState webrtc.ICEConnectionState) {
		log.Infof("Connection State has changed %s", connState.String())

		if connState == webrtc.ICEConnectionStateFailed {
			if closeErr := r.peer.Close(); closeErr != nil {
				panic(closeErr)
			}
		}
	})

	// Wait for the offer to be pasted
	offer := webrtc.SessionDescription{}
	decode(clientSession, &offer)
	fmt.Printf("Offer: %+v\n", offer)

	// Set the remote SessionDescription
	if err := r.peer.SetRemoteDescription(offer); err != nil {
		return "", fmt.Errorf("failed to set remote session description: %v", err)
	}

	// Create answer
	answer, err := r.peer.CreateAnswer(nil)
	if err != nil {
		return "", fmt.Errorf("failed to create answer: %v", err)
	}

	// Create channel that is blocked until ICE Gathering is complete
	gatherComplete := webrtc.GatheringCompletePromise(r.peer)

	// Sets the LocalDescription, and starts our UDP listeners
	if err = r.peer.SetLocalDescription(answer); err != nil {
		return "", fmt.Errorf("failed to set local description: %v", err)
	}

	// Block until ICE Gathering is complete, disabling trickle ICE
	// we do this because we only can exchange one signaling message
	// in a production application you should exchange ICE Candidates via OnICECandidate
	<-gatherComplete

	r.localSession = encode(r.peer.LocalDescription())
	return r.localSession, nil
}

func (r *RTC) Listen() error {
	// Read RTP packets forever and send them to the WebRTC Client
	inboundRTPPacket := make([]byte, 1600) // UDP MTU
	for {
		n, _, err := r.l.ReadFrom(inboundRTPPacket)
		if err != nil {
			return fmt.Errorf("error during read: %v", err)
		}

		if _, err = r.track.Write(inboundRTPPacket[:n]); err != nil {
			if errors.Is(err, io.ErrClosedPipe) {
				// The peerConnection has been closed.
				return ErrPeerClosedConn
			}

			return fmt.Errorf("failed to send RTP packet to client: %v", err)
		}
	}
}

// JSON encode + base64 a SessionDescription
func encode(obj *webrtc.SessionDescription) string {
	b, err := json.Marshal(obj)
	if err != nil {
		panic(err)
	}

	return base64.StdEncoding.EncodeToString(b)
}

// Decode a base64 and unmarshal JSON into a SessionDescription
func decode(in string, obj *webrtc.SessionDescription) {
	b, err := base64.StdEncoding.DecodeString(in)
	if err != nil {
		panic(err)
	}

	if err = json.Unmarshal(b, obj); err != nil {
		panic(err)
	}
}