rushlink/internal/db/paste.go

package db

import (
	"crypto/rand"
	"encoding/hex"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/pkg/errors"
	"gorm.io/gorm"
)

// PasteType describes the type of Paste (i.e. file, redirect, [...]).
type PasteType int

// PasteState describes the state of a Paste (i.e. present, deleted, [...]).
type PasteState int

// Paste describes the main Paste model in the database.
type Paste struct {
	ID          uint       `gorm:"primaryKey"`
	Type        PasteType  `gorm:"index"`
	State       PasteState `gorm:"index"`
  CreatedBy   uint       `gorm:"index"`
	Content     []byte
	Key         string `gorm:"uniqueIndex"`
	DeleteToken string
	CreatedAt   time.Time
	UpdatedAt   time.Time
	DeletedAt   gorm.DeletedAt
}

// ReservedPasteKeys keys are designated reserved, and will not be randomly chosen.
var ReservedPasteKeys = []string{"xd42", "example"}

// Note: we use iota here. That means removals of PasteType* are not allowed,
// because this changes the value of the constant. Please add the comment
// "// deprecated" if you want to remove the constant. Additions are only
// allowed at the bottom of this block, for the same reason.
const (
	PasteTypeUndef PasteType = iota
	// PasteTypePaste is as of yet unused.  It is still unclear if this type
	// will ever get a proper meaning.
	PasteTypePaste
	PasteTypeRedirect
	PasteTypeFileUpload
)

// Note: we use iota here. See the comment above PasteType*
const (
	PasteStateUndef PasteState = iota
	PasteStatePresent
	PasteStateDeleted
)

// minKeyLen specifies the mimimum length of a paste key.
const minKeyLen = 4

var (
	// ErrKeyInvalidChar occurs when a key contains an invalid character.
	ErrKeyInvalidChar = errors.New("invalid character in key")
	// ErrKeyInvalidLength occurs when a key embeds a length that is incorrect.
	ErrKeyInvalidLength = errors.New("key length encoding is incorrect")
	// ErrPasteDoesNotExist occurs when a key does not exist in the database.
	ErrPasteDoesNotExist = errors.New("url key not found in the database")
)

// ErrHTTPStatusCode returns the HTTP status code that should correspond to
// the provided error.
// server error, or false if it is not.
func ErrHTTPStatusCode(err error) int {
	switch err {
	case nil:
		return 0
	case gorm.ErrRecordNotFound, ErrKeyInvalidChar, ErrKeyInvalidLength, ErrPasteDoesNotExist:
		return http.StatusNotFound
	}
	return http.StatusInternalServerError
}

// Base64 encoding and decoding
var base64Alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"

func (t PasteType) String() string {
	switch t {
	case PasteTypeUndef:
		return "unknown"
	case PasteTypePaste:
		return "paste"
	case PasteTypeRedirect:
		return "redirect"
	case PasteTypeFileUpload:
		return "file"
	default:
		return "invalid"
	}
}

func (t PasteState) String() string {
	switch t {
	case PasteStateUndef:
		return "unknown"
	case PasteStatePresent:
		return "present"
	case PasteStateDeleted:
		return "deleted"
	default:
		return "invalid"
	}
}

// GetPaste retrieves a paste from the database.
func GetPaste(db *gorm.DB, key string) (*Paste, error) {
	if err := ValidatePasteKey(key); err != nil {
		return nil, err
	}
	return GetPasteNoValidate(db, key)
}

// ValidatePasteKey validates the format of the key that has
func ValidatePasteKey(key string) error {
	internalLen := minKeyLen
	countingOnes := true
	for _, ch := range key {
		limb := strings.IndexRune(base64Alphabet, ch)
		if limb == -1 {
			return ErrKeyInvalidChar
		}
		for i := 5; i >= 0 && countingOnes; i-- {
			if (limb>>uint(i))&0x1 == 0 {
				countingOnes = false
				break
			}
			internalLen++
		}
	}
	if internalLen != len(key) {
		return ErrKeyInvalidLength
	}
	return nil
}

// GetPasteNoValidate retrieves a paste from the database without validating
// the key format first.
func GetPasteNoValidate(db *gorm.DB, key string) (*Paste, error) {
	var ps []Paste
	if err := db.Unscoped().Limit(1).Where("key = ?", key).Find(&ps).Error; err != nil {
		return nil, err
	}
	if len(ps) == 0 {
		return nil, ErrPasteDoesNotExist
	}
	return &ps[0], nil
}

// Save saves this Paste to the database.
func (p *Paste) Save(db *gorm.DB) error {
	return db.Save(p).Error
}

