rushlink/internal/db/fileupload.go

package db

import (
	"bytes"
	"encoding/hex"
	"hash/crc32"
	"io"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"

	"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 of the FileUpload (present/deleted/etc).
	State FileUploadState

	// ID identifies this FileUpload.
	ID uuid.UUID

	// FileName contains the original filename of this FileUpload.
	FileName string

	// Content type as determined by http.DetectContentType.
	ContentType string

	// Checksum holds a crc32c checksum of the file.
	//
	// This checksum is only meant to allow for the detection of random
	// database corruption.
	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
}

// Path returns the path of the FileStore root.
func (fs *FileStore) Path() string {
	return fs.path
}

// filePath resolves the path of a file in the FileStore given some id and filename.
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)
}

// NewFileUpload creates a new FileUpload object.
//
// Internally, this function detects the type of the file stored in `r` using
// `http.DetectContentType`.
func NewFileUpload(fs *FileStore, r io.Reader, fileName string) (*FileUpload, error) {
	// Generate a file ID
	id, err := uuid.NewRandom()
	if err != nil {
		return nil, errors.Wrap(err, "generating UUID")
	}

	// Construct a checksum for this file
	hash := crc32.New(checksumTable)
	tee := io.TeeReader(r, hash)

	// Detect the file type
	var tmpBuf bytes.Buffer
	tmpBuf.Grow(512)
	io.CopyN(&tmpBuf, tee, 512)
	contentType := http.DetectContentType(tmpBuf.Bytes())

	// Open the file on disk for writing
	baseName := filepath.Base(fileName)
	filePath := fs.filePath(id, baseName)
	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()

	// Write the file to disk
	_, err = io.Copy(file, &tmpBuf)
	if err != nil {
		return nil, errors.Wrap(err, "writing to file")
	}
	_, err = io.Copy(file, tee)
	if err != nil {
		return nil, errors.Wrap(err, "writing to file")
	}

	fu := &FileUpload{
		State:       FileUploadStatePresent,
		ID:          id,
		FileName:    baseName,
		ContentType: contentType,
		Checksum:    hash.Sum32(),
	}
	return fu, nil
}

// GetFileUpload tries to retrieve a FileUpload object from the bolt database.
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
	}
	return decodeFileUpload(storedBytes)
}

// AllFileUploads tries to retrieve all FileUpload objects from the bolt database.
func AllFileUploads(tx *bolt.Tx) ([]FileUpload, error) {
	bucket := tx.Bucket([]byte(BucketFileUpload))
	if bucket == nil {
		return nil, errors.Errorf("bucket %v does not exist", BucketFileUpload)
	}
	var fus []FileUpload
	err := bucket.ForEach(func(_, storedBytes []byte) error {
		fu, err := decodeFileUpload(storedBytes)
		if err != nil {
			return err
		}
		fus = append(fus, *fu)
		return nil
	})
	if err != nil {
		return nil, err
	}
	return fus, nil
}

func decodeFileUpload(storedBytes []byte) (*FileUpload, error) {
	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 := fu.Path(fs)
	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)
}

// Path returns the path to this FileUpload in the FileStore provided in fs.
func (fu *FileUpload) Path(fs *FileStore) string {
	return fs.filePath(fu.ID, fu.FileName)
}

// 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
}

// Ext returns the extension of the file attached to this FileUpload.
func (fu *FileUpload) Ext() string {
	return filepath.Ext(fu.FileName)
}