rushlink/paste.go

package rushlink

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

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

const (
	pasteTypeUndef      pasteType = 0
	pasteTypePaste                = 1
	pasteTypeRedirect             = 2
	pasteTypeFileUpload           = 3
)

const (
	pasteStateUndef   pasteState = 0
	pasteStatePresent            = 1
	pasteStateDeleted            = 2
)

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

// 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 {
	// 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); 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
}

// Get 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
}

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