separate metadata/upload into separate files, handle mime types safely, preserve names on download (only)

Metadata.go

@@ -0,0 +1,45 @@
+package contented
+
+import (
+	"encoding/json"
+	"os"
+	"time"
+
+	"github.com/boltdb/bolt"
+)
+
+type Metadata struct {
+	FileSize   int64
+	UploadTime time.Time
+	UploadIP   string
+	Filename   string
+	MimeType   string
+}
+
+func (this *Server) Metadata(fileID string) (*Metadata, error) {
+	var m Metadata
+	err := this.db.View(func(tx *bolt.Tx) error {
+		content := tx.Bucket(this.metadataBucket).Get([]byte(fileID))
+		if len(content) == 0 {
+			return os.ErrNotExist
+		}
+
+		return json.Unmarshal(content, &m)
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &m, nil
+}
+
+func (this *Server) SetMetadata(fileID string, m Metadata) error {
+	jb, err := json.Marshal(m)
+	if err != nil {
+		return err
+	}
+
+	return this.db.Update(func(tx *bolt.Tx) error {
+		return tx.Bucket(this.metadataBucket).Put([]byte(fileID), jb)
+	})
+}

Server.go

@@ -1,17 +1,12 @@
 package contented
 
 import (
-	"crypto/sha512"
-	"encoding/hex"
 	"encoding/json"
-	"io"
 	"log"
-	"mime/multipart"
 	"net/http"
 	"os"
 	"path/filepath"
 	"strings"
-	"time"
 
 	"github.com/boltdb/bolt"
 )
@@ -55,99 +50,6 @@ func NewServer(opts *ServerOptions) (*Server, error) {
 	return s, nil
 }
 
-type Metadata struct {
-	FileSize   int64
-	UploadTime time.Time
-	UploadIP   string
-	Filename   string
-}
-
-func (this *Server) handleUpload(src multipart.File, hdr *multipart.FileHeader, UploadIP string) (string, error) {
-
-	// Get file length
-	srcLen, err := src.Seek(0, io.SeekEnd)
-	if err != nil {
-		return "", err
-	}
-
-	_, err = src.Seek(0, io.SeekStart)
-	if err != nil {
-		return "", err
-	}
-
-	// Get file hash
-	hasher := sha512.New512_256()
-	_, err = io.CopyN(hasher, src, int64(srcLen))
-	if err != nil {
-		return "", err
-	}
-
-	_, err = src.Seek(0, io.SeekStart)
-	if err != nil {
-		return "", err
-	}
-
-	// Save file to disk
-	fileID := hex.EncodeToString(hasher.Sum(nil))
-	dest, err := os.OpenFile(filepath.Join(this.opts.DataDirectory, fileID), os.O_CREATE|os.O_WRONLY, 0600)
-	if err != nil {
-		if os.IsExist(err) {
-			return fileID, nil // hash matches existing upload
-		}
-
-		return "", err
-	}
-	defer dest.Close()
-
-	_, err = io.CopyN(dest, src, int64(srcLen))
-	if err != nil {
-		return "", err
-	}
-
-	// Persist metadata to DB
-	m := Metadata{
-		Filename:   hdr.Filename,
-		UploadTime: time.Now(),
-		UploadIP:   UploadIP,
-		FileSize:   srcLen,
-	}
-	err = this.SetMetadata(fileID, m)
-	if err != nil {
-		return "", err
-	}
-
-	// Done
-	return fileID, nil
-}
-
-func (this *Server) Metadata(fileID string) (*Metadata, error) {
-	var m Metadata
-	err := this.db.View(func(tx *bolt.Tx) error {
-		content := tx.Bucket(this.metadataBucket).Get([]byte(fileID))
-		if len(content) == 0 {
-			return os.ErrNotExist
-		}
-
-		return json.Unmarshal(content, &m)
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	return &m, nil
-}
-
-func (this *Server) SetMetadata(fileID string, m Metadata) error {
-	jb, err := json.Marshal(m)
-	if err != nil {
-		return err
-	}
-
-	return this.db.Update(func(tx *bolt.Tx) error {
-		return tx.Bucket(this.metadataBucket).Put([]byte(fileID), jb)
-	})
-}
-
 func (this *Server) handleView(w http.ResponseWriter, r *http.Request, fileID string) error {
 	// Load file
 	f, err := os.Open(filepath.Join(this.opts.DataDirectory, fileID))
@@ -163,7 +65,15 @@ func (this *Server) handleView(w http.ResponseWriter, r *http.Request, fileID st
 		return err
 	}
 
-	http.ServeContent(w, r, m.Filename, m.UploadTime, f)
+	// ServeContent only uses the filename to get the mime type, which we can
+	// set accurately (including blacklist)
+	w.Header().Set(`Content-Type`, m.MimeType)
+	if m.MimeType == `application/octet-stream` {
+		w.Header().Set(`Content-Disposition`, `attachment; filename="`+m.Filename+`"`)
+	}
+
+	http.ServeContent(w, r, "", m.UploadTime, f)
+
 	return nil
 }
 

upload.go

@@ -0,0 +1,89 @@
+package contented
+
+import (
+	"crypto/sha512"
+	"encoding/hex"
+	"io"
+	"mime"
+	"mime/multipart"
+	"os"
+	"path"
+	"path/filepath"
+	"strings"
+	"time"
+)
+
+func (this *Server) handleUpload(src multipart.File, hdr *multipart.FileHeader, UploadIP string) (string, error) {
+
+	// Get file length
+	srcLen, err := src.Seek(0, io.SeekEnd)
+	if err != nil {
+		return "", err
+	}
+
+	_, err = src.Seek(0, io.SeekStart)
+	if err != nil {
+		return "", err
+	}
+
+	// Get file hash
+	hasher := sha512.New512_256()
+	_, err = io.CopyN(hasher, src, int64(srcLen))
+	if err != nil {
+		return "", err
+	}
+
+	_, err = src.Seek(0, io.SeekStart)
+	if err != nil {
+		return "", err
+	}
+
+	// Save file to disk
+	fileID := hex.EncodeToString(hasher.Sum(nil))
+	dest, err := os.OpenFile(filepath.Join(this.opts.DataDirectory, fileID), os.O_CREATE|os.O_WRONLY, 0600)
+	if err != nil {
+		if os.IsExist(err) {
+			return fileID, nil // hash matches existing upload
+		}
+
+		return "", err
+	}
+	defer dest.Close()
+
+	_, err = io.CopyN(dest, src, int64(srcLen))
+	if err != nil {
+		return "", err
+	}
+
+	// Determine mime type
+	ctype := hdr.Header.Get("Content-Type")
+	if ctype == "" {
+		ctype = mime.TypeByExtension(path.Ext(hdr.Filename))
+	}
+	if ctype == "" {
+		ctype = `application/octet-stream`
+	}
+
+	// Blacklisted mime-types
+	ctype = strings.ToLower(ctype)
+	if strings.HasPrefix(ctype, `text/html`) ||
+		strings.HasPrefix(ctype, `application/javascript`) {
+		ctype = `application/octet-stream`
+	}
+
+	// Persist metadata to DB
+	m := Metadata{
+		Filename:   hdr.Filename,
+		UploadTime: time.Now(),
+		UploadIP:   UploadIP,
+		FileSize:   srcLen,
+		MimeType:   ctype,
+	}
+	err = this.SetMetadata(fileID, m)
+	if err != nil {
+		return "", err
+	}
+
+	// Done
+	return fileID, nil
+}