rushlink/internal/db/db.go
2020-04-22 18:25:27 +02:00

205 lines
5.3 KiB
Go

package db
import (
"bytes"
"fmt"
"io"
"log"
"net/http"
"os"
"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 = 3
// 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, fs *FileStore) (*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(func(tx *bolt.Tx) error { return migrate(tx, fs) }); 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.Bolt.Close()
}
// Initialize and migrate the database to the current version
func migrate(tx *bolt.Tx, fs *FileStore) error {
// Guidelines for error handling:
// - Errors based on malformed *structure* should be fatal!
// - Errors based on malformed *data* should print a warning
// (and if possible try to fix the 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
}
}
if dbVersion < 3 {
log.Println("migrating database to version 3")
// In this version, we changed te way how Content-Types are being
// stored. Previously, we allowed clients to provide their own
// Content-Types for files, using the Content-Disposition header in
// multipart forms. The new way detects these types using
// http.DetectContentType.
//
// Scan through all the FileUploads and update their ContentTypes.
bucket := tx.Bucket([]byte(BucketFileUpload))
cursor := bucket.Cursor()
var id, storedBytes []byte
id, storedBytes = cursor.First()
for id != nil {
fu, err := decodeFileUpload(storedBytes)
if err != nil {
log.Print("error: ", errors.Wrapf(err, "corrupted FileUpload in database at '%v'", id))
id, storedBytes = cursor.Next()
continue
}
if fu.State != FileUploadStatePresent {
id, storedBytes = cursor.Next()
continue
}
filePath := fu.Path(fs)
file, err := os.Open(fu.Path(fs))
if err != nil {
log.Print("error: ", errors.Wrapf(err, "could not open file at '%v'", filePath))
id, storedBytes = cursor.Next()
continue
}
var buf bytes.Buffer
buf.Grow(512)
io.CopyN(&buf, file, 512)
contentType := http.DetectContentType(buf.Bytes())
if contentType != fu.ContentType {
fu.ContentType = contentType
fu.Save(tx)
cursor.Seek(id)
}
id, storedBytes = cursor.Next()
}
// Update the version number
if err := setDBVersion(tx, 3); 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)
}