rushlink/internal/db/paste.go

package db

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

	gobmarsh "gitea.hashru.nl/dsprenkels/rushlink/pkg/gobmarsh"
	"github.com/google/uuid"
	"github.com/pkg/errors"
	bolt "go.etcd.io/bbolt"
)

type PasteType int
type PasteState int

type Paste struct {
	Type        PasteType
	State       PasteState
	Content     []byte
	Key         string
	DeleteToken string
	TimeCreated time.Time
}

// 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
	PasteTypeRedirect
	PasteTypeFileUpload
)

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

// Base64 encoding and decoding
var base64Alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
var base64Encoder = base64.RawURLEncoding.WithPadding(base64.NoPadding)

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(tx *bolt.Tx, key string) (*Paste, error) {
	pastesBucket := tx.Bucket([]byte(BucketPastes))
	if pastesBucket == nil {
		return nil, errors.Errorf("bucket %v does not exist", BucketPastes)
	}
	storedBytes := pastesBucket.Get([]byte(key))
	if storedBytes == nil {
		return nil, nil
	}
	p := &Paste{}
	err := gobmarsh.Unmarshal(storedBytes, p)
	return p, err
}

func (p *Paste) Save(tx *bolt.Tx) error {
	bucket := tx.Bucket([]byte(BucketPastes))
	if bucket == nil {
		return errors.Errorf("bucket %v does not exist", BucketPastes)
	}

	buf, err := gobmarsh.Marshal(p)
	if err != nil {
		return errors.Wrap(err, "encoding for database failed")
	}
	if err := bucket.Put([]byte(p.Key), buf); err != nil {
		return errors.Wrap(err, "database transaction failed")
	}
	return nil
}

func (p *Paste) Delete(tx *bolt.Tx, 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(tx, fuID)
		if err != nil {
			return errors.Wrap(err, "failed to find file in database")
		}
		if err := fu.Delete(tx, fs); err != nil {
			return errors.Wrap(err, "failed to remove file")
		}
	}

	// Replace the old paste with a new empty paste
	p.Type = PasteTypeUndef
	p.State = PasteStateDeleted
	p.Content = []byte{}
	if err := p.Save(tx); 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 key until it is not 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.
func GeneratePasteKey(tx *bolt.Tx) (string, error) {
	pastesBucket := tx.Bucket([]byte(BucketPastes))
	if pastesBucket == nil {
		return "", errors.Errorf("bucket %v does not exist", BucketPastes)
	}

	epoch := 0
	var key string
	for {
		var err error
		key, err = generatePasteKeyInner(epoch)
		if err != nil {
			return "", errors.Wrap(err, "url-key generation failed")
		}

		found := pastesBucket.Get([]byte(key))
		if found == nil {
			break
		}

		isReserved := false
		for _, reservedKey := range ReservedPasteKeys {
			if strings.HasPrefix(key, reservedKey) {
				isReserved = true
				break
			}
		}
		if !isReserved {
			break
		}

		epoch++
	}
	return key, nil
}

func generatePasteKeyInner(epoch int) (string, error) {
	urlKey := make([]byte, 4+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
}

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