From 140014f2a68c65dc7b3b96bfa589a18bb2a03f2a Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Wed, 18 Jan 2023 15:04:06 +0300
Subject: [PATCH] Fix info for WS/MP4

---
 cmd/mp4/ws.go      |  6 +++++-
 pkg/mp4/segment.go | 22 +++++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/cmd/mp4/ws.go b/cmd/mp4/ws.go
index e5c4dc04..4944a848 100644
--- a/cmd/mp4/ws.go
+++ b/cmd/mp4/ws.go
@@ -71,7 +71,11 @@ func handlerWSMP4(tr *api.Transport, msg *api.Message) error {
 		return errors.New(api.StreamNotFound)
 	}
 
-	cons := &mp4.Segment{OnlyKeyframe: true}
+	cons := &mp4.Segment{
+		RemoteAddr:   tr.Request.RemoteAddr,
+		UserAgent:    tr.Request.UserAgent(),
+		OnlyKeyframe: true,
+	}
 
 	if codecs, ok := msg.Value.(string); ok {
 		log.Trace().Str("codecs", codecs).Msgf("[mp4] new WS/MP4 consumer")
diff --git a/pkg/mp4/segment.go b/pkg/mp4/segment.go
index 7411701d..9cc3a88a 100644
--- a/pkg/mp4/segment.go
+++ b/pkg/mp4/segment.go
@@ -1,18 +1,25 @@
 package mp4
 
 import (
+	"encoding/json"
 	"github.com/AlexxIT/go2rtc/pkg/h264"
 	"github.com/AlexxIT/go2rtc/pkg/h265"
 	"github.com/AlexxIT/go2rtc/pkg/streamer"
 	"github.com/pion/rtp"
+	"sync/atomic"
 )
 
 type Segment struct {
 	streamer.Element
 
-	Medias       []*streamer.Media
+	Medias     []*streamer.Media
+	UserAgent  string
+	RemoteAddr string
+
 	MimeType     string
 	OnlyKeyframe bool
+
+	send uint32
 }
 
 func (c *Segment) GetMedias() []*streamer.Media {
@@ -56,6 +63,7 @@ func (c *Segment) AddTrack(media *streamer.Media, track *streamer.Track) *stream
 				}
 
 				buf := muxer.Marshal(0, packet)
+				atomic.AddUint32(&c.send, uint32(len(buf)))
 				c.Fire(append(init, buf...))
 
 				return nil
@@ -73,6 +81,7 @@ func (c *Segment) AddTrack(media *streamer.Media, track *streamer.Track) *stream
 						buf = append(buf, b...)
 					}
 
+					atomic.AddUint32(&c.send, uint32(len(buf)))
 					c.Fire(buf)
 
 					buf = buf[:0]
@@ -106,6 +115,7 @@ func (c *Segment) AddTrack(media *streamer.Media, track *streamer.Track) *stream
 			}
 
 			buf := muxer.Marshal(0, packet)
+			atomic.AddUint32(&c.send, uint32(len(buf)))
 			c.Fire(append(init, buf...))
 
 			return nil
@@ -121,3 +131,13 @@ func (c *Segment) AddTrack(media *streamer.Media, track *streamer.Track) *stream
 
 	panic("unsupported codec")
 }
+
+func (c *Segment) MarshalJSON() ([]byte, error) {
+	info := &streamer.Info{
+		Type:       "WS/MP4 client",
+		RemoteAddr: c.RemoteAddr,
+		UserAgent:  c.UserAgent,
+		Send:       atomic.LoadUint32(&c.send),
+	}
+	return json.Marshal(info)
+}