From c873d4d3445ef324c4c7227426a62d488782ad84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Sun, 29 Jan 2023 21:29:16 +0100
Subject: [PATCH] screenshare POC.

---
 client/src/components/header.vue         |  34 ++++++
 client/src/neko/base.ts                  |  16 +++
 server/internal/capture/manager.go       |  19 ++++
 server/internal/capture/streamsrcsink.go | 137 +++++++++++++++++++++++
 server/internal/config/capture.go        |  18 +++
 server/internal/types/capture.go         |  12 ++
 server/internal/webrtc/webrtc.go         |  96 +++++++++++++++-
 7 files changed, 331 insertions(+), 1 deletion(-)
 create mode 100644 server/internal/capture/streamsrcsink.go

diff --git a/client/src/components/header.vue b/client/src/components/header.vue
index a072f80..23e952a 100644
--- a/client/src/components/header.vue
+++ b/client/src/components/header.vue
@@ -5,6 +5,14 @@
       <span><b>n</b>.eko</span>
     </a>
     <ul class="menu">
+      <li>
+        <button class="btn" @click="startShareScreen" v-if="!mediaStream">
+          START SCREEN SHARE
+        </button>
+        <button class="btn" @click="stopShareScreen" v-else>
+          STOP SCREEN SHARE
+        </button>
+      </li>
       <li>
         <i
           :class="[{ disabled: !admin }, { locked: isLocked('control') }, 'fas', 'fa-mouse']"
@@ -207,5 +215,31 @@
 
       return this.$t(`locks.${resource}.` + (this.isLocked(resource) ? `locked` : `unlocked`))
     }
+
+    //
+    // Screen Share
+    //
+    mediaStream: MediaStream | null = null
+    mediaRtcpSender: RTCRtpSender | null = null
+    async startShareScreen() {
+      // get media stream from user's browser
+      this.mediaStream = await navigator.mediaDevices
+        .getDisplayMedia({
+          video: true,
+          audio: false,
+        })
+      const mediaTrack = this.mediaStream.getVideoTracks()[0];
+      this.mediaRtcpSender = this.$client.addTrack(mediaTrack, this.mediaStream)
+    }
+    async stopShareScreen() {
+      if (this.mediaStream) {
+        this.mediaStream.getTracks().forEach(track => track.stop())
+        this.mediaStream = null
+      }
+      if (this.mediaRtcpSender) {
+        this.$client.removeTrack(this.mediaRtcpSender)
+        this.mediaRtcpSender = null
+      }
+    }
   }
 </script>
diff --git a/client/src/neko/base.ts b/client/src/neko/base.ts
index 57baef8..f519438 100644
--- a/client/src/neko/base.ts
+++ b/client/src/neko/base.ts
@@ -313,6 +313,22 @@ export abstract class BaseClient extends EventEmitter<BaseEvents> {
     this._peer.setRemoteDescription({ type: 'answer', sdp })
   }
 
+  public addTrack(track: MediaStreamTrack, ...streams: MediaStream[]): RTCRtpSender {
+    if (!this._peer) {
+      throw new Error('peer not connected')
+    }
+
+    return this._peer.addTrack(track, ...streams)
+  }
+
+  public removeTrack(sender: RTCRtpSender) {
+    if (!this._peer) {
+      throw new Error('peer not connected')
+    }
+
+    this._peer.removeTrack(sender)
+  }
+
   private async onMessage(e: MessageEvent) {
     const { event, ...payload } = JSON.parse(e.data) as WebSocketMessages
 
diff --git a/server/internal/capture/manager.go b/server/internal/capture/manager.go
index 014e1a6..d6681b7 100644
--- a/server/internal/capture/manager.go
+++ b/server/internal/capture/manager.go
@@ -2,12 +2,14 @@ package capture
 
 import (
 	"errors"
+	"fmt"
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 
 	"m1k1o/neko/internal/config"
 	"m1k1o/neko/internal/types"
+	"m1k1o/neko/internal/types/codec"
 )
 
 type CaptureManagerCtx struct {
@@ -18,6 +20,9 @@ type CaptureManagerCtx struct {
 	broadcast *BroacastManagerCtx
 	audio     *StreamSinkManagerCtx
 	video     *StreamSinkManagerCtx
+
+	// source-sinks
+	screenshare *StreamSrcSinkManagerCtx
 }
 
 func New(desktop types.DesktopManager, config *config.Capture) *CaptureManagerCtx {
@@ -43,6 +48,15 @@ func New(desktop types.DesktopManager, config *config.Capture) *CaptureManagerCt
 			}
 			return NewVideoPipeline(config.VideoCodec, config.Display, config.VideoPipeline, fps, config.VideoBitrate, config.VideoHWEnc)
 		}, "video"),
