rushlink/internal/db/fileupload.go
2019-12-16 10:51:21 +05:30

223 lines
5.6 KiB
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 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
}
// 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
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()
// 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: fileName,
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)
}
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)
}