From 398ce5b2302ae359c8e12ee252c9c729c87edbfb Mon Sep 17 00:00:00 2001
From: Pierre-Emmanuel Jacquier
 <15922119+pierre-emmanuelJ@users.noreply.github.com>
Date: Fri, 11 Sep 2020 16:20:12 +0200
Subject: [PATCH] Forward user-agent

Signed-off-by: Pierre-Emmanuel Jacquier <15922119+pierre-emmanuelJ@users.noreply.github.com>
---
 pkg/server/handlers.go           | 32 +++++++++++++++++++++++++++++---
 pkg/server/xtreamHandles.go      |  4 ++--
 pkg/xtream-proxy/xtream-proxy.go |  5 +++--
 3 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/pkg/server/handlers.go b/pkg/server/handlers.go
index 61f6bf95..670368e3 100644
--- a/pkg/server/handlers.go
+++ b/pkg/server/handlers.go
@@ -57,7 +57,17 @@ func (c *Config) stream(ctx *gin.Context, oriURL *url.URL) {
 		return
 	}
 
-	resp, err := http.Get(oriURL.String())
+	client := &http.Client{}
+
+	req, err := http.NewRequest("GET", oriURL.String(), nil)
+	if err != nil {
+		ctx.AbortWithError(http.StatusInternalServerError, err) // nolint: errcheck
+		return
+	}
+
+	req.Header.Set("User-Agent", ctx.Request.UserAgent())
+
+	resp, err := client.Do(req)
 	if err != nil {
 		ctx.AbortWithError(http.StatusInternalServerError, err) // nolint: errcheck
 		return
@@ -79,7 +89,15 @@ func (c *Config) hlsStream(ctx *gin.Context, oriURL *url.URL) {
 		},
 	}
 
-	resp, err := client.Get(oriURL.String())
+	req, err := http.NewRequest("GET", oriURL.String(), nil)
+	if err != nil {
+		ctx.AbortWithError(http.StatusInternalServerError, err) // nolint: errcheck
+		return
+	}
+
+	req.Header.Set("User-Agent", ctx.Request.UserAgent())
+
+	resp, err := client.Do(req)
 	if err != nil {
 		ctx.AbortWithError(http.StatusInternalServerError, err) // nolint: errcheck
 		return
@@ -98,7 +116,15 @@ func (c *Config) hlsStream(ctx *gin.Context, oriURL *url.URL) {
 			hlsChannelsRedirectURL[id] = *location
 			hlsChannelsRedirectURLLock.Unlock()
 
-			hlsResp, err := http.Get(location.String())
+			hlsReq, err := http.NewRequest("GET", location.String(), nil)
+			if err != nil {
+				ctx.AbortWithError(http.StatusInternalServerError, err) // nolint: errcheck
+				return
+			}
+
+			hlsReq.Header.Set("User-Agent", ctx.Request.UserAgent())
+
+			hlsResp, err := client.Do(hlsReq)
 			if err != nil {
 				ctx.AbortWithError(http.StatusInternalServerError, err) // nolint: errcheck
 				return
diff --git a/pkg/server/xtreamHandles.go b/pkg/server/xtreamHandles.go
index bb4b50d3..6665aa50 100644
--- a/pkg/server/xtreamHandles.go
+++ b/pkg/server/xtreamHandles.go
@@ -160,7 +160,7 @@ func (c *Config) xtreamPlayerAPI(ctx *gin.Context, q url.Values) {
 		action = q["action"][0]
 	}
 
-	client, err := xtreamapi.New(c.XtreamUser.String(), c.XtreamPassword.String(), c.XtreamBaseURL)
+	client, err := xtreamapi.New(c.XtreamUser.String(), c.XtreamPassword.String(), c.XtreamBaseURL, ctx.Request.UserAgent())
 	if err != nil {
 		ctx.AbortWithError(http.StatusInternalServerError, err) // nolint: errcheck
 		return
@@ -183,7 +183,7 @@ func (c *Config) xtreamPlayerAPI(ctx *gin.Context, q url.Values) {
 }
 
 func (c *Config) xtreamXMLTV(ctx *gin.Context) {
-	client, err := xtreamapi.New(c.XtreamUser.String(), c.XtreamPassword.String(), c.XtreamBaseURL)
+	client, err := xtreamapi.New(c.XtreamUser.String(), c.XtreamPassword.String(), c.XtreamBaseURL, ctx.Request.UserAgent())
 	if err != nil {
 		ctx.AbortWithError(http.StatusInternalServerError, err) // nolint: errcheck
 		return
diff --git a/pkg/xtream-proxy/xtream-proxy.go b/pkg/xtream-proxy/xtream-proxy.go
index 7a74d4b5..9cf64342 100644
--- a/pkg/xtream-proxy/xtream-proxy.go
+++ b/pkg/xtream-proxy/xtream-proxy.go
@@ -19,6 +19,7 @@
 package xtreamproxy
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -47,8 +48,8 @@ type Client struct {
 }
 
 // New new xtream client
-func New(user, password, baseURL string) (*Client, error) {
-	cli, err := xtream.NewClient(user, password, baseURL)
+func New(user, password, baseURL, userAgent string) (*Client, error) {
+	cli, err := xtream.NewClientWithUserAgent(context.Background(), user, password, baseURL, userAgent)
 	if err != nil {
 		return nil, err
 	}