From eec5e4def4349d7ad8dee123358d4aba3d76ca72 Mon Sep 17 00:00:00 2001
From: Daan Sprenkels <daan@dsprenkels.com>
Date: Tue, 10 Dec 2019 11:16:18 +0100
Subject: [PATCH] Detect file types instead of trusting clients

---
 handlers.go               |  7 ++---
 internal/db/fileupload.go | 60 ++++++++++++++++++++++++++++-----------
 2 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/handlers.go b/handlers.go
index 66ff5e9..554a93f 100644
--- a/handlers.go
+++ b/handlers.go
@@ -55,11 +55,11 @@ func (rl *rushlink) uploadFileGetHandler(w http.ResponseWriter, r *http.Request)
 		panic(err)
 	}
 
-	filePath := rl.fs.FilePath(fu.ID, fu.FileName)
+	filePath := fu.Path(rl.fs)
 	file, err := os.Open(filePath)
 	if err != nil {
 		if os.IsNotExist(err) {
-			log.Printf("error: %v should exist according to the database, but it doesn't", filePath)
+			log.Printf("error: '%v' should exist according to the database, but it doesn't", filePath)
 			renderError(w, r, http.StatusNotFound, "file not found")
 			return
 		}
@@ -192,8 +192,7 @@ func (rl *rushlink) newFileUploadPasteHandler(w http.ResponseWriter, r *http.Req
 	var paste *db.Paste
 	if err := rl.db.Bolt.Update(func(tx *bolt.Tx) error {
 		var err error
-		// Create the fileUpload in the database
-		fu, err = db.NewFileUpload(rl.fs, file, header.Filename, header.Header.Get("Content-Type"))
+		fu, err = db.NewFileUpload(rl.fs, file, header.Filename)
 		if err != nil {
 			panic(errors.Wrap(err, "creating fileUpload"))
 		}
diff --git a/internal/db/fileupload.go b/internal/db/fileupload.go
index f69be39..0d42bbe 100644
--- a/internal/db/fileupload.go
+++ b/internal/db/fileupload.go
@@ -1,9 +1,11 @@
 package db
 
 import (
+	"bytes"
 	"encoding/hex"
 	"hash/crc32"
 	"io"
+	"net/http"
 	"net/url"
 	"os"
 	"path"
@@ -76,14 +78,42 @@ func OpenFileStore(path string) (*FileStore, error) {
 	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.
-func NewFileUpload(fs *FileStore, r io.Reader, fileName string, contentType string) (*FileUpload, error) {
+//
+// 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")
 	}
 
-	filePath := fs.FilePath(id, fileName)
+	// 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")
 	}
@@ -93,8 +123,11 @@ func NewFileUpload(fs *FileStore, r io.Reader, fileName string, contentType stri
 	}
 	defer file.Close()
 
-	hash := crc32.New(checksumTable)
-	tee := io.TeeReader(r, hash)
+	// 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")
@@ -110,17 +143,7 @@ func NewFileUpload(fs *FileStore, r io.Reader, fileName string, contentType stri
 	return fu, nil
 }
 
-func (fs *FileStore) Path() string {
-	return fs.path
-}
-
-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)
-}
-
+// 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 {
@@ -155,7 +178,7 @@ func (fu *FileUpload) Save(tx *bolt.Tx) error {
 // Delete deletes a FileUpload from the database.
 func (fu *FileUpload) Delete(tx *bolt.Tx, fs *FileStore) error {
 	// Remove the file in the backend
-	filePath := fs.FilePath(fu.ID, fu.FileName)
+	filePath := fu.Path(fs)
 	if err := os.Remove(filePath); err != nil {
 		return err
 	}
@@ -173,6 +196,11 @@ func (fu *FileUpload) Delete(tx *bolt.Tx, fs *FileStore) error {
 	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