forked from electricdusk/rushlink
Refactor database login into a separate module
This commit is contained in:
Daan Sprenkels
147
internal/db/db.go
Normal file
147
internal/db/db.go
Normal file
@@ -0,0 +1,147 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
|
||||
bolt "go.etcd.io/bbolt"
|
||||
|
||||
gobmarsh "gitea.hashru.nl/dsprenkels/rushlink/pkg/gobmarsh"
|
||||
)
|
||||
|
||||
// Database is the main rushlink database type.
|
||||
//
|
||||
// Open a database using DB.Open() and close it in the end using DB.Close().
|
||||
// Only one instance of DB should exist in a program at any moment.
|
||||
type Database struct {
|
||||
Bolt *bolt.DB
|
||||
}
|
||||
|
||||
// CurrentMigrateVersion holds the current "migrate version".
|
||||
//
|
||||
// If we alter the database format, we bump this number and write a new
|
||||
// database migration in migrate().
|
||||
const CurrentMigrateVersion = 2
|
||||
|
||||
// BucketConf holds the name for the "configuration" bucket.
|
||||
//
|
||||
// This bucket holds the database version, secret site-wide keys, etc.
|
||||
const BucketConf = "conf"
|
||||
|
||||
// BucketPastes holds the name for the pastes bucket.
|
||||
const BucketPastes = "pastes"
|
||||
|
||||
// BucketFileUpload holds the name for the file-upload bucket.
|
||||
const BucketFileUpload = "fileUpload"
|
||||
|
||||
// KeyMigrateVersion stores the current migration version. If this value is less than
|
||||
// CurrentMigrateVersion, the database has to be migrated.
|
||||
const KeyMigrateVersion = "migrate_version"
|
||||
|
||||
// OpenDB opens a database file located at path.
|
||||
func OpenDB(path string) (*Database, error) {
|
||||
if path == "" {
|
||||
return nil, errors.New("database not set")
|
||||
}
|
||||
|
||||
db, err := bolt.Open(path, 0666, &bolt.Options{Timeout: 1 * time.Second})
|
||||
if err != nil {
|
||||
return nil, errors.Wrapf(err, "failed to open database at '%v'", path)
|
||||
}
|
||||
if err := db.Update(migrate); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Database{db}, nil
|
||||
}
|
||||
|
||||
// Close the bolt database
|
||||
func (db *Database) Close() error {
|
||||
if db == nil {
|
||||
panic("no open database")
|
||||
}
|
||||
return db.Close()
|
||||
}
|
||||
|
||||
// Initialize and migrate the database to the current version
|
||||
func migrate(tx *bolt.Tx) error {
|
||||
dbVersion, err := dbVersion(tx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Migrate the database to version 1
|
||||
if dbVersion < 1 {
|
||||
log.Println("migrating database to version 1")
|
||||
// Create conf bucket
|
||||
_, err := tx.CreateBucket([]byte(BucketConf))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Create paste bucket
|
||||
_, err = tx.CreateBucket([]byte(BucketPastes))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Update the version number
|
||||
if err := setDBVersion(tx, 1); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if dbVersion < 2 {
|
||||
log.Println("migrating database to version 2")
|
||||
// Create fileUpload bucket
|
||||
_, err := tx.CreateBucket([]byte(BucketFileUpload))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Update the version number
|
||||
if err := setDBVersion(tx, 2); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get the current migrate version from the database
|
||||
func dbVersion(tx *bolt.Tx) (int, error) {
|
||||
conf := tx.Bucket([]byte(BucketConf))
|
||||
if conf == nil {
|
||||
return 0, nil
|
||||
}
|
||||
dbVersionBytes := conf.Get([]byte(KeyMigrateVersion))
|
||||
if dbVersionBytes == nil {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Version was already stored
|
||||
var dbVersion int
|
||||
if err := gobmarsh.Unmarshal(dbVersionBytes, &dbVersion); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if dbVersion == 0 {
|
||||
return 0, fmt.Errorf("database version is invalid (%v)", dbVersion)
|
||||
}
|
||||
if dbVersion > CurrentMigrateVersion {
|
||||
return 0, fmt.Errorf("database version is too recent (%v > %v)", dbVersion, CurrentMigrateVersion)
|
||||
}
|
||||
return dbVersion, nil
|
||||
}
|
||||
|
||||
// Update the current migrate version in the database
|
||||
func setDBVersion(tx *bolt.Tx, version int) error {
|
||||
conf, err := tx.CreateBucketIfNotExists([]byte(BucketConf))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
versionBytes, err := gobmarsh.Marshal(version)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return conf.Put([]byte(KeyMigrateVersion), versionBytes)
|
||||
}
|
||||
184
internal/db/fileupload.go
Normal file
184
internal/db/fileupload.go
Normal file
@@ -0,0 +1,184 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"encoding/hex"
|
||||
"hash/crc32"
|
||||
"io"
|
||||
"net/url"
|
||||
"os"
|
||||
"path"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/pkg/errors"
|
||||
bolt "go.etcd.io/bbolt"
|
||||
|
||||
gobmarsh "gitea.hashru.nl/dsprenkels/rushlink/pkg/gobmarsh"
|
||||
)
|
||||
|
||||
// Use the Castagnoli checksum because of the acceleration on Intel CPUs
|
||||
var checksumTable = crc32.MakeTable(crc32.Castagnoli)
|
||||
|
||||
// FileStore holds the path to a file storage location.
|
||||
type FileStore struct {
|
||||
path string
|
||||
}
|
||||
|
||||
// FileUploadState determines the current state of a FileUpload object.
|
||||
type FileUploadState int
|
||||
|
||||
// FileUpload models an uploaded file.
|
||||
type FileUpload struct {
|
||||
State FileUploadState
|
||||
ID uuid.UUID
|
||||
FileName string
|
||||
ContentType string
|
||||
Checksum uint32
|
||||
}
|
||||
|
||||
const (
|
||||
dirMode os.FileMode = 0750
|
||||
fileMode os.FileMode = 0640
|
||||
)
|
||||
|
||||
const (
|
||||
// FileUploadStateUndef is an undefined FileUpload.
|
||||
FileUploadStateUndef FileUploadState = 0
|
||||
// FileUploadStatePresent denotes the normal (existing) state.
|
||||
FileUploadStatePresent FileUploadState = 1
|
||||
// FileUploadStateDeleted denotes a deleted state.
|
||||
FileUploadStateDeleted FileUploadState = 2
|
||||
)
|
||||
|
||||
func (t FileUploadState) String() string {
|
||||
switch t {
|
||||
case FileUploadStateUndef:
|
||||
return "unknown"
|
||||
case FileUploadStatePresent:
|
||||
return "present"
|
||||
case FileUploadStateDeleted:
|
||||
return "deleted"
|
||||
default:
|
||||
return "invalid"
|
||||
}
|
||||
}
|
||||
|
||||
// OpenFileStore opens the file storage at path.
|
||||
func OpenFileStore(path string) (*FileStore, error) {
|
||||
if path == "" {
|
||||
return nil, errors.New("file-store not set")
|
||||
}
|
||||
|
||||
// Try to create the file store directory if it does not yet exist
|
||||
if err := os.MkdirAll(path, dirMode); err != nil {
|
||||
return nil, errors.Wrap(err, "creating file store directory")
|
||||
}
|
||||
|
||||
return &FileStore{path[:]}, nil
|
||||
}
|
||||
|
||||
// NewFileUpload creates a new FileUpload object.
|
||||
func NewFileUpload(fs *FileStore, r io.Reader, fileName string, contentType string) (*FileUpload, error) {
|
||||
id, err := uuid.NewRandom()
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "generating UUID")
|
||||
}
|
||||
|
||||
filePath := fs.FilePath(id, fileName)
|
||||
if err := os.Mkdir(path.Dir(filePath), dirMode); err != nil {
|
||||
return nil, errors.Wrap(err, "creating file dir")
|
||||
}
|
||||
file, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE, fileMode)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "opening file")
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
hash := crc32.New(checksumTable)
|
||||
tee := io.TeeReader(r, hash)
|
||||
_, err = io.Copy(file, tee)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "writing to file")
|
||||
}
|
||||
|
||||
fu := &FileUpload{
|
||||
State: FileUploadStatePresent,
|
||||
ID: id,
|
||||
FileName: fileName,
|
||||
ContentType: contentType,
|
||||
Checksum: hash.Sum32(),
|
||||
}
|
||||
return fu, nil
|
||||
}
|
||||
|
||||
func (fs *FileStore) Path() string {
|
||||
return fs.path
|
||||
}
|
||||
|
||||
func (fs *FileStore) FilePath(id uuid.UUID, fileName string) string {
|
||||
if fs.path == "" {
|
||||
panic("fileStoreDir called while the file store path has not been set")
|
||||
}
|
||||
return path.Join(fs.path, hex.EncodeToString(id[:]), fileName)
|
||||
}
|
||||
|
||||
func GetFileUpload(tx *bolt.Tx, id uuid.UUID) (*FileUpload, error) {
|
||||
bucket := tx.Bucket([]byte(BucketFileUpload))
|
||||
if bucket == nil {
|
||||
return nil, errors.Errorf("bucket %v does not exist", BucketFileUpload)
|
||||
}
|
||||
storedBytes := bucket.Get(id[:])
|
||||
if storedBytes == nil {
|
||||
return nil, nil
|
||||
}
|
||||
fu := &FileUpload{}
|
||||
err := gobmarsh.Unmarshal(storedBytes, fu)
|
||||
return fu, err
|
||||
}
|
||||
|
||||
// Save saves a FileUpload in the database.
|
||||
func (fu *FileUpload) Save(tx *bolt.Tx) error {
|
||||
bucket := tx.Bucket([]byte(BucketFileUpload))
|
||||
if bucket == nil {
|
||||
return errors.Errorf("bucket %v does not exist", BucketFileUpload)
|
||||
}
|
||||
|
||||
buf, err := gobmarsh.Marshal(fu)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "encoding for database failed")
|
||||
}
|
||||
if err := bucket.Put(fu.ID[:], buf); err != nil {
|
||||
return errors.Wrap(err, "database transaction failed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Delete deletes a FileUpload from the database.
|
||||
func (fu *FileUpload) Delete(tx *bolt.Tx, fs *FileStore) error {
|
||||
// Remove the file in the backend
|
||||
filePath := fs.FilePath(fu.ID, fu.FileName)
|
||||
if err := os.Remove(filePath); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Update the file in the server
|
||||
if err := (&FileUpload{
|
||||
ID: fu.ID,
|
||||
State: FileUploadStateDeleted,
|
||||
}).Save(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Cleanup the parent directory
|
||||
wrap := "deletion succeeded, but removing the file directory has failed"
|
||||
return errors.Wrap(os.Remove(path.Dir(filePath)), wrap)
|
||||
}
|
||||
|
||||
// URL returns the URL for the FileUpload.
|
||||
func (fu *FileUpload) URL() *url.URL {
|
||||
rawurl := "/uploads/" + hex.EncodeToString(fu.ID[:]) + "/" + fu.FileName
|
||||
urlParse, err := url.Parse(rawurl)
|
||||
if err != nil {
|
||||
panic("could not construct /uploads/ url")
|
||||
}
|
||||
return urlParse
|
||||
}
|
||||
236
internal/db/paste.go
Normal file
236
internal/db/paste.go
Normal file
@@ -0,0 +1,236 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user