enhancement: better YouTube headers handling

internetarchive · Sep 16, 2024 · f526b93 · f526b93
1 parent 07690ee
commit f526b93
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 19 deletions.
diff --git a/internal/pkg/crawl/dependencies/ytdlp/model.go b/internal/pkg/crawl/dependencies/ytdlp/model.go
@@ -27,23 +27,23 @@ type Video struct {
 		Fragments     []struct {
 			URL string `json:"url"`
 		} `json:"fragments"`
-		HasDrm             bool        `json:"has_drm"`
-		Height             float64     `json:"height"`
-		HTTPHeaders        HTTPHeaders `json:"http_headers"`
-		Language           interface{} `json:"language"`
-		LanguagePreference float64     `json:"language_preference"`
-		Preference         interface{} `json:"preference"`
-		Protocol           string      `json:"protocol"`
-		Quality            float64     `json:"quality"`
-		Resolution         string      `json:"resolution"`
-		SourcePreference   float64     `json:"source_preference"`
-		Tbr                float64     `json:"tbr"`
-		URL                string      `json:"url"`
-		Vbr                float64     `json:"vbr,omitempty"`
-		Vcodec             string      `json:"vcodec"`
-		VideoExt           string      `json:"video_ext"`
-		Width              float64     `json:"width"`
-		Abr                float64     `json:"abr,omitempty"`
+		HasDrm             bool              `json:"has_drm"`
+		Height             float64           `json:"height"`
+		HTTPHeaders        map[string]string `json:"http_headers"`
+		Language           interface{}       `json:"language"`
+		LanguagePreference float64           `json:"language_preference"`
+		Preference         interface{}       `json:"preference"`
+		Protocol           string            `json:"protocol"`
+		Quality            float64           `json:"quality"`
+		Resolution         string            `json:"resolution"`
+		SourcePreference   float64           `json:"source_preference"`
+		Tbr                float64           `json:"tbr"`
+		URL                string            `json:"url"`
+		Vbr                float64           `json:"vbr,omitempty"`
+		Vcodec             string            `json:"vcodec"`
+		VideoExt           string            `json:"video_ext"`
+		Width              float64           `json:"width"`
+		Abr                float64           `json:"abr,omitempty"`
 	} `json:"requested_formats"`
 	Formats []struct {
 		Acodec      string  `json:"acodec"`

diff --git a/internal/pkg/crawl/dependencies/ytdlp/ytdlp.go b/internal/pkg/crawl/dependencies/ytdlp/ytdlp.go
@@ -8,7 +8,9 @@ import (
 	"strconv"
 )
 
-func GetJSON(port int) (URLs []string, rawJSON string, HTTPHeaders HTTPHeaders, err error) {
+func GetJSON(port int) (URLs []string, rawJSON string, HTTPHeaders map[string]string, err error) {
+	HTTPHeaders = make(map[string]string)
+
 	// Prepare the command
 	cmd := exec.Command("yt-dlp", "--dump-json", "http://localhost:"+strconv.Itoa(port), "-f", "bv[protocol=https]+ba[protocol=https]")
 

diff --git a/internal/pkg/crawl/sitespecific/youtube/youtube.go b/internal/pkg/crawl/sitespecific/youtube/youtube.go
@@ -12,7 +12,9 @@ func IsYouTubeWatchPage(URL *url.URL) bool {
 	return strings.Contains(URL.Host, "youtube.com") && (strings.Contains(URL.Path, "/watch") || strings.Contains(URL.Path, "/v/"))
 }
 
-func Parse(body io.ReadCloser) (URLs []*url.URL, rawJSON string, HTTPHeaders ytdlp.HTTPHeaders, err error) {
+func Parse(body io.ReadCloser) (URLs []*url.URL, rawJSON string, HTTPHeaders map[string]string, err error) {
+	HTTPHeaders = make(map[string]string)
+
 	// Create a temporary server to serve the body and call ytdlp on it
 	port, stopChan, err := ytdlp.ServeBody(body)
 	if err != nil {