rtc: connection, connectionsvc, service, streamer

go.mod

@@ -4,8 +4,10 @@ go 1.19
 
 require (
 	github.com/gen2brain/x264-go v0.2.4
+	github.com/google/uuid v1.1.1
 	github.com/kbinani/screenshot v0.0.0-20210720154843-7d3a670d8329
 	github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646
+	github.com/pion/sdp v1.3.0
 	github.com/pion/webrtc/v2 v2.2.26
 )
 
@@ -14,7 +16,6 @@ require (
 	github.com/gen2brain/shm v0.0.0-20200228170931-49f9650110c5 // indirect
 	github.com/gen2brain/x264-go/x264c v0.0.0-20221204084822-82ee2951dea2 // indirect
 	github.com/gen2brain/x264-go/yuv v0.0.0-20221204084822-82ee2951dea2 // indirect
-	github.com/google/uuid v1.1.1 // indirect
 	github.com/jezek/xgb v0.0.0-20210312150743-0e0f116e1240 // indirect
 	github.com/lucas-clemente/quic-go v0.7.1-0.20190401152353-907071221cf9 // indirect
 	github.com/lxn/win v0.0.0-20210218163916-a377121e959e // indirect

go.sum

@@ -61,6 +61,8 @@ github.com/pion/rtp v1.6.0 h1:4Ssnl/T5W2LzxHj9ssYpGVEQh3YYhQFNVmSWO88MMwk=
 github.com/pion/rtp v1.6.0/go.mod h1:QgfogHsMBVE/RFNno467U/KBqfUywEH+HK+0rtnwsdI=
 github.com/pion/sctp v1.7.10 h1:o3p3/hZB5Cx12RMGyWmItevJtZ6o2cpuxaw6GOS4x+8=
 github.com/pion/sctp v1.7.10/go.mod h1:EhpTUQu1/lcK3xI+eriS6/96fWetHGCvBi9MSsnaBN0=
+github.com/pion/sdp v1.3.0 h1:21lpgEILHyolpsIrbCBagZaAPj4o057cFjzaFebkVOs=
+github.com/pion/sdp v1.3.0/go.mod h1:ceA2lTyftydQTuCIbUNoH77aAt6CiQJaRpssA4Gee8I=
 github.com/pion/sdp/v2 v2.4.0 h1:luUtaETR5x2KNNpvEMv/r4Y+/kzImzbz4Lm1z8eQNQI=
 github.com/pion/sdp/v2 v2.4.0/go.mod h1:L2LxrOpSTJbAns244vfPChbciR/ReU1KWfG04OpkR7E=
 github.com/pion/srtp v1.5.1 h1:9Q3jAfslYZBt+C69SI/ZcONJh9049JUHZWYRRf5KEKw=

rtc/connection.go

@@ -0,0 +1,199 @@
+package rtc
+
+import (
+	"fmt"
+	"image"
+	"log"
+	"math/rand"
+	"strconv"
+	"strings"
+
+	"github.com/google/uuid"
+	"github.com/krishpranav/remote-desktop/encoders"
+	"github.com/krishpranav/remote-desktop/rdisplay"
+	"github.com/pion/sdp"
+	"github.com/pion/webrtc/v2"
+)
+
+type RemoteScreenPeerConn struct {
+	connection *webrtc.PeerConnection
+	stunServer string
+	track      *webrtc.Track
+	streamer   videoStreamer
+	grabber    rdisplay.ScreenGrabber
+	encService encoders.Service
+}
+
+func findBestCodec(sdp *sdp.SessionDescription, encService encoders.Service, h264Profile string) (*webrtc.RTPCodec, encoders.VideoCodec, error) {
+	var h264Codec *webrtc.RTPCodec
+	var vp8Codec *webrtc.RTPCodec
+	for _, md := range sdp.MediaDescriptions {
+		for _, format := range md.MediaName.Formats {
+			intPt, err := strconv.Atoi(format)
+			payloadType := uint8(intPt)
+			sdpCodec, err := sdp.GetCodecForPayloadType(payloadType)
+			if err != nil {
+				return nil, encoders.NoCodec, fmt.Errorf("Can't find codec for %d", payloadType)
+			}
+
+			if sdpCodec.Name == webrtc.H264 && h264Codec == nil {
+				packetSupport := strings.Contains(sdpCodec.Fmtp, "packetization-mode=1")
+				supportsProfile := strings.Contains(sdpCodec.Fmtp, fmt.Sprintf("profile-level-id=%s", h264Profile))
+				if packetSupport && supportsProfile {
+					h264Codec = webrtc.NewRTPH264Codec(payloadType, sdpCodec.ClockRate)
+					h264Codec.SDPFmtpLine = sdpCodec.Fmtp
+				}
+			} else if sdpCodec.Name == webrtc.VP8 && vp8Codec == nil {
+				vp8Codec = webrtc.NewRTPVP8Codec(payloadType, sdpCodec.ClockRate)
+				vp8Codec.SDPFmtpLine = sdpCodec.Fmtp
+			}
+		}
+	}
+	if vp8Codec != nil && encService.Supports(encoders.VP8Codec) {
+		return vp8Codec, encoders.VP8Codec, nil
+	}
+	if h264Codec != nil && encService.Supports(encoders.H264Codec) {
+		return h264Codec, encoders.H264Codec, nil
+	}
+	return nil, encoders.NoCodec, fmt.Errorf("Couldn't find a matching codec")
+}
+
+func newRemoteScreenPeerConn(stunServer string, grabber rdisplay.ScreenGrabber, encService encoders.Service) *RemoteScreenPeerConn {
+	return &RemoteScreenPeerConn{
+		stunServer: stunServer,
+		grabber:    grabber,
+		encService: encService,
+	}
+}
+
+func getTrackDirection(sdp *sdp.SessionDescription) webrtc.RTPTransceiverDirection {
+	for _, mediaDesc := range sdp.MediaDescriptions {
+		if mediaDesc.MediaName.Media == "video" {
+			if _, recvOnly := mediaDesc.Attribute("recvonly"); recvOnly {
+				return webrtc.RTPTransceiverDirectionRecvonly
+			} else if _, sendRecv := mediaDesc.Attribute("sendrecv"); sendRecv {
+				return webrtc.RTPTransceiverDirectionSendrecv
+			}
+		}
+	}
+	return webrtc.RTPTransceiverDirectionInactive
+}
+
+func (p *RemoteScreenPeerConn) ProcessOffer(strOffer string) (string, error) {
+	sdp := sdp.SessionDescription{}
+	err := sdp.Unmarshal(strOffer)
+	if err != nil {
+		return "", err
+	}
+
+	webrtcCodec, encCodec, err := findBestCodec(&sdp, p.encService, "42e01f")
+	if err != nil {
+		return "", err
+	}
+	mediaEngine := webrtc.MediaEngine{}
+	mediaEngine.RegisterCodec(webrtcCodec)
+
+	api := webrtc.NewAPI(webrtc.WithMediaEngine(mediaEngine))
+
+	pcconf := webrtc.Configuration{
+		ICEServers: []webrtc.ICEServer{
+			webrtc.ICEServer{
+				URLs: []string{p.stunServer},
+			},
+		},
+		SDPSemantics: webrtc.SDPSemanticsUnifiedPlan,
+	}
+
+	peerConn, err := api.NewPeerConnection(pcconf)
+	if err != nil {
+		return "", err
+	}
+	p.connection = peerConn
+
+	peerConn.OnICEConnectionStateChange(func(connState webrtc.ICEConnectionState) {
+		if connState == webrtc.ICEConnectionStateConnected {
+			p.start()
+		}
+		if connState == webrtc.ICEConnectionStateDisconnected {
+			p.Close()
+		}
+		log.Printf("Connection state: %s \n", connState.String())
+	})
+
+	track, err := peerConn.NewTrack(
+		webrtcCodec.PayloadType,
+		uint32(rand.Int31()),
+		uuid.New().String(),
+		fmt.Sprintf("remote-screen"),
+	)
+
+	log.Printf("Using codec %s (%d) %s", webrtcCodec.Name, webrtcCodec.PayloadType, webrtcCodec.SDPFmtpLine)
+
+	direction := getTrackDirection(&sdp)
+
+	if direction == webrtc.RTPTransceiverDirectionSendrecv {
+		_, err = peerConn.AddTrack(track)
+	} else if direction == webrtc.RTPTransceiverDirectionRecvonly {
+		_, err = peerConn.AddTransceiverFromTrack(track, webrtc.RtpTransceiverInit{
+			Direction: webrtc.RTPTransceiverDirectionSendonly,
+		})
+	} else {
+		return "", fmt.Errorf("Unsupported transceiver direction")
+	}
+
+	offerSdp := webrtc.SessionDescription{
+		SDP:  strOffer,
+		Type: webrtc.SDPTypeOffer,
+	}
+	err = peerConn.SetRemoteDescription(offerSdp)
+	if err != nil {
+		return "", err
+	}
+
+	p.track = track
+
+	answer, err := peerConn.CreateAnswer(nil)
+	if err != nil {
+		return "", err
+	}
+
+	screen := p.grabber.Screen()
+	sourceSize := image.Point{
+		screen.Bounds.Dx(),
+		screen.Bounds.Dy(),
+	}
+
+	encoder, err := p.encService.NewEncoder(encCodec, sourceSize, p.grabber.Fps())
+	if err != nil {
+		return "", err
+	}
+
+	size, err := encoder.VideoSize()
+	if err != nil {
+		return "", err
+	}
+
+	p.streamer = newRTCStreamer(p.track, &p.grabber, &encoder, size)
+
+	err = peerConn.SetLocalDescription(answer)
+	if err != nil {
+		return "", err
+	}
+	return answer.SDP, nil
+}
+
+func (p *RemoteScreenPeerConn) start() {
+	p.streamer.start()
+}
+
+func (p *RemoteScreenPeerConn) Close() error {
+
+	if p.streamer != nil {
+		p.streamer.close()
+	}
+
+	if p.connection != nil {
+		return p.connection.Close()
+	}
+	return nil
+}

rtc/connectionsvc.go

@@ -0,0 +1,58 @@
+package rtc
+
+import (
+	"fmt"
+
+	"github.com/krishpranav/remote-desktop/encoders"
+	"github.com/krishpranav/remote-desktop/rdisplay"
+)
+
+type RemoteScreenService struct {
+	stunServer      string
+	videoService    rdisplay.Service
+	encodingService encoders.Service
+}
+
+func NewRemoteScreenService(stun string, video rdisplay.Service, enc encoders.Service) Service {
+	return &RemoteScreenService{
+		stunServer:      stun,
+		videoService:    video,
+		encodingService: enc,
+	}
+}
+
+func hasElement(haystack []string, needle string) bool {
+	for _, item := range haystack {
+		if item == needle {
+			return true
+		}
+	}
+	return false
+}
+
+func (svc *RemoteScreenService) CreateRemoteScreenConnection(screenIx int, fps int) (RemoteScreenConnection, error) {
+	screens, err := svc.videoService.Screens()
+	if err != nil {
+		return nil, err
+	}
+
+	if screenIx < 0 || screenIx > len(screens) {
+		screenIx = 0
+	}
+	screen := screens[screenIx]
+	screenGrabber, err := svc.videoService.CreateScreenGrabber(screen, fps)
+	if err != nil {
+		return nil, err
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	if len(screens) == 0 {
+		return nil, fmt.Errorf("No available screens")
+	}
+
+	rtcPeer := newRemoteScreenPeerConn(svc.stunServer, screenGrabber, svc.encodingService)
+	return rtcPeer, nil
+}

rtc/service.go

@@ -0,0 +1,19 @@
+package rtc
+
+import (
+	"io"
+)
+
+type videoStreamer interface {
+	start()
+	close()
+}
+
+type RemoteScreenConnection interface {
+	io.Closer
+	ProcessOffer(offer string) (string, error)
+}
+
+type Service interface {
+	CreateRemoteScreenConnection(screenIx int, fps int) (RemoteScreenConnection, error)
+}

rtc/streamer.go

@@ -0,0 +1,76 @@
+package rtc
+
+import (
+	"fmt"
+	"image"
+
+	"github.com/krishpranav/remote-desktop/encoders"
+	"github.com/krishpranav/remote-desktop/rdisplay"
+	"github.com/nfnt/resize"
+	"github.com/pion/webrtc/v2"
+	"github.com/pion/webrtc/v2/pkg/media"
+)
+
+func resizeImage(src *image.RGBA, target image.Point) *image.RGBA {
+	return resize.Resize(uint(target.X), uint(target.Y), src, resize.Lanczos3).(*image.RGBA)
+}
+
+type rtcStreamer struct {
+	track   *webrtc.Track
+	stop    chan struct{}
+	screen  *rdisplay.ScreenGrabber
+	encoder *encoders.Encoder
+	size    image.Point
+}
+
+func newRTCStreamer(track *webrtc.Track, screen *rdisplay.ScreenGrabber, encoder *encoders.Encoder, size image.Point) videoStreamer {
+	return &rtcStreamer{
+		track:   track,
+		stop:    make(chan struct{}),
+		screen:  screen,
+		encoder: encoder,
+		size:    size,
+	}
+}
+
+func (s *rtcStreamer) start() {
+	go s.startStream()
+}
+
+func (s *rtcStreamer) startStream() {
+	screen := *s.screen
+	screen.Start()
+	frames := screen.Frames()
+	for {
+		select {
+		case <-s.stop:
+			screen.Stop()
+			return
+		case frame := <-frames:
+			err := s.stream(frame)
+			if err != nil {
+				fmt.Printf("Streamer: %v\n", err)
+				return
+			}
+		}
+	}
+}
+
+func (s *rtcStreamer) stream(frame *image.RGBA) error {
+	resized := resizeImage(frame, s.size)
+	payload, err := (*s.encoder).Encode(resized)
+	if err != nil {
+		return err
+	}
+	if payload == nil {
+		return nil
+	}
+	return s.track.WriteSample(media.Sample{
+		Data:    payload,
+		Samples: 1,
+	})
+}
+
+func (s *rtcStreamer) close() {
+	close(s.stop)
+}