morffix/server/upload.go

package server

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"git.lastassault.de/speatzle/morffix/constants"
)

func handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.Header.Get(constants.SHARED_SECRET_HEADER) != conf.SharedSecret {
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	strid := r.PathValue("id")
	if strid == "" {
		http.Error(w, "No id", http.StatusBadRequest)
		return
	}

	id, err := strconv.Atoi(strid)
	if err != nil {
		errorUpload(r, w, 0, "Convert File ID to int", err)
		return
	}

	if r.Header.Get(constants.WORKER_VERSION_HEADER) != constants.WORKER_VERSION {
		http.Error(w, "Wrong Worker Version", http.StatusBadRequest)
		return
	}

	taskid, err := strconv.Atoi(r.Header.Get(constants.TASK_ID_HEADER))
	if err != nil {
		errorUpload(r, w, 0, "Convert Task ID to int", err)
		return
	}

	var fileid int
	var workerid string
	var status constants.TaskStatus
	err = db.QueryRow(r.Context(), "SELECT file_id, worker_id, status FROM tasks WHERE id = $1", taskid).Scan(&fileid, &workerid, &status)
	if err != nil {
		errorUpload(r, w, 0, "Getting Task for Upload", err)
		return
	}

	if workerid != r.Header.Get(constants.UUID_HEADER) {
		errorUpload(r, w, taskid, "Worker is not Assigned to Task", err)
		return
	}

	if fileid != id {
		errorUpload(r, w, taskid, "File is not for this Task", err)
		return
	}

	if !(status == constants.TASK_STATUS_RUNNING || status == constants.TASK_STATUS_UNKNOWN) {
		errorUpload(r, w, taskid, "File can only be uploaded when task status is running or unknown", err)
		return
	}

	hash, err := base64.StdEncoding.DecodeString(r.Header.Get(constants.HASH_HEADER))
	if err != nil {
		errorUpload(r, w, taskid, "Decode Hash", err)
		return
	}

	var fPath string
	var libraryID int
	err = db.QueryRow(r.Context(), "SELECT path, library_id FROM files WHERE id = $1", id).Scan(&fPath, &libraryID)
	if err != nil {
		errorUpload(r, w, taskid, "Query File Path", err)
		return
	}

	var lPath string
	var tPath string
	var tReplace bool
	err = db.QueryRow(r.Context(), "SELECT path, transcode_path, transcode_replace FROM libraries WHERE id = $1", libraryID).Scan(&lPath, &tPath, &tReplace)
	if err != nil {
		errorUpload(r, w, taskid, "Query Library Path", err)
		return
	}

	// When replacing, send to temp, otherwise write directly
	var path string
	if tReplace {
		path = filepath.Join(lPath, fPath+constants.TEMP_FILE_EXTENSION)
	} else {
		path = filepath.Join(tPath, fPath)
	}

	slog.InfoContext(r.Context(), "Starting multipart Parsing for file", "id", id, "path", path)

	err = r.ParseMultipartForm(100 << 20)
	if err != nil {
		errorUpload(r, w, taskid, "Parse Multipart Form", err)
		return
	}

	srcFile, _, err := r.FormFile(constants.FORM_FILE_KEY)
	if err != nil {
		errorUpload(r, w, taskid, "Form File", err)
		return
	}

	src := http.MaxBytesReader(w, srcFile, 1000000<<20)
	// MaxBytesReader closes the underlying io.Reader on its Close() is called
	defer src.Close()

	// if we are not replacing then we dont know if the destination folder exists
	if !tReplace {
		err = os.MkdirAll(filepath.Dir(path), 0775)
		if err != nil {
			errorUpload(r, w, taskid, "Creating Folder", err)
			return
		}
	}

	out, err := os.Create(path)
	if err != nil {
		errorUpload(r, w, taskid, "Creating File", err)
		return
	}
	defer out.Close()

	written, err := io.Copy(out, src)
	if err != nil {
		var maxBytesError *http.MaxBytesError
		if errors.As(err, &maxBytesError) {
			slog.ErrorContext(r.Context(), "File to Large", "err", err)
			http.Error(w, "File to Large: "+err.Error(), http.StatusRequestEntityTooLarge)
			return
		}
		errorUpload(r, w, taskid, "Failed to write the uploaded content", err)
		return
	}

	slog.InfoContext(r.Context(), "upload done", "written", written)

	out.Close()

	// if we are replacing, then we need to change the original file to the uploaded file and update the hash in the database
	if tReplace {
		tx, err := db.Begin(r.Context())
		if err != nil {
			errorUpload(r, w, taskid, "Begin Transaction", err)
			return
		}

		dstPath := filepath.Join(lPath, fPath)
		slog.InfoContext(r.Context(), "Replacing Original file With Uploaded File", "fileid", id, "path", path, "dstPath", dstPath)

		err = os.Rename(path, dstPath)
		if err != nil {
			errorUpload(r, w, taskid, "Replace File", err)
			return
		}

		info, err := os.Stat(dstPath)
		if err != nil {
			errorUpload(r, w, taskid, "Stat File", err)
			return
		}
		modTime := info.ModTime().UTC().Round(time.Second)
		size := info.Size()

		_, err = tx.Exec(r.Context(), "UPDATE files SET hash = $2, size = $3, mod_time = $3 WHERE id = $1", fileid, hash, size, modTime)
		if err != nil {
			errorUpload(r, w, taskid, "Set File Hash", err)
			return
		}

		err = tx.Commit(r.Context())
		if err != nil {
			errorUpload(r, w, taskid, "Commit File Hash", err)
			return
		}

		slog.InfoContext(r.Context(), "Original file Replaced with Uploaded File", "fileid", id, "dstPath", dstPath)
	}
}

func errorUpload(r *http.Request, w http.ResponseWriter, taskID int, msg string, err error) {
	slog.ErrorContext(r.Context(), msg, "err", err)
	http.Error(w, msg+": "+err.Error(), http.StatusInternalServerError)
	if taskID != 0 {
		_, err2 := db.Exec(context.TODO(), "UPDATE tasks SET log = log || $2 WHERE id = $1", taskID, []string{fmt.Sprintf("%v MASTER: upload error:"+msg+":"+err.Error(), time.Now())})
		if err != nil {
			slog.ErrorContext(r.Context(), "Updating task log with upload error", "err", err2)
		}
	}
}