threadr.lostcave.ddnss.de/handlers/profile_edit.go

package handlers

import (
	"bytes"
	"crypto/sha256"
	"errors"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	"image/png"
	_ "image/png"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"threadr/models"

	"github.com/gorilla/sessions"
)

const maxProfileImageBytes = 2 << 20

func ProfileEditHandler(app *App) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		session := r.Context().Value("session").(*sessions.Session)
		userID, ok := session.Values["user_id"].(int)
		if !ok {
			http.Redirect(w, r, app.Config.ThreadrDir+"/login/", http.StatusFound)
			return
		}

		if r.Method == http.MethodPost {
			if !app.validateCSRFToken(r, session) {
				http.Error(w, "Invalid CSRF token", http.StatusForbidden)
				return
			}

			r.Body = http.MaxBytesReader(w, r.Body, maxProfileImageBytes+(256<<10))

			// Handle file upload
			file, handler, err := r.FormFile("pfp")
			if err == nil {
				defer file.Close()

				fileHash, fileID, err := saveProfileImageUpload(app, file)
				if err != nil {
					if errors.Is(err, errInvalidProfileImage) || errors.Is(err, errProfileImageTooLarge) {
						http.Error(w, err.Error(), http.StatusBadRequest)
						return
					}
					log.Printf("Error saving profile image: %v", err)
					http.Error(w, "Failed to save file", http.StatusInternalServerError)
					return
				}

				// Create file record in the database
				fileRecord := models.File{
					OriginalName:  sanitizeOriginalFileName(handler.Filename),
					Hash:          fileHash,
					HashAlgorithm: "sha256",
				}
				createdFileID, err := models.CreateFile(app.DB, fileRecord)
				if err != nil {
					log.Printf("Error creating file record: %v", err)
					http.Error(w, "Failed to save file information", http.StatusInternalServerError)
					_ = os.Remove(fileID)
					return
				}

				finalPath := filepath.Join(app.Config.FileStorageDir, models.ProfileImageStorageName(createdFileID))
				if err := os.Rename(fileID, finalPath); err != nil {
					_ = os.Remove(fileID)
					_ = models.DeleteFileByID(app.DB, createdFileID)
					log.Printf("Error moving file on disk: %v", err)
					http.Error(w, "Failed to save file", http.StatusInternalServerError)
					return
				}

				// Update user's pfp_file_id
				err = models.UpdateUserPfp(app.DB, userID, createdFileID)
				if err != nil {
					_ = os.Remove(finalPath)
					_ = models.DeleteFileByID(app.DB, createdFileID)
					log.Printf("Error updating user pfp: %v", err)
					http.Error(w, "Failed to update profile", http.StatusInternalServerError)
					return
				}
			} else if err != nil && !errors.Is(err, http.ErrMissingFile) {
				var maxBytesErr *http.MaxBytesError
				if errors.As(err, &maxBytesErr) || strings.Contains(err.Error(), "request body too large") {
					http.Error(w, errProfileImageTooLarge.Error(), http.StatusBadRequest)
					return
				}
				log.Printf("Error reading upload: %v", err)
				http.Error(w, "Failed to process file upload", http.StatusBadRequest)
				return
			}

			// Update other profile fields
			displayName := r.FormValue("display_name")
			bio := r.FormValue("bio")
			err = models.UpdateUserProfile(app.DB, userID, displayName, bio)
			if err != nil {
				log.Printf("Error updating profile: %v", err)
				http.Error(w, "Failed to update profile", http.StatusInternalServerError)
				return
			}

			http.Redirect(w, r, app.Config.ThreadrDir+"/profile/", http.StatusFound)
			return
		}

		user, err := models.GetUserByID(app.DB, userID)
		if err != nil {
			log.Printf("Error fetching user: %v", err)
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}
		if user == nil {
			http.Error(w, "User not found", http.StatusNotFound)
			return
		}

		data := struct {
			PageData
			User models.User
		}{
			PageData: PageData{
				Title:            "ThreadR - Edit Profile",
				Navbar:           "profile",
				LoggedIn:         true,
				ShowCookieBanner: false,
				BasePath:         app.Config.ThreadrDir,
				StaticPath:       app.Config.ThreadrDir + "/static",
				CurrentURL:       r.URL.RequestURI(),
				CSRFToken:        app.csrfToken(session),
			},
			User: *user,
		}
		if err := app.Tmpl.ExecuteTemplate(w, "profile_edit", data); err != nil {
			log.Printf("Error executing template in ProfileEditHandler: %v", err)
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}
	}
}

var (
	errInvalidProfileImage  = errors.New("Profile picture must be a PNG, JPEG, or GIF image")
	errProfileImageTooLarge = errors.New("Profile picture must be 2 MB or smaller")
)

func saveProfileImageUpload(app *App, file multipart.File) (string, string, error) {
	limitedReader := io.LimitReader(file, maxProfileImageBytes+1)
	data, err := io.ReadAll(limitedReader)
	if err != nil {
		return "", "", err
	}
	if int64(len(data)) > maxProfileImageBytes {
		return "", "", errProfileImageTooLarge
	}

	contentType := http.DetectContentType(data)
	if !isAllowedProfileImageType(contentType) {
		return "", "", errInvalidProfileImage
	}

	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return "", "", errInvalidProfileImage
	}

	tmpFile, err := os.CreateTemp(app.Config.FileStorageDir, "pfp-*.png")
	if err != nil {
		return "", "", err
	}
	defer func() {
		_ = tmpFile.Close()
	}()

	hash := sha256.Sum256(data)
	if err := png.Encode(tmpFile, img); err != nil {
		_ = os.Remove(tmpFile.Name())
		return "", "", err
	}

	if err := tmpFile.Close(); err != nil {
		_ = os.Remove(tmpFile.Name())
		return "", "", err
	}

	return fmt.Sprintf("%x", hash[:]), tmpFile.Name(), nil
}

func sanitizeOriginalFileName(name string) string {
	base := filepath.Base(strings.TrimSpace(name))
	if base == "." || base == string(filepath.Separator) || base == "" {
		return "profile.png"
	}
	return base
}

func isAllowedProfileImageType(contentType string) bool {
	switch contentType {
	case "image/png", "image/jpeg", "image/gif":
		return true
	default:
		return false
	}
}