rushlink/paste.go

package rushlink

import (
	"crypto/rand"
	"encoding/hex"
	"strings"
	"time"

	"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
}

const (
	typeUndef    pasteType = 0
	typePaste              = 1
	typeRedirect           = 2
)

const (
	stateUndef   pasteState = 0
	statePresent            = 1
	stateDeleted            = 2
)

// Retrieve a paste from the database
func getPaste(tx *bolt.Tx, key string) (*paste, error) {
	pastesBucket := tx.Bucket([]byte(BUCKET_PASTES))
	if pastesBucket == nil {
		return nil, errors.Errorf("bucket %v does not exist", BUCKET_PASTES)
	}
	storedBytes := pastesBucket.Get([]byte(key))
	if storedBytes == nil {
		return nil, nil
	}
	p := &paste{}
	err := Unmarshal(storedBytes, p)
	return p, err
}

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

	buf, err := 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) error {
	// Replace the old paste with a new empty paste
	return (&paste{
		Key:         p.Key,
		State:       stateDeleted,
		DeleteToken: p.DeleteToken,
	}).save(tx)
}

// Generate a key until it is not in the database, this occurs 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(BUCKET_PASTES))
	if pastesBucket == nil {
		return "", errors.Errorf("bucket %v does not exist", BUCKET_PASTES)
	}

	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
}