+
+		// source-sinks
+		screenshare: streamSrcSinkNew(config.ScreenshareEnabled, map[string]string{
+			codec.VP8().Name: "appsrc format=time is-live=true do-timestamp=true name=appsrc " +
+				fmt.Sprintf("! application/x-rtp, payload=%d, encoding-name=VP8-DRAFT-IETF-01 ", codec.VP8().PayloadType) +
+				"! rtpvp8depay " +
+				"! appsink name=appsink",
+			// TODO: Add support for more codecs.
+		}, "webcam"),
 	}
 }
 
@@ -95,6 +109,7 @@ func (manager *CaptureManagerCtx) Start() {
 func (manager *CaptureManagerCtx) Shutdown() error {
 	manager.logger.Info().Msgf("shutdown")
 
+	manager.screenshare.shutdown()
 	manager.broadcast.shutdown()
 
 	manager.audio.shutdown()
@@ -114,3 +129,7 @@ func (manager *CaptureManagerCtx) Audio() types.StreamSinkManager {
 func (manager *CaptureManagerCtx) Video() types.StreamSinkManager {
 	return manager.video
 }
+
+func (manager *CaptureManagerCtx) Screenshare() types.StreamSrcSinkManager {
+	return manager.screenshare
+}
diff --git a/server/internal/capture/streamsrcsink.go b/server/internal/capture/streamsrcsink.go
new file mode 100644
index 0000000..d723cfb
--- /dev/null
+++ b/server/internal/capture/streamsrcsink.go
@@ -0,0 +1,137 @@
+package capture
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+
+	"m1k1o/neko/internal/capture/gst"
+	"m1k1o/neko/internal/types"
+	"m1k1o/neko/internal/types/codec"
+)
+
+type StreamSrcSinkManagerCtx struct {
+	logger        zerolog.Logger
+	sampleChannel chan types.Sample
+
+	enabled       bool
+	codecPipeline map[string]string // codec -> pipeline
+
+	codec       codec.RTPCodec
+	pipeline    *gst.Pipeline
+	pipelineMu  sync.Mutex
+	pipelineStr string
+}
+
+func streamSrcSinkNew(enabled bool, codecPipeline map[string]string, video_id string) *StreamSrcSinkManagerCtx {
+	logger := log.With().
+		Str("module", "capture").
+		Str("submodule", "stream-src-sink").
+		Str("video_id", video_id).Logger()
+
+	return &StreamSrcSinkManagerCtx{
+		logger:        logger,
+		enabled:       enabled,
+		codecPipeline: codecPipeline,
+		sampleChannel: make(chan types.Sample),
+	}
+}
+
+func (manager *StreamSrcSinkManagerCtx) shutdown() {
+	manager.logger.Info().Msgf("shutdown")
+
+	manager.Stop()
+}
+
+func (manager *StreamSrcSinkManagerCtx) Codec() codec.RTPCodec {
+	manager.pipelineMu.Lock()
+	defer manager.pipelineMu.Unlock()
+
+	return manager.codec
+}
+
+func (manager *StreamSrcSinkManagerCtx) Start(codec codec.RTPCodec) error {
+	manager.pipelineMu.Lock()
+	defer manager.pipelineMu.Unlock()
+
+	if manager.pipeline != nil {
+		return types.ErrCapturePipelineAlreadyExists
+	}
+
+	if !manager.enabled {
+		return errors.New("stream-src-sink not enabled")
+	}
+
+	found := false
+	for codecName, pipeline := range manager.codecPipeline {
+		if codecName == codec.Name {
+			manager.pipelineStr = pipeline
+			manager.codec = codec
+			found = true
+			break
+		}
+	}
+
+	if !found {
+		return errors.New("no pipeline found for a codec")
+	}
+
+	var err error
+
+	manager.logger.Info().
+		Str("codec", manager.codec.Name).
+		Str("src", manager.pipelineStr).
+		Msgf("creating pipeline")
+
+	manager.pipeline, err = gst.CreatePipeline(manager.pipelineStr)
+	if err != nil {
+		return err
+	}
+
+	manager.pipeline.AttachAppsrc("appsrc")
+	manager.pipeline.AttachAppsink("appsink", manager.sampleChannel)
+	manager.pipeline.Play()
+
+	return nil
+}
+
+func (manager *StreamSrcSinkManagerCtx) Stop() {
+	manager.pipelineMu.Lock()
+	defer manager.pipelineMu.Unlock()
+
+	if manager.pipeline == nil {
+		return
+	}
+
+	manager.pipeline.Destroy()
+	manager.pipeline = nil
+
+	manager.logger.Info().
+		Str("codec", manager.codec.Name).
+		Str("src", manager.pipelineStr).
+		Msgf("destroying pipeline")
+}
+
+func (manager *StreamSrcSinkManagerCtx) Push(bytes []byte) {
+	manager.pipelineMu.Lock()
+	defer manager.pipelineMu.Unlock()
+
+	if manager.pipeline == nil {
+		return
+	}
+
+	manager.pipeline.Push(bytes)
+}
+
+func (manager *StreamSrcSinkManagerCtx) Started() bool {
+	manager.pipelineMu.Lock()
+	defer manager.pipelineMu.Unlock()
+
+	return manager.pipeline != nil
+}
+
+func (manager *StreamSrcSinkManagerCtx) GetSampleChannel() chan types.Sample {
+	return manager.sampleChannel
+}
diff --git a/server/internal/config/capture.go b/server/internal/config/capture.go
index 820557a..7360f4d 100644
--- a/server/internal/config/capture.go
+++ b/server/internal/config/capture.go
@@ -27,6 +27,9 @@ type Capture struct {
 	// broadcast
 	BroadcastPipeline string
 	BroadcastUrl      string
+
+	// screenshare
+	ScreenshareEnabled bool
 }
 
 func (Capture) Init(cmd *cobra.Command) error {
@@ -151,6 +154,15 @@ func (Capture) Init(cmd *cobra.Command) error {
 		return err
 	}
 
+	//
+	// screenshare
+	//
+
+	cmd.PersistentFlags().Bool("screenshare.enabled", true, "enable screenshare")
+	if err := viper.BindPFlag("screenshare.enabled", cmd.PersistentFlags().Lookup("screenshare.enabled")); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -230,4 +242,10 @@ func (s *Capture) Set() {
 
 	s.BroadcastPipeline = viper.GetString("broadcast_pipeline")
 	s.BroadcastUrl = viper.GetString("broadcast_url")
+
+	//
+	// screenshare
+	//
+
+	s.ScreenshareEnabled = viper.GetBool("screenshare.enabled")
 }
diff --git a/server/internal/types/capture.go b/server/internal/types/capture.go
index 36e54b8..5fc5beb 100644
--- a/server/internal/types/capture.go
+++ b/server/internal/types/capture.go
@@ -28,6 +28,17 @@ type StreamSinkManager interface {
 	GetSampleChannel() chan Sample
 }
 
+type StreamSrcSinkManager interface {
+	Codec() codec.RTPCodec
+
+	Start(codec codec.RTPCodec) error
+	Stop()
+
+	Push(bytes []byte)
+	Started() bool
+	GetSampleChannel() chan Sample
+}
+
 type CaptureManager interface {
 	Start()
 	Shutdown() error
@@ -35,4 +46,5 @@ type CaptureManager interface {
 	Broadcast() BroadcastManager
 	Audio() StreamSinkManager
 	Video() StreamSinkManager
+	Screenshare() StreamSrcSinkManager
 }
diff --git a/server/internal/webrtc/webrtc.go b/server/internal/webrtc/webrtc.go
index 693f572..382ec58 100644
--- a/server/internal/webrtc/webrtc.go
+++ b/server/internal/webrtc/webrtc.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/pion/ice/v2"
 	"github.com/pion/interceptor"
+	"github.com/pion/rtcp"
 	"github.com/pion/webrtc/v3"
 	"github.com/pion/webrtc/v3/pkg/media"
 	"github.com/rs/zerolog"
@@ -18,6 +19,7 @@ import (
 
 	"m1k1o/neko/internal/config"
 	"m1k1o/neko/internal/types"
+	"m1k1o/neko/internal/types/codec"
 	"m1k1o/neko/internal/webrtc/pionlog"
 )
 
@@ -40,6 +42,8 @@ type WebRTCManager struct {
 	desktop    types.DesktopManager
 	config     *config.WebRTC
 	api        *webrtc.API
+
+	screenshareStop *func()
 }
 
 func (manager *WebRTCManager) Start() {
@@ -82,7 +86,19 @@ func (manager *WebRTCManager) Start() {
 
 	go func() {
 		for {
-			sample, ok := <-manager.capture.Video().GetSampleChannel()
+			var sample types.Sample
+			var ok bool
+
+			select {
+			case sample, ok = <-manager.capture.Video().GetSampleChannel():
+				// if screenshare is active, we need to drop all video samples
+				// ideally we would stop the video capture meanwhile.
+				if manager.capture.Screenshare().Started() {
+					continue
+				}
+			case sample, ok = <-manager.capture.Screenshare().GetSampleChannel():
+			}
+
 			if !ok {
 				manager.logger.Debug().Msg("video capture channel is closed")
 				continue
@@ -305,6 +321,84 @@ func (manager *WebRTCManager) CreatePeer(id string, session types.Session) (type
 		}
 	})
 
+	connection.OnTrack(func(track *webrtc.TrackRemote, receiver *webrtc.RTPReceiver) {
+		logger := manager.logger.With().
+			Str("kind", track.Kind().String()).
+			Str("mime", track.Codec().RTPCodecCapability.MimeType).
+			Logger()
+
+		logger.Info().Msgf("received new remote track")
+
+		// parse codec from remote track
+		codec, ok := codec.ParseRTC(track.Codec())
+		if !ok {
+			logger.Warn().Msg("remote track with unknown codec")
+			receiver.Stop()
+			return
+		}
+
+		var srcSinkManager types.StreamSrcSinkManager
+
+		stopped := false
+		stopFn := func() {
+			if stopped {
+				return
+			}
+
+			stopped = true
+			receiver.Stop()
+			srcSinkManager.Stop()
+			logger.Info().Msg("remote track stopped")
+		}
+
+		logger.Info().Msgf("found codec %s", codec.Name)
+
+		if track.Kind() == webrtc.RTPCodecTypeVideo {
+			// video -> webcam
+			srcSinkManager = manager.capture.Screenshare()
+			defer stopFn()
+
+			if manager.screenshareStop != nil {
+				(*manager.screenshareStop)()
+			}
+			manager.screenshareStop = &stopFn
+		} else {
+			logger.Warn().Msg("expected only video tracks")
+			receiver.Stop()
+			return
+		}
+
+		logger.Info().Msg("starting srcSinkManager")
+		err := srcSinkManager.Start(codec)
+		if err != nil {
+			logger.Err(err).Msg("failed to start pipeline")
+			return
+		}
+
+		ticker := time.NewTicker(3 * time.Second)
+		defer ticker.Stop()
+
+		go func() {
+			for range ticker.C {
+				err := connection.WriteRTCP([]rtcp.Packet{&rtcp.PictureLossIndication{MediaSSRC: uint32(track.SSRC())}})
+				if err != nil {
+					logger.Err(err).Msg("remote track rtcp send err")
+				}
+			}
+		}()
+
+		buf := make([]byte, 1400)
+		for {
+			i, _, err := track.Read(buf)
+			if err != nil {
+				logger.Warn().Err(err).Msg("failed read from remote track")
+				break
+			}
+
+			srcSinkManager.Push(buf[:i])
+		}
+	})
+
 	if err := session.SetPeer(peer); err != nil {
 		return nil, err
 	}