From 62e82d831ee1d578387d8b0e862fd3056ffb4ca4 Mon Sep 17 00:00:00 2001
From: Daan Sprenkels <daan@dsprenkels.com>
Date: Tue, 10 Dec 2019 11:59:02 +0100
Subject: [PATCH] db: Migrate FileUpload.ContentTypes to auto-detect

---
 cmd/rushlink/main.go      | 10 +++---
 internal/db/db.go         | 65 ++++++++++++++++++++++++++++++++++++---
 internal/db/fileupload.go |  4 +++
 3 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/cmd/rushlink/main.go b/cmd/rushlink/main.go
index 32dd43a..98cad7a 100644
--- a/cmd/rushlink/main.go
+++ b/cmd/rushlink/main.go
@@ -18,15 +18,15 @@ var (
 func main() {
 	flag.Parse()
 
-	database, err := db.OpenDB(*databasePath)
-	if err != nil {
-		log.Fatalln(err)
-	}
-	defer database.Close()
 	filestore, err := db.OpenFileStore(*fileStorePath)
 	if err != nil {
 		log.Fatalln(err)
 	}
+	database, err := db.OpenDB(*databasePath, filestore)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	defer database.Close()
 
 	go rushlink.StartMetricsServer(*metricsListen, database, filestore)
 	rushlink.StartMainServer(*httpListen, database, filestore)
diff --git a/internal/db/db.go b/internal/db/db.go
index 9531fa1..72d216d 100644
--- a/internal/db/db.go
+++ b/internal/db/db.go
@@ -1,8 +1,12 @@
 package db
 
 import (
+	"bytes"
 	"fmt"
+	"io"
 	"log"
+	"net/http"
+	"os"
 	"time"
 
 	"github.com/pkg/errors"
@@ -24,7 +28,7 @@ type Database struct {
 //
 // If we alter the database format, we bump this number and write a new
 // database migration in migrate().
-const CurrentMigrateVersion = 2
+const CurrentMigrateVersion = 3
 
 // BucketConf holds the name for the "configuration" bucket.
 //
@@ -42,7 +46,7 @@ const BucketFileUpload = "fileUpload"
 const KeyMigrateVersion = "migrate_version"
 
 // OpenDB opens a database file located at path.
-func OpenDB(path string) (*Database, error) {
+func OpenDB(path string, fs *FileStore) (*Database, error) {
 	if path == "" {
 		return nil, errors.New("database not set")
 	}
@@ -51,7 +55,7 @@ func OpenDB(path string) (*Database, error) {
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to open database at '%v'", path)
 	}
-	if err := db.Update(migrate); err != nil {
+	if err := db.Update(func(tx *bolt.Tx) error { return migrate(tx, fs) }); err != nil {
 		return nil, err
 	}
 	return &Database{db}, nil
@@ -66,7 +70,12 @@ func (db *Database) Close() error {
 }
 
 // Initialize and migrate the database to the current version
-func migrate(tx *bolt.Tx) error {
+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
@@ -104,6 +113,54 @@ func migrate(tx *bolt.Tx) error {
 		}
 	}
 
+	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
 }
 
diff --git a/internal/db/fileupload.go b/internal/db/fileupload.go
index 0d42bbe..45ceaf3 100644
--- a/internal/db/fileupload.go
+++ b/internal/db/fileupload.go
@@ -153,6 +153,10 @@ func GetFileUpload(tx *bolt.Tx, id uuid.UUID) (*FileUpload, error) {
 	if storedBytes == nil {
 		return nil, nil
 	}
+	return decodeFileUpload(storedBytes)
+}
+
+func decodeFileUpload(storedBytes []byte) (*FileUpload, error) {
 	fu := &FileUpload{}
 	err := gobmarsh.Unmarshal(storedBytes, fu)
 	return fu, err