Added limits for uploadingupload

1 files changed, 17 insertions, 3 deletions

diff --git a/main.go b/main.go
index e6c1def..0ebbdc4 100644
--- a/main.go
+++ b/main.go
@@ -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) {