client: add OnRequest callback (#163)

client.go

@@ -41,6 +41,9 @@ type ClientOnDownloadPartFunc func(url string)
 // ClientOnDecodeErrorFunc is the prototype of Client.OnDecodeError.
 type ClientOnDecodeErrorFunc func(err error)
 
+// ClientOnRequestFunc is the prototype of the function passed to OnRequest().
+type ClientOnRequestFunc func(*http.Request)
+
 // ClientOnTracksFunc is the prototype of the function passed to OnTracks().
 type ClientOnTracksFunc func([]*Track) error
 
@@ -85,6 +88,8 @@ type Client struct {
 	//
 	// callbacks (all optional)
 	//
+	// called when sending a request to the server.
+	OnRequest ClientOnRequestFunc
 	// called when tracks are available.
 	OnTracks ClientOnTracksFunc
 	// called before downloading a primary playlist.
@@ -119,6 +124,9 @@ func (c *Client) Start() error {
 	if c.HTTPClient == nil {
 		c.HTTPClient = http.DefaultClient
 	}
+	if c.OnRequest == nil {
+		c.OnRequest = func(_ *http.Request) {}
+	}
 	if c.OnTracks == nil {
 		c.OnTracks = func(_ []*Track) error {
 			return nil
@@ -229,6 +237,7 @@ func (c *Client) runInner() error {
 	c.primaryDownloader = &clientPrimaryDownloader{
 		primaryPlaylistURL:        c.playlistURL,
 		httpClient:                c.HTTPClient,
+		onRequest:                 c.OnRequest,
 		onDownloadPrimaryPlaylist: c.OnDownloadPrimaryPlaylist,
 		onDownloadStreamPlaylist:  c.OnDownloadStreamPlaylist,
 		onDownloadSegment:         c.OnDownloadSegment,

client_primary_downloader.go

@@ -43,13 +43,16 @@ func cloneURL(ur *url.URL) *url.URL {
 func clientDownloadPlaylist(
 	ctx context.Context,
 	httpClient *http.Client,
+	onRequest ClientOnRequestFunc,
 	ur *url.URL,
 ) (playlist.Playlist, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, ur.String(), nil)
 	if err != nil {
 		return nil, err
 	}
 
+	onRequest(req)
+
 	res, err := httpClient.Do(req)
 	if err != nil {
 		return nil, err
@@ -124,6 +127,7 @@ type streamTracksEntry struct {
 type clientPrimaryDownloader struct {
 	primaryPlaylistURL        *url.URL
 	httpClient                *http.Client
+	onRequest                 ClientOnRequestFunc
 	onDownloadPrimaryPlaylist ClientOnDownloadPrimaryPlaylistFunc
 	onDownloadStreamPlaylist  ClientOnDownloadStreamPlaylistFunc
 	onDownloadSegment         ClientOnDownloadSegmentFunc
@@ -153,7 +157,7 @@ func (d *clientPrimaryDownloader) initialize() {
 func (d *clientPrimaryDownloader) run(ctx context.Context) error {
 	d.onDownloadPrimaryPlaylist(d.primaryPlaylistURL.String())
 
-	pl, err := clientDownloadPlaylist(ctx, d.httpClient, d.primaryPlaylistURL)
+	pl, err := clientDownloadPlaylist(ctx, d.httpClient, d.onRequest, d.primaryPlaylistURL)
 	if err != nil {
 		return err
 	}
@@ -165,6 +169,7 @@ func (d *clientPrimaryDownloader) run(ctx context.Context) error {
 		ds := &clientStreamDownloader{
 			isLeading:                true,
 			httpClient:               d.httpClient,
+			onRequest:                d.onRequest,
 			onDownloadStreamPlaylist: d.onDownloadStreamPlaylist,
 			onDownloadSegment:        d.onDownloadSegment,
 			onDownloadPart:           d.onDownloadPart,
@@ -195,6 +200,7 @@ func (d *clientPrimaryDownloader) run(ctx context.Context) error {
 		ds := &clientStreamDownloader{
 			isLeading:                true,
 			httpClient:               d.httpClient,
+			onRequest:                d.onRequest,
 			onDownloadStreamPlaylist: d.onDownloadStreamPlaylist,
 			onDownloadSegment:        d.onDownloadSegment,
 			onDownloadPart:           d.onDownloadPart,
@@ -224,6 +230,7 @@ func (d *clientPrimaryDownloader) run(ctx context.Context) error {
 
 				ds := &clientStreamDownloader{
 					isLeading:                false,
+					onRequest:                d.onRequest,
 					httpClient:               d.httpClient,
 					onDownloadStreamPlaylist: d.onDownloadStreamPlaylist,
 					onDownloadSegment:        d.onDownloadSegment,

client_stream_downloader.go

@@ -52,6 +52,7 @@ func dateTimeOfPreloadHint(pl *playlist.Media) *time.Time {
 type clientStreamDownloader struct {
 	isLeading                bool
 	httpClient               *http.Client
+	onRequest                ClientOnRequestFunc
 	onDownloadStreamPlaylist ClientOnDownloadStreamPlaylistFunc
 	onDownloadSegment        ClientOnDownloadSegmentFunc
 	onDownloadPart           ClientOnDownloadPartFunc
@@ -189,7 +190,7 @@ func (d *clientStreamDownloader) downloadPlaylist(
 
 	d.onDownloadStreamPlaylist(ur.String())
 
-	pl, err := clientDownloadPlaylist(ctx, d.httpClient, ur)
+	pl, err := clientDownloadPlaylist(ctx, d.httpClient, d.onRequest, ur)
 	if err != nil {
 		return nil, err
 	}
@@ -223,6 +224,8 @@ func (d *clientStreamDownloader) downloadPreloadHint(
 			"-"+strconv.FormatUint(preloadHint.ByteRangeStart+*preloadHint.ByteRangeLength-1, 10))
 	}
 
+	d.onRequest(req)
+
 	res, err := d.httpClient.Do(req)
 	if err != nil {
 		return nil, err
@@ -268,6 +271,8 @@ func (d *clientStreamDownloader) downloadSegment(
 			"-"+strconv.FormatUint(*start+*length-1, 10))
 	}
 
+	d.onRequest(req)
+
 	res, err := d.httpClient.Do(req)
 	if err != nil {
 		return nil, err

client_test.go

@@ -73,6 +73,11 @@ y++U32uuSFiXDcSLarfIsE992MEJLSAynbF1Rsgsr3gXbGiuToJRyxbIeVy7gwzD
 -----END RSA PRIVATE KEY-----
 `)
 
+const (
+	testHeaderKey   = "X-Test-Header"
+	testHeaderValue = "test-value"
+)
+
 func writeTempFile(byts []byte) (string, error) {
 	tmpf, err := os.CreateTemp(os.TempDir(), "rtsp-")
 	if err != nil {
@@ -117,6 +122,11 @@ func TestClient(t *testing.T) {
 			t.Run(mode+"_"+format, func(t *testing.T) {
 				httpServ := &http.Server{
 					Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+						// reject requests that do not have out custom header - key pair to ensure OnRequest was called
+						if r.Header.Get(testHeaderKey) != testHeaderValue {
+							w.WriteHeader(http.StatusUnauthorized)
+							return
+						}
 						if format == "mpegts" {
 							switch {
 							case r.Method == http.MethodGet && r.URL.Path == "/stream.m3u8":
@@ -390,6 +400,9 @@ func TestClient(t *testing.T) {
 				c = &Client{
 					URI:        prefix + "://localhost:5780/stream.m3u8",
 					HTTPClient: &http.Client{Transport: tr},
+					OnRequest: func(r *http.Request) {
+						r.Header.Set(testHeaderKey, testHeaderValue)
+					},
 					OnTracks: func(tracks []*Track) error {
 						var sps []byte
 						var pps []byte
@@ -476,11 +489,23 @@ func TestClient(t *testing.T) {
 				err = c.Start()
 				require.NoError(t, err)
 
-				<-videoRecv
-				<-audioRecv
-
-				err = <-c.Wait()
-				require.Equal(t, ErrClientEOS, err)
+				// loop the channels with a select so we can catch the error if something fails
+				// for example if you get a 401 then we need to handle it
+				for {
+					exit := false
+					select {
+					case <-videoRecv:
+						continue
+					case <-audioRecv:
+						continue
+					case err := <-c.Wait():
+						require.Equal(t, ErrClientEOS, err)
+						exit = true
+					}
+					if exit {
+						break
+					}
+				}
 
 				c.Close()
 			})