rushlink/internal/db/paste.go
2019-12-06 18:03:56 +01:00

237 lines
5.6 KiB
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
}