depot/web/fup/fuphttp/fngen/fngen.go

// SPDX-FileCopyrightText: 2021 Luke Granger-Brown <depot@lukegb.com>
//
// SPDX-License-Identifier: Apache-2.0

package fngen

import (
	"context"
	"crypto/rand"
	"fmt"
	"strings"

	petname "github.com/dustinkirkland/golang-petname"
)

// FilenameGenerator generates a new filename based on the provided prefix and extension.
// The full filename, including the provided filenameExt, should be returned.
// attempt will be 0 the first time the filename is generated;
type FilenameGenerator func(filenamePrefix string, filenameExt string, attempt int) string

// IdentityGenerator returns a FilenameGenerator that, on the first attempt, will use the provided filename.
// If it collides with an existing filename, then the provided fallback FilenameGenerator will be used.
func IdentityGenerator(fallback FilenameGenerator) FilenameGenerator {
	return func(filenamePrefix string, filenameExt string, attempt int) string {
		if attempt > 0 {
			return fallback(filenamePrefix, filenameExt, attempt-1)
		}
		return filenamePrefix + filenameExt
	}
}

func init() {
	petname.NonDeterministicMode()
}

// PetnameGenerator uses RFC1178 to generate filenames.
// This is sort-of "imgur style".
func PetnameGenerator(filenamePrefix string, filenameExt string, attempt int) string {
	bits := strings.Split(petname.Generate(3, "-"), "-")
	return fmt.Sprintf("%s%s%s%s",
		strings.Title(bits[0]),
		strings.Title(bits[1]),
		strings.Title(bits[2]),
		filenameExt)
}

var boringAlphabet = []byte("0123456789abcdefghijklmnopqrstuvwxyz")

const boringLength = 8

// BoringGenerator just uses a simple lowercase alphabet to generate filenames.
func BoringGenerator(filenamePrefix string, filenameExt string, attempt int) string {
	out := make([]byte, boringLength+len(filenameExt))
	if _, err := rand.Read(out[:boringLength]); err != nil {
		panic(fmt.Sprintf("BoringGenerator's rand.Read: %v", err))
	}
	for n := 0; n < boringLength; n++ {
		out[n] = boringAlphabet[out[n]%byte(len(boringAlphabet))]
	}
	for n := boringLength; n < boringLength+len(filenameExt); n++ {
		out[n] = filenameExt[n-boringLength]
	}
	return string(out)
}

// UniqueName repeatedly attempts to generate a name until existsFn returns false.
func UniqueName(ctx context.Context, existsFn func(ctx context.Context, name string) (bool, error), filenamePrefix, extension string, generator FilenameGenerator) (string, error) {
	// There's a race condition here, because we might generate a currently-not-taken filename and collide with someone uploading with the same name.
	// However, given the usecase of this software, I've decided not to care. I'm unlikely to be uploading enough things concurrently that this is a real issue.
	attempt := 0
	for {
		name := generator(filenamePrefix, extension, attempt)
		if name == "" {
			return "", fmt.Errorf("filename generator returned the empty string")
		}

		exists, err := existsFn(ctx, name)
		if err != nil {
			return "", fmt.Errorf("unable to check for existence of %q: %v", name, err)
		}
		if !exists {
			return name, nil
		}

		attempt++
	}
}