diff --git a/api/token.go b/api/token.go index a8ac8d7..266f430 100644 --- a/api/token.go +++ b/api/token.go @@ -7,6 +7,8 @@ import ( "net/http" "path/filepath" + "os" + "github.com/gin-gonic/gin" "github.com/gotify/location" "github.com/gotify/server/auth" @@ -80,6 +82,9 @@ func (a *TokenAPI) DeleteApplication(ctx *gin.Context) { withID(ctx, "id", func(id uint) { if app := a.DB.GetApplicationByID(id); app != nil && app.UserID == auth.GetUserID(ctx) { a.DB.DeleteApplicationByID(id) + if app.Image != "" { + os.Remove(a.ImageDir + app.Image) + } } else { ctx.AbortWithError(404, fmt.Errorf("app with id %d doesn't exists", id)) } @@ -117,14 +122,23 @@ func (a *TokenAPI) UploadApplicationImage(ctx *gin.Context) { return } - name := auth.GenerateImageName() ext := filepath.Ext(file.Filename) + + name := auth.GenerateImageName() + for exist(a.ImageDir + name + ext) { + name = auth.GenerateImageName() + } + err = ctx.SaveUploadedFile(file, a.ImageDir+name+ext) if err != nil { ctx.AbortWithError(500, err) return } + if app.Image != "" { + os.Remove(a.ImageDir + app.Image) + } + app.Image = name + ext a.DB.UpdateApplication(app) ctx.JSON(200, withAbsoluteURL(ctx, app)) @@ -134,6 +148,13 @@ func (a *TokenAPI) UploadApplicationImage(ctx *gin.Context) { }) } +func exist(path string) bool { + if _, err := os.Stat(path); os.IsNotExist(err) { + return false + } + return true +} + func withAbsoluteURL(ctx *gin.Context, app *model.Application) *model.Application { url := location.Get(ctx) diff --git a/api/token_test.go b/api/token_test.go index 243017a..77cdef3 100644 --- a/api/token_test.go +++ b/api/token_test.go @@ -15,6 +15,8 @@ import ( "os" "reflect" + "io/ioutil" + "github.com/bouk/monkey" "github.com/gin-gonic/gin" "github.com/gotify/server/mode" @@ -359,7 +361,56 @@ func (s *TokenSuite) Test_UploadAppImage_WithImageFile_expectSuccess() { assert.Equal(s.T(), 200, s.recorder.Code) _, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png") assert.Nil(s.T(), err) - assert.Nil(s.T(), os.Remove("PorrUa5b1IIK3yKo_Pp6ww_9v.png")) + + s.a.DeleteApplication(s.ctx) + + _, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png") + assert.True(s.T(), os.IsNotExist(err)) +} + +func (s *TokenSuite) Test_UploadAppImage_WithImageFile_DeleteExstingImageAndGenerateNewName() { + s.db.User(5) + s.db.CreateApplication(&model.Application{UserID: 5, ID: 1, Image: "PorrUa5b1IIK3yKo_Pp6ww_9v.png"}) + + cType, buffer, err := upload(map[string]*os.File{"file": mustOpen("../test/assets/image.png")}) + assert.Nil(s.T(), err) + s.ctx.Request = httptest.NewRequest("POST", "/irrelevant", &buffer) + s.ctx.Request.Header.Set("Content-Type", cType) + test.WithUser(s.ctx, 5) + s.ctx.Params = gin.Params{{Key: "id", Value: "1"}} + fakeImage(s.T(), "PorrUa5b1IIK3yKo_Pp6ww_9v.png") + + s.a.UploadApplicationImage(s.ctx) + + assert.Equal(s.T(), 200, s.recorder.Code) + + _, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png") + assert.True(s.T(), os.IsNotExist(err)) + _, err = os.Stat("Zal6-ySIuL-T3EMLCcFtityHn.png") + assert.Nil(s.T(), err) + assert.Nil(s.T(), os.Remove("Zal6-ySIuL-T3EMLCcFtityHn.png")) +} + +func (s *TokenSuite) Test_UploadAppImage_WithImageFile_DeleteExistingImage() { + s.db.User(5) + s.db.CreateApplication(&model.Application{UserID: 5, ID: 1, Image: "existing.png"}) + + fakeImage(s.T(), "existing.png") + cType, buffer, err := upload(map[string]*os.File{"file": mustOpen("../test/assets/image.png")}) + assert.Nil(s.T(), err) + s.ctx.Request = httptest.NewRequest("POST", "/irrelevant", &buffer) + s.ctx.Request.Header.Set("Content-Type", cType) + test.WithUser(s.ctx, 5) + s.ctx.Params = gin.Params{{Key: "id", Value: "1"}} + + s.a.UploadApplicationImage(s.ctx) + + assert.Equal(s.T(), 200, s.recorder.Code) + + _, err = os.Stat("existing.png") + assert.True(s.T(), os.IsNotExist(err)) + + os.Remove("PorrUa5b1IIK3yKo_Pp6ww_9v.png") } func (s *TokenSuite) Test_UploadAppImage_WithTextFile_expectBadRequest() { @@ -444,3 +495,11 @@ func mustOpen(f string) *os.File { } return r } + +func fakeImage(t *testing.T, path string) { + data, err := ioutil.ReadFile("../test/assets/image.png") + assert.Nil(t, err) + // Write data to dst + err = ioutil.WriteFile(path, data, 0644) + assert.Nil(t, err) +}