diff --git a/Metadata.go b/Metadata.go new file mode 100644 index 0000000..193e3c7 --- /dev/null +++ b/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) + }) +} diff --git a/Server.go b/Server.go index 78a2c3f..f16c689 100644 --- a/Server.go +++ b/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 } diff --git a/upload.go b/upload.go new file mode 100644 index 0000000..871bc1c --- /dev/null +++ b/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 +}