Reserve media IDs to avoid re-upload attacks

controllers/maintenance_controller/maintainance_controller.go

@@ -288,6 +288,12 @@ func doPurge(media *types.Media, ctx context.Context, log *logrus.Entry) error {
 		return err
 	}
 
+	metadataDb := storage.GetDatabase().GetMetadataStore(ctx, log)
+	err = metadataDb.ReserveMediaId(media.Origin, media.MediaId, "purged / deleted")
+	if err != nil {
+		return err
+	}
+
 	mediaDb := storage.GetDatabase().GetMediaStore(ctx, log)
 	err = mediaDb.Delete(media.Origin, media.MediaId)
 	if err != nil {

controllers/upload_controller/upload_controller.go

@@ -2,6 +2,7 @@ package upload_controller
 
 import (
 	"context"
+	"database/sql"
 	"io"
 	"io/ioutil"
 	"strconv"
@@ -86,9 +87,42 @@ func UploadMedia(contents io.ReadCloser, contentLength int64, contentType string
 		data = contents
 	}
 
-	mediaId, err := util.GenerateRandomString(64)
-	if err != nil {
-		return nil, err
+	metadataDb := storage.GetDatabase().GetMetadataStore(ctx, log)
+	mediaDb := storage.GetDatabase().GetMediaStore(ctx, log)
+
+	mediaTaken := true
+	var mediaId string
+	var err error
+	attempts := 0
+	for mediaTaken {
+		attempts += 1
+		if attempts > 10 {
+			return nil, errors.New("failed to generate a media ID after 10 rounds")
+		}
+
+		mediaId, err = util.GenerateRandomString(64)
+		if err != nil {
+			return nil, err
+		}
+
+		mediaTaken, err = metadataDb.IsReserved(origin, mediaId)
+		if err != nil {
+			return nil, err
+		}
+
+		if !mediaTaken {
+			// Double check it isn't already in use
+			var media *types.Media
+			media, err = mediaDb.Get(origin, mediaId)
+			if err == sql.ErrNoRows {
+				mediaTaken = false
+				continue
+			}
+			if err != nil {
+				return nil, err
+			}
+			mediaTaken = media != nil
+		}
 	}
 
 	return StoreDirect(data, contentLength, contentType, filename, userId, origin, mediaId, ctx, log)

migrations/11_add_reserved_ids_table_down.sql

@@ -0,0 +1,2 @@
+DROP INDEX reserved_media_index;
+DROP TABLE reserved_media;

migrations/11_add_reserved_ids_table_up.sql

@@ -0,0 +1,6 @@
+CREATE TABLE IF NOT EXISTS reserved_media (
+	origin TEXT NOT NULL,
+	media_id TEXT NOT NULL,
+	reason TEXT NOT NULL
+);
+CREATE UNIQUE INDEX IF NOT EXISTS reserved_media_index ON reserved_media (media_id, origin);

storage/stores/metadata_store.go

@@ -27,6 +27,8 @@ const insertNewBackgroundTask = "INSERT INTO background_tasks (task, params, sta
 const selectBackgroundTask = "SELECT id, task, params, start_ts, end_ts FROM background_tasks WHERE id = $1"
 const updateBackgroundTask = "UPDATE background_tasks SET end_ts = $2 WHERE id = $1"
 const selectAllBackgroundTasks = "SELECT id, task, params, start_ts, end_ts FROM background_tasks"
+const insertReservation = "INSERT INTO reserved_media (origin, media_id, reason) VALUES ($1, $2, $3);"
+const selectReservation = "SELECT origin, media_id, reason FROM reserved_media WHERE origin = $1 AND media_id = $2;"
 
 type metadataStoreStatements struct {
 	upsertLastAccessed                            *sql.Stmt
@@ -42,6 +44,8 @@ type metadataStoreStatements struct {
 	selectBackgroundTask                          *sql.Stmt
 	updateBackgroundTask                          *sql.Stmt
 	selectAllBackgroundTasks                      *sql.Stmt
+	insertReservation                             *sql.Stmt
+	selectReservation                             *sql.Stmt
 }
 
 type MetadataStoreFactory struct {
@@ -101,6 +105,12 @@ func InitMetadataStore(sqlDb *sql.DB) (*MetadataStoreFactory, error) {
 	if store.stmts.selectAllBackgroundTasks, err = store.sqlDb.Prepare(selectAllBackgroundTasks); err != nil {
 		return nil, err
 	}
+	if store.stmts.insertReservation, err = store.sqlDb.Prepare(insertReservation); err != nil {
+		return nil, err
+	}
+	if store.stmts.selectReservation, err = store.sqlDb.Prepare(selectReservation); err != nil {
+		return nil, err
+	}
 
 	return &store, nil
 }
@@ -314,3 +324,27 @@ func (s *MetadataStore) GetAllBackgroundTasks() ([]*types.BackgroundTask, error)
 
 	return results, nil
 }
+
+func (s *MetadataStore) ReserveMediaId(origin string, mediaId string, reason string) error {
+	_, err := s.statements.insertReservation.ExecContext(s.ctx, origin, mediaId, reason)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (s *MetadataStore) IsReserved(origin string, mediaId string) (bool, error) {
+	r := s.statements.selectReservation.QueryRowContext(s.ctx, origin, mediaId)
+	var dbOrigin string
+	var dbMediaId string
+	var dbReason string
+
+	err := r.Scan(&dbOrigin, &dbMediaId, &dbReason)
+	if err == sql.ErrNoRows {
+		return false, nil
+	}
+	if err != nil {
+		return true, err
+	}
+	return true, nil
+}