From 08611f394e3540834b83619d4154abaa39d62c92 Mon Sep 17 00:00:00 2001
From: Lukas Schulte Pelkum <kbrt@protonmail.com>
Date: Sun, 11 Jul 2021 17:57:45 +0200
Subject: [PATCH] Implement paste modification endpoint (#10)

---
 internal/storage/mongodb/mongodb_driver.go   |  5 ++-
 internal/storage/postgres/postgres_driver.go | 10 ++++-
 internal/web/controllers/v2/pastes.go        | 39 ++++++++++++++++++++
 3 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/internal/storage/mongodb/mongodb_driver.go b/internal/storage/mongodb/mongodb_driver.go
index 2831422..5ba4af8 100644
--- a/internal/storage/mongodb/mongodb_driver.go
+++ b/internal/storage/mongodb/mongodb_driver.go
@@ -117,8 +117,9 @@ func (driver *MongoDBDriver) Save(paste *shared.Paste) error {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
-	// Insert the paste object
-	_, err := collection.InsertOne(ctx, paste)
+	// Upsert the paste object
+	filter := bson.M{"_id": paste.ID}
+	_, err := collection.UpdateOne(ctx, filter, paste, options.Update().SetUpsert(true))
 	return err
 }
 
diff --git a/internal/storage/postgres/postgres_driver.go b/internal/storage/postgres/postgres_driver.go
index 4ef5925..859b1ad 100644
--- a/internal/storage/postgres/postgres_driver.go
+++ b/internal/storage/postgres/postgres_driver.go
@@ -93,7 +93,15 @@ func (driver *PostgresDriver) Get(id string) (*shared.Paste, error) {
 
 // Save saves a paste
 func (driver *PostgresDriver) Save(paste *shared.Paste) error {
-	query := "INSERT INTO pastes VALUES ($1, $2, $3, $4, $5)"
+	query := `
+		INSERT INTO pastes (id, content, modificationToken, created, autoDelete)
+		VALUES ($1, $2, $3, $4, $5)
+		ON CONFLICT (id) DO UPDATE
+			SET content = excluded.token,
+				modificationToken = excluded.modificationToken,
+				created = excluded.created,
+				autoDelete = excluded.autoDelete
+	`
 
 	_, err := driver.pool.Exec(context.Background(), query, paste.ID, paste.Content, paste.ModificationToken, paste.Created, paste.AutoDelete)
 	return err
diff --git a/internal/web/controllers/v2/pastes.go b/internal/web/controllers/v2/pastes.go
index 70724cc..096da51 100644
--- a/internal/web/controllers/v2/pastes.go
+++ b/internal/web/controllers/v2/pastes.go
@@ -19,6 +19,7 @@ func InitializePastesController(group *router.Group, rateLimiterMiddleware *limi
 	// moms spaghetti
 	group.GET("/{id}", rateLimiterMiddleware.Handle(middlewareInjectPaste(endpointGetPaste)))
 	group.POST("/", rateLimiterMiddleware.Handle(endpointCreatePaste))
+	group.PATCH("/{id}", rateLimiterMiddleware.Handle(middlewareInjectPaste(middlewareValidateModificationToken(endpointModifyPaste))))
 	group.DELETE("/{id}", rateLimiterMiddleware.Handle(middlewareInjectPaste(middlewareValidateModificationToken(endpointDeletePaste))))
 }
 
@@ -163,6 +164,44 @@ func endpointCreatePaste(ctx *fasthttp.RequestCtx) {
 	ctx.SetBody(jsonData)
 }
 
+type endpointModifyPastePayload struct {
+	Content *string `json:"content"`
+}
+
+// endpointModifyPaste handles the 'PATCH /v2/pastes/{id}' endpoint
+func endpointModifyPaste(ctx *fasthttp.RequestCtx) {
+	// Read, parse and validate the request payload
+	payload := new(endpointModifyPastePayload)
+	if err := json.Unmarshal(ctx.PostBody(), payload); err != nil {
+		ctx.SetStatusCode(fasthttp.StatusInternalServerError)
+		ctx.SetBodyString(err.Error())
+		return
+	}
+	if payload.Content != nil && *payload.Content == "" {
+		ctx.SetStatusCode(fasthttp.StatusBadRequest)
+		ctx.SetBodyString("missing paste content")
+		return
+	}
+	if payload.Content != nil && config.Current.LengthCap > 0 && len(*payload.Content) > config.Current.LengthCap {
+		ctx.SetStatusCode(fasthttp.StatusBadRequest)
+		ctx.SetBodyString("too large paste content")
+		return
+	}
+
+	// Modify the paste itself
+	paste := ctx.UserValue("_paste").(*shared.Paste)
+	if payload.Content != nil {
+		paste.Content = *payload.Content
+	}
+
+	// Save the modified paste
+	if err := storage.Current.Save(paste); err != nil {
+		ctx.SetStatusCode(fasthttp.StatusInternalServerError)
+		ctx.SetBodyString(err.Error())
+		return
+	}
+}
+
 // endpointDeletePaste handles the 'DELETE /v2/pastes/{id}' endpoint
 func endpointDeletePaste(ctx *fasthttp.RequestCtx) {
 	paste := ctx.UserValue("_paste").(*shared.Paste)