From c1df3ba9e97c7e9c8c0f1afbf7015cf348373be6 Mon Sep 17 00:00:00 2001 From: Krisna Pranav Date: Sat, 15 Apr 2023 17:49:13 +0530 Subject: [PATCH] rtc: connection, connectionsvc, service, streamer --- go.mod | 3 +- go.sum | 2 + rtc/connection.go | 199 +++++++++++++++++++++++++++++++++++++++++++ rtc/connectionsvc.go | 58 +++++++++++++ rtc/service.go | 19 +++++ rtc/streamer.go | 76 +++++++++++++++++ 6 files changed, 356 insertions(+), 1 deletion(-) create mode 100644 rtc/connection.go create mode 100644 rtc/connectionsvc.go create mode 100644 rtc/service.go create mode 100644 rtc/streamer.go diff --git a/go.mod b/go.mod index a092a76..8314afe 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 0e9845f..36e433b 100644 --- a/go.sum +++ b/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= diff --git a/rtc/connection.go b/rtc/connection.go new file mode 100644 index 0000000..f049375 --- /dev/null +++ b/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 +} diff --git a/rtc/connectionsvc.go b/rtc/connectionsvc.go new file mode 100644 index 0000000..33efcd2 --- /dev/null +++ b/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 +} diff --git a/rtc/service.go b/rtc/service.go new file mode 100644 index 0000000..d672efa --- /dev/null +++ b/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) +} diff --git a/rtc/streamer.go b/rtc/streamer.go new file mode 100644 index 0000000..bdaa29b --- /dev/null +++ b/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) +}