diff options
| -rw-r--r-- | config.json | 5 | ||||
| -rw-r--r-- | main.go | 20 |
2 files changed, 21 insertions, 4 deletions
diff --git a/config.json b/config.json index 60269dd..0bf434a 100644 --- a/config.json +++ b/config.json @@ -2,5 +2,8 @@ "listen": "localhost:8080", "db": "user=julian dbname=noteman sslmode=disable", "tmpl": "", - "data": "data" + + "data": "data", + "max_upload": 10, + "max_pages": 5 } @@ -25,7 +25,10 @@ type Config struct { Listen string `json:"listen"` DBStr string `json:"db"` TmplPath string `json:"tmpl"` + DataPath string `json:"data"` + MaxSize uint `json:"max_upload"` + MaxPages uint `json:"max_pages"` } type Server struct { @@ -228,6 +231,8 @@ func (s *Server) httpUpload(w http.ResponseWriter, r *http.Request) { return } + r.Body = http.MaxBytesReader(w, r.Body, int64(s.conf.MaxSize) * 1000000) + // Head hash hash := r.FormValue("h") _, err := strconv.ParseUint(hash, 16, 16) @@ -253,7 +258,12 @@ func (s *Server) httpUpload(w http.ResponseWriter, r *http.Request) { } // Create file - fname := s.allocFile(hash, mtype) + fname, err := s.allocFile(hash, mtype) + if err != nil { + s.Error(w, r, err.Error(), http.StatusBadRequest) + return + } + f, err := os.Create(path.Join(s.conf.DataPath, fname)) if err != nil { s.Error(w, r, err.Error(), http.StatusInternalServerError) @@ -273,18 +283,22 @@ func (s *Server) httpUpload(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, "/", http.StatusSeeOther) } -func (s *Server) allocFile(hash string, t string) string { +func (s *Server) allocFile(hash string, t string) (string, error) { s.fslock.RLock() existing := s.filestore[hash] s.fslock.RUnlock() + if len(existing) >= int(s.conf.MaxPages) { + return "", fmt.Errorf("No more than %d pages is allowed", s.conf.MaxPages) + } + var ext string extarr, _ := mime.ExtensionsByType(t) if len(extarr) > 0 { ext = extarr[0] } - return fmt.Sprintf("%s.%d%s", hash, len(existing), ext) + return fmt.Sprintf("%s.%d%s", hash, len(existing), ext), nil } func (s *Server) saveFile(hash string, fname string) { |