-package main

-

-import (

- "crypto/sha1"

- "encoding/base64"

- "encoding/hex"

- "errors"

- "io"

- "io/ioutil"

- "math/rand"

- "mime"

- "os"

- "path"

- "path/filepath"

- "strconv"

- "strings"

- "time"

-)

-

-const (

- MaxIdTries = 64

-

- DefaultIdCharset = "abcdefghijklmnopqrstuvwxyz"

- DefaultIdLength = 6

- DefaultMaxSize = 50 * 1024 * 1024

-)

-

-type Storage struct {

- Folder string

- IdCharset string

- IdLength int

- MaxSize int64

- FilterMime []string

- FilterExt []string

- Whitelist bool

-}

-

-type ErrForbidden struct{ Type string }

-

-func (e ErrForbidden) Error() string { return "forbidden type: " + e.Type }

-

-type ErrTooLarge struct{ Size int64 }

-

-func (e ErrTooLarge) Error() string {

- return "file exceeds maximum allowed size of " + strconv.FormatInt(e.Size, 10) + " bytes"

-}

-

-type ErrNotFound struct{ Name string }

-

-func (e ErrNotFound) Error() string { return "file " + e.Name + " not found" }

-

-func NewStorage(folder string) *Storage {

- if err := os.MkdirAll(path.Join(folder, "temp"), 0755); err != nil {

- panic(err)

- }

- if err := os.MkdirAll(path.Join(folder, "files"), 0755); err != nil {

- panic(err)

- }

- if err := os.MkdirAll(path.Join(folder, "ids"), 0755); err != nil {

- panic(err)

- }

-

- return &Storage{

- Folder: folder,

- IdCharset: DefaultIdCharset,

- IdLength: DefaultIdLength,

- MaxSize: DefaultMaxSize,

- }

-}

-

-func (s *Storage) Get(id string) (file *os.File, hash string, size int64, modtime time.Time, err error) {

- ext := path.Ext(id)

- id = id[:len(id)-len(ext)]

- for i := 0; i < len(id); i++ {

- if !strings.ContainsRune(s.IdCharset, rune(id[i])) {

- err = errors.New("invalid ID: " + id)

- return

- }

- }

- folder := s.idToFolder("ids", id)

- files, err := ioutil.ReadDir(folder)

- if err != nil {

- err = ErrNotFound{id + ext}

- return

- }

- if len(files) < 1 {

- err = errors.New("internal storage error")

- return

- }

- fn := files[0].Name()

- fp := path.Join(folder, fn)

- target, err := os.Readlink(fp)

- if err != nil {

- return

- }

- bhash, err := base64.RawURLEncoding.DecodeString(path.Base(path.Dir(target)))

- if err != nil {

- return

- }

- hash = hex.EncodeToString(bhash)

- if path.Ext(fn) != ext {

- err = ErrNotFound{id + ext}

- return

- }

- stat, err := os.Lstat(fp)

- if err != nil {

- return

- }

- modtime = stat.ModTime()

- size = stat.Size()

- file, err = os.Open(fp)

- return

-}

-

-var errFileExists = errors.New("file exists")

-

-func (s *Storage) New(r io.Reader, name string) (id, hash string, size int64, err error) {

- temp, err := ioutil.TempFile(path.Join(s.Folder, "temp"), "file")

- if err != nil {

- return

- }

- defer func() {

- if temp != nil {

- temp.Close()

- os.Remove(temp.Name())

- }

- }()

-

- hash, size, err = s.readInput(temp, r)

- if err != nil {

- return

- }

- _, ext, err := s.getMimeExt(temp.Name(), name)

- if err != nil {

- return

- }

- id, err = s.storeFile(temp, hash, ext)

- if err == nil {

- temp = nil // prevent deletion

- } else if err == errFileExists {

- err = nil

- }

-

- return

-}

-

-func (s *Storage) randomId() string {

- id := make([]byte, s.IdLength)

- for i := 0; i < len(id); i++ {

- id[i] = s.IdCharset[rand.Intn(len(s.IdCharset))]

- }

- return string(id)

-}

-

-func (s *Storage) idToFolder(subfolder, id string) string {

- name := id

- for len(name) < 4 {

- name = "_" + name

- }

- return path.Join(s.Folder, subfolder, name[0:2], name[2:4], id)

-}

-

-func (s *Storage) readInput(w io.Writer, r io.Reader) (hash string, size int64, err error) {

- h := sha1.New()

- w = io.MultiWriter(h, w)

- if s.MaxSize > 0 {

- r = io.LimitReader(r, s.MaxSize+1)

- }

- size, err = io.Copy(w, r)

- if err != nil {

- return

- }

- if lr, ok := r.(*io.LimitedReader); ok && lr.N == 0 {

- err = ErrTooLarge{s.MaxSize}

- return

- }

- hash = base64.RawURLEncoding.EncodeToString(h.Sum(nil))

- return

-}

-

-func (s *Storage) getMimeExt(fpath string, name string) (mimetype, ext string, err error) {

- mimetype, err = GetMimeType(fpath)

- if err != nil {

- return

- }

-

- // choose file extension, prefer the user-provided one

- ext = path.Ext(name)

- exts, err := mime.ExtensionsByType(mimetype)

- if err != nil {

- return

- }

- if ext != "" && find(exts, ext) == "" {

- ext = ""

- if len(exts) > 0 {

- ext = exts[0]

- }

- }

-

- filtered, ok := s.findFilter(exts, mimetype)

- if !ok && s.Whitelist { // whitelist: reject if not on filters

- err = ErrForbidden{mimetype}

- } else if ok && !s.Whitelist { // blacklist: reject if filtered

- // only block application/octet-stream if explicitly requested

- if mimetype != "application/octet-stream" || find(s.FilterMime, mimetype) != "" {

- err = ErrForbidden{filtered}

- }

- }

-

- return

-}

-

-func (s *Storage) findFilter(exts []string, mimetype string) (match string, ok bool) {

- if m := find(s.FilterMime, mimetype); m == mimetype {

- return m, true

- }

- for _, ext := range exts {

- if e := find(s.FilterExt, ext); e == ext {

- return e, true

- }

- }

- return "", false

-}

-

-func (s *Storage) storeFile(file *os.File, hash, ext string) (id string, err error) {

- hfolder := s.idToFolder("files", hash)

- hpath := path.Join(hfolder, "file")

- fexists := false

-

- os.MkdirAll(path.Dir(hfolder), 0755)

- err = os.Mkdir(hfolder, 0755)

- if err != nil {

- if _, err = os.Stat(hpath); err != nil {

- err = errors.New("internal storage error")

- }

- fexists = true

- } else {

- err = os.Rename(file.Name(), hpath)

- os.Chmod(hpath, 0644)

- }

- if err != nil {

- return

- }

-

- fpath := ""

- for i := 0; i < MaxIdTries; i++ {

- id = s.randomId()

- dir := s.idToFolder("ids", id)

- os.MkdirAll(path.Dir(dir), 0755)

- err = os.Mkdir(dir, 0755)

- if err == nil {

- fpath = path.Join(dir, "file"+ext)

- id += ext

- break

- }

- }

- if fpath == "" {

- err = errors.New("internal storage error")

- return

- }

- rhpath, err := filepath.Rel(path.Dir(fpath), hpath)

- if err != nil {

- return

- }

- err = os.Symlink(rhpath, fpath)

-

- if fexists && err == nil {

- err = errFileExists

- }

- return

-}

-

-func find(ss []string, search string) string {

- for _, s := range ss {

- if s == search {

- return s

- }

- }

- return ""

-}