diff --git a/catalyst-uploader.go b/catalyst-uploader.go index 6eff256..8cc8915 100644 --- a/catalyst-uploader.go +++ b/catalyst-uploader.go @@ -42,6 +42,7 @@ func run() int { segTimeout := fs.Duration("segment-timeout", 5*time.Minute, "Segment write timeout") disableRecording := CommaSliceFlag(fs, "disable-recording", `Comma-separated list of playbackIDs to disable recording for`) disableThumbs := CommaSliceFlag(fs, "disable-thumbs", `Comma-separated list of playbackIDs to disable thumbs for`) + thumbsURLReplacement := CommaMapFlag(fs, "thumbs-replace-urls", `Map of space separated playbackIDs to space separated URL replacement to use when saving thumbnails. E.g. playbackID1 playbackID2=oldURL newURL`) defaultConfigFile := "/etc/livepeer/catalyst_uploader.conf" if _, err := os.Stat(defaultConfigFile); os.IsNotExist(err) { @@ -50,6 +51,7 @@ func run() int { _ = fs.String("config", defaultConfigFile, "config file (optional)") err = ff.Parse(fs, os.Args[1:], + ff.WithIgnoreUndefined(true), ff.WithConfigFileFlag("config"), ff.WithConfigFileParser(ff.PlainParser), ff.WithEnvVarPrefix("CATALYST_UPLOADER"), @@ -110,7 +112,7 @@ func run() int { } start := time.Now() - out, err := core.Upload(os.Stdin, uri, WaitBetweenWrites, *timeout, *storageFallbackURLs, *segTimeout, *disableThumbs) + out, err := core.Upload(os.Stdin, uri, WaitBetweenWrites, *timeout, *storageFallbackURLs, *segTimeout, *disableThumbs, *thumbsURLReplacement) if err != nil { glog.Errorf("Uploader failed for %s: %s", uri.Redacted(), err) return 1 diff --git a/core/uploader.go b/core/uploader.go index 1f4b918..d2af9f6 100644 --- a/core/uploader.go +++ b/core/uploader.go @@ -53,7 +53,7 @@ var expiryField = map[string]string{ "Object-Expires": "+168h", // Objects will be deleted after 7 days } -func Upload(input io.Reader, outputURI *url.URL, waitBetweenWrites, writeTimeout time.Duration, storageFallbackURLs map[string]string, segTimeout time.Duration, disableThumbs []string) (*drivers.SaveDataOutput, error) { +func Upload(input io.Reader, outputURI *url.URL, waitBetweenWrites, writeTimeout time.Duration, storageFallbackURLs map[string]string, segTimeout time.Duration, disableThumbs []string, thumbsURLReplacement map[string]string) (*drivers.SaveDataOutput, error) { ext := filepath.Ext(outputURI.Path) inputFile, err := os.CreateTemp("", "upload-*"+ext) if err != nil { @@ -78,7 +78,7 @@ func Upload(input io.Reader, outputURI *url.URL, waitBetweenWrites, writeTimeout return nil, fmt.Errorf("failed to upload video %s: (%d bytes) %w", outputURI.Redacted(), bytesWritten, err) } - if err = extractThumb(outputURI, inputFileName, storageFallbackURLs, disableThumbs); err != nil { + if err = extractThumb(outputURI, inputFileName, storageFallbackURLs, disableThumbs, thumbsURLReplacement); err != nil { glog.Errorf("extracting thumbnail failed for %s: %v", outputURI.Redacted(), err) } return out, nil @@ -229,13 +229,33 @@ func uploadFile(outputURI *url.URL, fileName string, fields *drivers.FilePropert return out, bytesWritten, err } -func extractThumb(outputURI *url.URL, segmentFileName string, storageFallbackURLs map[string]string, disableThumbs []string) error { +func extractThumb(outputURI *url.URL, segmentFileName string, storageFallbackURLs map[string]string, disableThumbs []string, thumbsURLReplacement map[string]string) error { for _, playbackID := range disableThumbs { if strings.Contains(outputURI.Path, playbackID) { glog.Infof("Thumbnails disabled for %s", outputURI.Redacted()) return nil } } + for playbackIDs, replacement := range thumbsURLReplacement { + for _, playbackID := range strings.Split(playbackIDs, " ") { + if strings.Contains(outputURI.Path, playbackID) { + outputURIStr := outputURI.String() + split := strings.Split(replacement, " ") + if len(split) != 2 { + break + } + original, replaceWith := split[0], split[1] + + newURI, err := url.Parse(strings.Replace(outputURIStr, original, replaceWith, 1)) + if err != nil { + return fmt.Errorf("failed to parse thumbnail URL: %w", err) + } + outputURI = newURI + glog.Infof("Replaced thumbnail location for %s", outputURI.Redacted()) + break + } + } + } tmpDir, err := os.MkdirTemp(os.TempDir(), "thumb-*") if err != nil { diff --git a/core/uploader_test.go b/core/uploader_test.go index a917f00..e7a18e2 100644 --- a/core/uploader_test.go +++ b/core/uploader_test.go @@ -36,7 +36,7 @@ func TestItWritesSlowInputIncrementally(t *testing.T) { go func() { u, err := url.Parse(outputFile.Name()) require.NoError(t, err) - _, err = Upload(slowReader, u, 100*time.Millisecond, time.Second, nil, time.Minute, nil) + _, err = Upload(slowReader, u, 100*time.Millisecond, time.Second, nil, time.Minute, nil, nil) require.NoError(t, err, "") }()