// Delete deletes this Paste from the database.
func (p *Paste) Delete(db *gorm.DB, fs *FileStore) error {
	// Remove the (maybe) attached file
	if p.Type == PasteTypeFileUpload {
		fuID, err := uuid.FromBytes(p.Content)
		if err != nil {
			return errors.Wrap(err, "failed to parse uuid")
		}
		fu, err := GetFileUpload(db, fuID)
		if err != nil {
			return errors.Wrap(err, "failed to find file in database")
		}
		if err := fu.Delete(db, fs); err != nil {
			return errors.Wrap(err, "failed to remove file")
		}
	}

	// Wipe the old paste
	p.Type = PasteTypeUndef
	p.State = PasteStateDeleted
	p.Content = []byte{}
	if err := db.Save(&p).Error; err != nil {
		return errors.Wrap(err, "failed to wipe paste in database")
	}
	// Soft-delete the paste as well
	if err := db.Delete(&p).Error; err != nil {
		return errors.Wrap(err, "failed to delete paste in database")
	}
	return nil
}

// RedirectURL returns the URL from this paste.
//
// This function assumes that the paste is valid.  If the paste struct is
// corrupted in some way, this function will panic.
func (p *Paste) RedirectURL() *url.URL {
	if p.Type != PasteTypeRedirect {
		panic("expected p.Type to be PasteTypeRedirect")
	}
	rawurl := string(p.Content)
	urlParse, err := url.Parse(rawurl)
	if err != nil {
		panic(errors.Wrapf(err, "invalid URL ('%v') in database for key '%v'", rawurl, p.Key))
	}
	return urlParse
}

// GeneratePasteKey generates a new paste key. It will ensure that the newly
// generated paste key does not already exist in the database.
// The running time of this function is in O(log N), where N is the amount of
// keys stored in the url-shorten database.
// In tx, a Bolt transaction is given. Use minimumEntropy to set the mimimum
// guessing entropy of the generated key.
func GeneratePasteKey(db *gorm.DB, minimumEntropy int) (string, error) {
	epoch := 0
	var key string
	for {
		var err error
		key, err = generatePasteKeyInner(epoch, minimumEntropy)
		if err != nil {
			return "", errors.Wrap(err, "url-key generation failed")
		}

		var count int64
		db.Unscoped().Model(&Paste{}).Where("key = ?", []byte(key)).Count(&count)
		if err != nil {
			return "", errors.Wrap(err, "failed to check if key already exists")
		}
		alreadyInUse := count != 0

		isReserved := false
		for _, reservedKey := range ReservedPasteKeys {
			if strings.HasPrefix(key, reservedKey) {
				isReserved = true
				break
			}
		}
		if !alreadyInUse && !isReserved {
			break
		}
		epoch++
	}
	return key, nil
}

// generatePasteKeyInner generates a new paste key, but leaves the
// uniqueness and is-reserved checks to the caller. That is, it only
// generates a random key in the correct (syntactical) format.
// Both epoch and entropy can be used to set the key length. Epoch is used
// to prevent collisions in retrying to generate new keys. Entropy (in bits)
// is used to ensure that a new key has at least some amount of guessing
// entropy.
func generatePasteKeyInner(epoch, entropy int) (string, error) {
	entropyEpoch := entropy
	entropyEpoch -= minKeyLen * 6         // First 4 characters provide 24 bits.
	entropyEpoch++                        // One bit less because of '0' bit.
	entropyEpoch = (entropyEpoch-1)/5 + 1 // 5 bits for every added epoch.
	if epoch < entropyEpoch {
		epoch = entropyEpoch
	}
	urlKey := make([]byte, minKeyLen+epoch)
	_, err := rand.Read(urlKey)
	if err != nil {
		return "", err
	}
	// Put all the values in the range 0..64 for easier base64-encoding
	for i := 0; i < len(urlKey); i++ {
		urlKey[i] &= 0x3F
	}
	// Implement truncate-resistance by forcing the prefix to
	//     0b111110xxxxxxxxxx
	//       ^----- {epoch} ones followed by a single 0
	//
	// Example when epoch is 1: prefix is 0b10.
	i := 0
	for i < epoch {
		// Set this bit to 1
		limb := i / 6
		bit := i % 6
		urlKey[limb] |= 1 << uint(5-bit)
		i++
	}
	// Finally set the next bit to 0
	limb := i / 6
	bit := i % 6
	urlKey[limb] &= ^(1 << uint(5-bit))

	// Convert this ID to a canonical base64 notation
	for i := range urlKey {
		urlKey[i] = base64Alphabet[urlKey[i]]
	}
	return string(urlKey), nil
}

// GenerateDeleteToken generates a new (random) delete token.
func GenerateDeleteToken() (string, error) {
	var deleteToken [16]byte
	_, err := rand.Read(deleteToken[:])
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(deleteToken[:]), nil
}