stash/pkg/manager/downloads.go

package manager

import (
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/stashapp/stash/pkg/logger"
	"github.com/stashapp/stash/pkg/utils"
)

// DownloadStore manages single-use generated files for the UI to download.
type DownloadStore struct {
	m     map[string]*storeFile
	mutex sync.Mutex
}

type storeFile struct {
	path        string
	contentType string
	keep        bool
	wg          sync.WaitGroup
	once        sync.Once
}

func NewDownloadStore() *DownloadStore {
	return &DownloadStore{
		m: make(map[string]*storeFile),
	}
}

func (s *DownloadStore) RegisterFile(fp string, contentType string, keep bool) string {
	const keyLength = 4
	const attempts = 100

	// keep generating random keys until we get a free one
	// prevent infinite loop by only attempting a finite amount of times
	var hash string
	generate := true
	a := 0

	s.mutex.Lock()
	for generate && a < attempts {
		hash = utils.GenerateRandomKey(keyLength)
		_, generate = s.m[hash]
		a++
	}

	s.m[hash] = &storeFile{
		path:        fp,
		contentType: contentType,
		keep:        keep,
	}
	s.mutex.Unlock()

	return hash
}

func (s *DownloadStore) Serve(hash string, w http.ResponseWriter, r *http.Request) {
	s.mutex.Lock()
	f, ok := s.m[hash]

	if !ok {
		s.mutex.Unlock()
		http.NotFound(w, r)
		return
	}

	if !f.keep {
		s.waitAndRemoveFile(hash, &w, r)
	}

	s.mutex.Unlock()

	if f.contentType != "" {
		w.Header().Add("Content-Type", f.contentType)
	}
	http.ServeFile(w, r, f.path)
}

func (s *DownloadStore) waitAndRemoveFile(hash string, w *http.ResponseWriter, r *http.Request) {
	f := s.m[hash]
	notify := r.Context().Done()
	f.wg.Add(1)

	go func() {
		<-notify
		s.mutex.Lock()
		defer s.mutex.Unlock()

		f.wg.Done()
	}()

	go f.once.Do(func() {
		// leave it up for 30 seconds after the first request to allow for multiple requests
		time.Sleep(30 * time.Second)

		f.wg.Wait()
		s.mutex.Lock()
		defer s.mutex.Unlock()

		delete(s.m, hash)
		err := os.Remove(f.path)
		if err != nil {
			logger.Errorf("error removing %s after downloading: %s", f.path, err.Error())
		}
	})
}