Upload hooks

api/api.go

@@ -6,10 +6,12 @@ import (
 	"github.com/fasthttp/websocket"
 	"github.com/jinzhu/gorm"
 	_ "github.com/jinzhu/gorm/dialects/postgres"
+	"github.com/robfig/cron"
 	"github.com/sirupsen/logrus"
 	"github.com/valyala/fasthttp"
 	"os"
 	"path/filepath"
+	"time"
 )
 
 var WorkDir = getWorkDir()
@@ -28,6 +30,7 @@ type WebApi struct {
 	server      fasthttp.Server
 	db          *gorm.DB
 	MotdMessage *websocket.PreparedMessage
+	Cron        *cron.Cron
 }
 
 func Index(ctx *fasthttp.RequestCtx) {
@@ -67,6 +70,7 @@ func New(db *gorm.DB) *WebApi {
 	}
 
 	api := &WebApi{}
+	api.Cron = cron.New()
 
 	logrus.SetLevel(getLogLevel())
 
@@ -90,6 +94,7 @@ func New(db *gorm.DB) *WebApi {
 
 func (api *WebApi) Run() {
 	address := GetServerAddress()
+	api.setupCronJobs()
 
 	logrus.WithFields(logrus.Fields{
 		"addr": address,
@@ -100,7 +105,20 @@ func (api *WebApi) Run() {
 		logrus.Fatalf("Error in ListenAndServe: %s", err)
 	}
 }
+func (api *WebApi) setupCronJobs() {
+	duration, _ := time.ParseDuration("5m")
+	api.Cron = cron.New()
+	schedule := cron.Every(duration)
+	api.Cron.Schedule(schedule, cron.FuncJob(api.DisposeStaleUploadSlots))
 
+	api.Cron.Start()
+
+	logrus.WithFields(logrus.Fields{
+		"every": duration,
+	}).Info("Scheduled job for DisposeStaleUploadSlots")
+}
+
+//TODO: Move those to a different file/package
 func GetServerAddress() string {
 	serverAddress := os.Getenv("WS_BUCKET_ADDR")
 	if serverAddress == "" {

api/maintenance.go

@@ -0,0 +1,30 @@
+package api
+
+import (
+	"github.com/sirupsen/logrus"
+	"os"
+	"path/filepath"
+	"time"
+)
+
+func (api *WebApi) DisposeStaleUploadSlots() {
+
+	var toDispose []UploadSlot
+	api.db.Where("? >= to_dispose_date", time.Now().Unix()).Find(&toDispose)
+
+	for _, slot := range toDispose {
+		path := filepath.Join(WorkDir, slot.FileName)
+
+		err := os.Remove(path)
+		api.db.Where("token = ?", slot.Token).Delete(UploadSlot{})
+
+		logrus.WithFields(logrus.Fields{
+			"fileName": slot.FileName,
+			"err":      err,
+		}).Trace("Deleted file")
+	}
+
+	logrus.WithFields(logrus.Fields{
+		"staleUploadSlots": len(toDispose),
+	}).Info("Disposed stale upload slots")
+}

api/models.go

@@ -12,9 +12,11 @@ type GenericResponse struct {
 }
 
 type AllocateUploadSlotRequest struct {
-	Token    string `json:"token"`
-	MaxSize  int64  `json:"max_size"`
-	FileName string `json:"file_name"`
+	Token         string `json:"token"`
+	MaxSize       int64  `json:"max_size"`
+	FileName      string `json:"file_name"`
+	ToDisposeDate int64  `json:"to_dispose_date"`
+	UploadHook    string `json:"upload_hook"`
 }
 
 func (req *AllocateUploadSlotRequest) IsValid() bool {
@@ -45,9 +47,11 @@ func (req *AllocateUploadSlotRequest) IsValid() bool {
 }
 
 type UploadSlot struct {
-	MaxSize  int64  `json:"max_size"`
-	Token    string `gorm:"primary_key",json:"token"`
-	FileName string `json:"file_name"`
+	MaxSize       int64  `json:"max_size"`
+	Token         string `gorm:"primary_key",json:"token"`
+	FileName      string `json:"file_name"`
+	ToDisposeDate int64  `json:"to_dispose_date"`
+	UploadHook    string `json:"upload_hook"`
 }
 
 type WebsocketMotd struct {

api/slot.go

@@ -7,7 +7,9 @@ import (
 	"github.com/valyala/fasthttp"
 	"io"
 	"os"
+	"os/exec"
 	"path/filepath"
+	"strings"
 	"sync"
 )
 
@@ -133,12 +135,32 @@ func (api *WebApi) Upload(ctx *fasthttp.RequestCtx) {
 			logrus.WithError(err).Error("Error while closing file")
 		}
 		mu.(*sync.RWMutex).Unlock()
+		mu.(*sync.RWMutex).RLock()
+
+		executeUploadHook(slot)
+
+		mu.(*sync.RWMutex).RUnlock()
 	})
 	if err != nil {
 		logrus.WithError(err).Error("Error while upgrading connexion")
 	}
 }
 
+func executeUploadHook(slot UploadSlot) {
+
+	path := filepath.Join(WorkDir, slot.FileName)
+
+	commandStr := strings.Replace(slot.UploadHook, "$1", "\""+path+"\"", -1)
+	cmd := exec.Command("bash", "-c", commandStr)
+	output, err := cmd.CombinedOutput()
+
+	logrus.WithFields(logrus.Fields{
+		"output":     string(output),
+		"err":        err,
+		"commandStr": commandStr,
+	}).Info("Execute upload hook")
+}
+
 func (api *WebApi) ReadUploadSlot(ctx *fasthttp.RequestCtx) {
 
 	tokenStr := string(ctx.Request.Header.Peek("X-Upload-Token"))
@@ -191,9 +213,11 @@ func (api *WebApi) ReadUploadSlot(ctx *fasthttp.RequestCtx) {
 func (api *WebApi) allocateUploadSlot(req *AllocateUploadSlotRequest) error {
 
 	slot := &UploadSlot{
-		MaxSize:  req.MaxSize,
-		FileName: req.FileName,
-		Token:    req.Token,
+		MaxSize:       req.MaxSize,
+		FileName:      req.FileName,
+		Token:         req.Token,
+		ToDisposeDate: req.ToDisposeDate,
+		UploadHook:    req.UploadHook,
 	}
 
 	logrus.WithFields(logrus.Fields{