rushlink/internal/db/db.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, 0660, &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)
}