make image module way more robust

- Create a better (more) re-usable flow for downloading images to disk and inserting images into the db
- Validate images properly by reading them fully and using DecodeConfig
- Fix some security issues by re-encoding images
- Always outputs one format `jpeg`, which keeps everything "normalized" and adds compression to images (mainly user uploads which might be big).

server/database/db.go

@@ -16,6 +16,10 @@ import (
 // Also runs migrations, etc, before returning connection.
 // Any error returned from this func should always make our app Exit (caller
 // handled).
+//
+// NOTE: Our mock db used in tests mimics this function, so if this func is
+// changed, you should look at the mock db to make it match if it makes
+// sense, so our tests stay accurate to prod.
 func New() (*gorm.DB, error) {
 	slog.Info("New: Opening new database connection")
 	// Open the database.
@@ -27,13 +31,27 @@ func New() (*gorm.DB, error) {
 		slog.Error("New: Opening database failed.")
 		return nil, err
 	}
-	if err := configure(db); err != nil {
-		slog.Error("New: Configuring connection failed!", "error", err)
+	// Setup the db (migrations, etc)
+	if err := Setup(db); err != nil {
+		slog.Error("New: Setting up connection failed!", "error", err)
 		return nil, err
 	}
+	return db, nil
+}
+
+// Setup configures our db connection and applies migrations.
+//
+// NOTE: This exists as a separate function so it can be reused by our testutil
+// package that we want to have configured in the same way as the main db so
+// that tests reflect real life.
+func Setup(db *gorm.DB) error {
+	if err := configure(db); err != nil {
+		slog.Error("Setup: Configuring connection failed!", "error", err)
+		return err
+	}
 	// Perform auto migration.
-	slog.Info("New: AutoMigrating")
-	err = db.AutoMigrate(
+	slog.Info("Setup: AutoMigrating")
+	err := db.AutoMigrate(
 		&migrate.MigrationRecord{},
 		&entity.User{},
 		&entity.UserServices{},
@@ -50,21 +68,21 @@ func New() (*gorm.DB, error) {
 		&entity.Tag{},
 	)
 	if err != nil {
-		slog.Error("New: Auto migration failed.")
-		return nil, err
+		slog.Error("Setup: Auto migration failed.")
+		return err
 	}
-	slog.Info("New: AutoMigrated")
+	slog.Info("Setup: AutoMigrated")
 	// Perform our manual migrations.
 	if err := migrate.Now(db); err != nil {
-		slog.Error("New: Manual migrations failed.", "error", err)
-		return nil, err
+		slog.Error("Setup: Manual migrations failed.", "error", err)
+		return err
 	}
 	// Optimize database.
 	if err := optimize(db); err != nil {
-		slog.Error("New: Optimizing database failed.", "error", err)
-		return nil, err
+		slog.Error("Setup: Optimizing database failed.", "error", err)
+		return err
 	}
-	return db, nil
+	return nil
 }
 
 // Configure our SQLite database connection.

server/feature/game/games.go

@@ -35,7 +35,10 @@ func (s *Service) saveGame(c *entity.Game, onlyUpdate bool) error {
 		return errors.New("game missing id or title")
 	}
 	if c.CoverID != "" {
-		p, err := image.DownloadAndInsertImage(s.db, "https://images.igdb.com/igdb/image/upload/t_cover_big/"+c.CoverID+".png", "games")
+		p, err := image.DownloadAndInsertFromUrl(
+			s.db,
+			"https://images.igdb.com/igdb/image/upload/t_cover_big/"+c.CoverID+".png",
+			"games")
 		if err != nil {
 			slog.Error("saveGame: Failed to cache game cover.", "error", err)
 		} else {

server/feature/user/router.go

@@ -136,7 +136,9 @@ func (r *Router) UpdateAvatar(c *gin.Context) {
 	userId := c.MustGet("userId").(uint)
 	response, err := r.service.UploadUserAvatar(c, userId)
 	if err != nil {
-		c.JSON(http.StatusInternalServerError, router.ErrorResponse{Error: err.Error()})
+		c.JSON(
+			http.StatusInternalServerError,
+			router.ErrorResponse{Error: err.Error()})
 		return
 	}
 	c.JSON(http.StatusOK, response)

server/feature/user/user.go

@@ -1,16 +1,10 @@
 package user
 
 import (
-	"crypto/sha256"
-	"encoding/hex"
 	"errors"
-	"io"
 	"log/slog"
-	"path"
-	"path/filepath"
 
 	"github.com/gin-gonic/gin"
-	"github.com/sbondCo/Watcharr/config"
 	"github.com/sbondCo/Watcharr/database/entity"
 	"github.com/sbondCo/Watcharr/image"
 	"gorm.io/gorm"
@@ -133,58 +127,41 @@ func (s *Service) UserUpdateBio(userId uint, newBio string) error {
 	return nil
 }
 
-func (s *Service) UploadUserAvatar(c *gin.Context, userId uint) (entity.Image, error) {
+func (s *Service) UploadUserAvatar(
+	c *gin.Context,
+	userId uint,
+) (entity.Image, error) {
 	file, err := c.FormFile("avatar")
 	if err != nil {
 		slog.Error("failed to get file", "error", err)
 		return entity.Image{}, errors.New("no file found")
 	}
 
-	slog.Debug("an avatar is being uploaded", "name", file.Filename)
+	slog.Debug("UploadUserAvatar: An avatar is being uploaded",
+		"name", file.Filename)
 
 	f, _ := file.Open()
-	if err := image.IsValidImageType(f); err != nil {
-		return entity.Image{}, errors.New("invalid image type")
-	}
-	h := sha256.New()
-	if _, err := io.Copy(h, f); err != nil {
-		slog.Error("uploadUserAvatar: Copy failed!", "error", err)
-		return entity.Image{}, errors.New("copy failed")
-	}
-	hs := hex.EncodeToString(h.Sum(nil))
+	defer f.Close()
 
-	slog.Debug("image hash calculated", "hash", hs, "first_letter", hs[0:1])
-
-	// Upload the file to specific dst.
-	outp := path.Join("img/up/", hs[0:1], hs+filepath.Ext(file.Filename))
-	c.SaveUploadedFile(file, path.Join(config.DataPath, outp))
-
-	_, err = f.Seek(0, 0)
+	img, err := image.DownloadAndInsert(s.db, f, "up")
 	if err != nil {
-		slog.Error("uploadUserAvatar seeking back to start of image failed", "error", err)
+		slog.Error("UploadUserAvatar: DownloadAndInsert failed!",
+			"error", err)
+		return entity.Image{}, errors.New("processing image failed")
 	}
 
-	// No need to remove old image, the daily cleanup task will handle removing unused ones.
-	var img entity.Image
-	err = s.db.Transaction(func(tx *gorm.DB) error {
-		// Insert avatar into db
-		img, err = image.InsertImage(s.db, hs, outp, f)
-		if err != nil {
-			return err
-		}
-		if img.ID == 0 {
-			return errors.New("image has no id")
-		}
-		// Update users avatar to newly inserted
-		if err := tx.Where("id = ?", userId).Updates(&entity.User{AvatarID: img.ID}).Error; err != nil {
-			return err
-		}
-		// commit transaction if no errors
-		return nil
-	})
-	if err != nil {
-		slog.Error("uploadUserAvatar failed!", "error", err)
-		return entity.Image{}, errors.New("uploadUserAvatar transaction failed")
+	// No need to remove old image, the daily cleanup task will handle removing
+	// unused ones.
+
+	// Update users avatar to newly inserted
+	res := s.db.
+		Where("id = ?", userId).
+		Updates(&entity.User{AvatarID: img.ID})
+	if res.Error != nil {
+		slog.Error("UploadUserAvatar: Updating the users avatar in db failed!",
+			"error", err)
+		return entity.Image{}, errors.New("updating user failed")
 	}
+
 	return img, nil
 }

server/image/image.go

@@ -7,28 +7,40 @@ import (
 	"errors"
 	"fmt"
 	"image"
-	_ "image/jpeg"
-	_ "image/png"
+	"image/gif"
+	"image/jpeg"
+	"image/png"
 	"io"
 	"log/slog"
-	"mime/multipart"
 	"net/http"
 	"os"
 	"path"
-	"path/filepath"
 
 	"github.com/buckket/go-blurhash"
 	"github.com/sbondCo/Watcharr/config"
 	"github.com/sbondCo/Watcharr/database/entity"
+	"github.com/sbondCo/Watcharr/util"
 	"gorm.io/gorm"
 )
 
+const (
+	defMaxSize        int64 = 10 * util.Mebibyte
+	defMaxWidthHeight int64 = 7680
+	defMaxPixels      int64 = 10_000_000
+)
+
 // TODO now that this file is in the image package it no longer needs to have "image(s)"
 // in the name of all the functions..
 
 // Insert an image into database
-func InsertImage(db *gorm.DB, hash string, path string, f io.Reader) (entity.Image, error) {
-	bh, _ := GetBlurHash(f)
+func Insert(
+	db *gorm.DB,
+	hash string,
+	path string,
+	b []byte,
+) (entity.Image, error) {
+	br := bytes.NewReader(b)
+	bh, _ := GetBlurHash(br)
 	img := entity.Image{
 		Hash:     hash,
 		Path:     path,
@@ -96,25 +108,144 @@ WHERE NOT EXISTS (
 	}
 }
 
-func IsValidImageType(f multipart.File) error {
-	// Read first 512 bytes, since that is all `DetectContentType` will evaluate on.
-	// Reading whole file is a waste.
-	buff := make([]byte, 512)
-	if _, err := f.Read(buff); err != nil {
-		slog.Error("isValidImageType: failed to read file into buffer", "error", err)
-		return errors.New("failed to verify if image is valid")
+// Validate the image safely.
+// Re-encodes the image in our own desired format, which helps verify the file
+// is a valid image and any undesirable data (eg think xss; appended html,
+// exit, etc) is not kept in the final image file we store.
+// We currently prefer jpg for the format we re-encode to, which also helps us
+// save storage space (adds compression, which the image we are validating
+// could lack or have a higher quality setting, etc).
+// Returns the new re-encoded image data, file extension, and error.
+// NOTE: Never use a user-set file extension (always validate we allow it).
+func Validate(
+	b []byte,
+) ([]byte, string, error) {
+	br := bytes.NewReader(b)
+
+	// Check image header for config/format.
+	cfg, format, err := image.DecodeConfig(br)
+	if err != nil {
+		slog.Error("Validate: Failed to DecodeConfig", "error", err)
+		return []byte{}, "", errors.New("invalid or bad image")
 	}
-	t := http.DetectContentType(buff)
-	slog.Debug("isValidImageType", "type", t)
-	if t != "image/png" && t != "image/jpeg" && t != "image/webp" && t != "image/gif" {
-		slog.Debug("isValidImageType: rejecting file as not valid (supported) image type")
-		return errors.New("invalid file type")
+	slog.Debug("Validate", "cfg", cfg, "format", format)
+	if int64(cfg.Width) > defMaxWidthHeight ||
+		int64(cfg.Height) > defMaxWidthHeight {
+		return []byte{}, "", errors.New("dimensions too large")
 	}
-	return nil
+	// Protect against images that max out the allowed width/height.
+	// If we assume each pixel is 4 bytes, someone maxing out 8000x8000 would
+	// mean ~235mb of data we need to decode into memory (i think!), but instead
+	// of limiting the max width/height values too much, we can limit the max
+	// amount of pixels to restrict the max size of the img pixels we'd allow.
+	// Then weird aspect ratios are still allowed.
+	// I'm definitely over-engineering this feature for a self-hosted movie list app lol.
+	if int64(cfg.Width)*int64(cfg.Height) > defMaxPixels {
+		return []byte{}, "", errors.New("i can't handle all those pixels")
+	}
+
+	// Seek back, we are reading again below for decode.
+	if _, err = br.Seek(0, 0); err != nil {
+		slog.Error("Validate: Seeking reader to start failed", "error", err)
+		return []byte{}, "", err
+	}
+
+	// Decode image.
+	// This should catch any malformed image files.
+	var img image.Image
+	switch format {
+	case "png":
+		img, err = png.Decode(br)
+	case "jpeg":
+		img, err = jpeg.Decode(br)
+	case "gif":
+		img, err = gif.Decode(br)
+	default:
+		return []byte{}, "", errors.New("unsupported image type")
+	}
+	if err != nil {
+		slog.Error("full image decode failed", "error", err, "format", format)
+		return []byte{}, "", errors.New("invalid or corrupt image")
+	}
+
+	// Re-encode the image from our decoded data (any extra included, possibly
+	// malicious data not part of the image should be gone now).
+	// exif data, etc should also be gone now too which is good.
+	// We just re-encode as jpeg right now, but if we wanted to, in the future
+	// it's possible to encode different formats based on if we want to perserve
+	// transparency from png, etc. OR maybe using webp will be easier and we
+	// can just use that format since it supports transparency and animations.
+	var buf bytes.Buffer
+	if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: 75}); err != nil {
+		slog.Error("Validate: Failed to encode jpeg", "error", err)
+		return []byte{}, "", errors.New("failed to encode image")
+	}
+	return buf.Bytes(), ".jpg", nil
 }
 
-func DownloadAndInsertImage(db *gorm.DB, url string, imgSubPath string) (entity.Image, error) {
-	slog.Debug("Attempting to download image", "url", url)
+// Creates the image file on disk.
+func save(b []byte, subPath string) (string, string, error) {
+	br := bytes.NewReader(b)
+	// First we get a hash of the files contents.
+	// Using the hash for filename has the benefit of us not storing duplicate
+	// files just because their filename provided to us is different.
+	h := sha256.New()
+	if _, err := io.Copy(h, br); err != nil {
+		slog.Error("save: Copy failed!", "error", err)
+		return "", "", errors.New("copy failed")
+	}
+	hs := hex.EncodeToString(h.Sum(nil))
+	slog.Debug("save: image hash calculated",
+		"hash", hs,
+		"first_letter", hs[0:1])
+
+	// Validate the file.
+	// We always validate, even from "trusted sources!"
+	b, ext, err := Validate(b)
+	if err != nil {
+		slog.Error("save: Validate failed!", "error", err)
+		return "", "", err
+	}
+	br = bytes.NewReader(b)
+
+	// Create paths for file.
+	imgPath := path.Join(
+		// Always outputs to `img/` dir.
+		"img/",
+		// Any sub path for separating images.
+		subPath,
+		// Sub-separate images by the starting character of their hash.
+		hs[0:1],
+		// File name is whole hash then the file extension.
+		hs+ext)
+	fullOutPath := path.Join(config.DataPath, imgPath)
+	slog.Debug("save: Built path", "path", imgPath)
+
+	// Save file
+	err = os.MkdirAll(path.Dir(fullOutPath), 0764)
+	if err != nil {
+		return "", "", err
+	}
+	out, err := os.Create(fullOutPath)
+	if err != nil {
+		return "", "", err
+	}
+	defer out.Close()
+	_, err = io.Copy(out, br)
+	if err != nil {
+		return "", "", err
+	}
+
+	return hs, imgPath, nil
+}
+
+// Download an image from `url` and insert it.
+func DownloadAndInsertFromUrl(
+	db *gorm.DB,
+	url string,
+	imgSubPath string,
+) (entity.Image, error) {
+	slog.Debug("DownloadAndInsertFromUrl: Running.", "url", url)
 
 	// Get the data
 	resp, err := http.Get(url)
@@ -128,66 +259,36 @@ func DownloadAndInsertImage(db *gorm.DB, url string, imgSubPath string) (entity.
 		return entity.Image{}, fmt.Errorf("bad status: %s", resp.Status)
 	}
 
-	// Read body into byte array, then create new reader
-	// So we have the ability to seek.
-	b, err := io.ReadAll(resp.Body)
-	if err != nil {
-		slog.Error("downloadAndInsertImage failed to read response into byte array", "error", err)
-		return entity.Image{}, err
-	}
-	br := bytes.NewReader(b)
-
-	h := sha256.New()
-	if _, err := io.Copy(h, br); err != nil {
-		slog.Error("DownloadAndInsertImage: Copy failed!", "error", err)
-		return entity.Image{}, errors.New("copy failed")
-	}
-	hs := hex.EncodeToString(h.Sum(nil))
-
-	// Seek back for file
-	_, err = br.Seek(0, 0)
-	if err != nil {
-		slog.Error("downloadAndInsertImage seeking back to start of br failed", "error", err)
-		return entity.Image{}, err
-	}
-
-	outp := path.Join("img/", imgSubPath, hs[0:1], hs+filepath.Ext(resp.Request.URL.Path))
-	dataOutP := path.Join(config.DataPath, outp)
-
-	// Create the file
-	out, err := os.Create(dataOutP)
-	if err != nil {
-		if os.IsNotExist(err) {
-			err = os.MkdirAll(path.Dir(dataOutP), 0764)
-			if err != nil {
-				return entity.Image{}, err
-			}
-			// If dirs made, try making file again
-			out, err = os.Create(dataOutP)
-			if err != nil {
-				return entity.Image{}, err
-			}
-		} else {
-			return entity.Image{}, err
-		}
-	}
-	defer out.Close()
-	_, err = io.Copy(out, br)
-	if err != nil {
-		return entity.Image{}, err
-	}
-
-	// Seek back for insertImage
-	_, err = br.Seek(0, 0)
-	if err != nil {
-		slog.Error("downloadAndInsertImage seeking back to start of br failed", "error", err)
-		return entity.Image{}, err
-	}
-
-	img, err := InsertImage(db, hs, outp, br)
-	if err != nil {
-		return entity.Image{}, err
-	}
-
-	return img, nil
+	return DownloadAndInsert(db, resp.Body, imgSubPath)
+}
+
+func DownloadAndInsert(
+	db *gorm.DB,
+	r io.Reader,
+	imgSubPath string,
+) (entity.Image, error) {
+	slog.Debug("DownloadAndInsert: Running.")
+
+	// Read all into memory.
+	b, err := util.LimitedReadAll(r, defMaxSize)
+	if err != nil {
+		slog.Error("DownloadAndInsert: Failed to read response!", "error", err)
+		return entity.Image{}, err
+	}
+
+	// Save the file.
+	imgHash, imgPath, err := save(b, imgSubPath)
+	if err != nil {
+		slog.Error("DownloadAndInsert: Failed to save file!", "error", err)
+		return entity.Image{}, err
+	}
+
+	// Insert image into db.
+	imge, err := Insert(db, imgHash, imgPath, b)
+	if err != nil {
+		slog.Error("DownloadAndInsert: Insert into db failed!", "error", err)
+		return entity.Image{}, err
+	}
+
+	return imge, nil
 }

server/image/image_test.go

@@ -0,0 +1,60 @@
+package image
+
+import (
+	"os"
+	"path"
+	"path/filepath"
+	"testing"
+
+	"github.com/sbondCo/Watcharr/config"
+	"github.com/sbondCo/Watcharr/database/entity"
+	"github.com/sbondCo/Watcharr/internal/testutil"
+)
+
+func TestDownloadAndInsertFromUrl(t *testing.T) {
+	testutil.SetupLogging()
+	db := testutil.SetupDB(t)
+
+	i, err := DownloadAndInsertFromUrl(
+		db,
+		"https://github.com/sbondCo/Watcharr/raw/dev/screenshot/homepage.png",
+		"test")
+	if err != nil {
+		t.Fatalf("DownloadAndInsert call failed: %v", err)
+	}
+
+	if i.ID == 0 || i.Path == "" || i.BlurHash == "" {
+		t.Fatal("returned entity.Image doesn't have certain fields!",
+			"id", i.ID, "path", i.Path, "blurhash", i.BlurHash)
+	}
+
+	fullImgDataPath := path.Join(config.DataPath, i.Path)
+
+	// Verify file exists and looks right.
+	if fi, err := os.Stat(fullImgDataPath); err != nil {
+		t.Fatalf("os.stat failed %v", err)
+	} else if fi.Size() <= 1 {
+		t.Fatalf("image file size doesn't seem right: %v", fi.Size())
+	} else if filepath.Ext(fi.Name()) != ".jpg" {
+		t.Fatalf("image file name doesn't have .jpg ext: %s", fi.Name())
+	}
+
+	if err := os.Remove(fullImgDataPath); err != nil {
+		// Not a fatal error because this is extra logic only for testing,
+		// but failing test because something in the real logic might possibly
+		// have something to do with it failing and we should probably know.
+		t.Errorf("removing the image file errored: %v", err)
+	}
+
+	// Verify image is in db
+	var c int64
+	if res := db.
+		Model(&entity.Image{}).
+		Where(&entity.Image{ID: i.ID}).
+		Count(&c); res.Error != nil {
+		t.Fatalf("verification query failed: %v", res.Error)
+	}
+	if c != 1 {
+		t.Fatalf("count of images in db doesn't look right: %v", c)
+	}
+}

server/internal/README.md

@@ -0,0 +1,3 @@
+We don't really use the `internal` folder because none of our code is anything we'd ever expect anyone to import and rely on in their own codebase.
+
+However, for some packages, it might make sense to put it here just as a signal to us (while developing) that this package is not code that goes into a prod build, etc (eg: `testutil`). This feels nicer, avoiding a scenario where our root folder has a bunch of real packages and ones that should never see prod mixed together (which is probably confusing).

server/internal/testutil/testutil.go

@@ -0,0 +1,55 @@
+// testutil is for testing code that we want to reuse for tests.
+package testutil
+
+import (
+	"log/slog"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/sbondCo/Watcharr/database"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+// SetupLogging will configure the default slog logger.
+// Since our main app uses `slog`, this is useful for getting
+// debug logs to show, or hiding all, etc.
+// Controlled by env var `WTEST_LOG_LEVEL` (accepts: `debug` or `error`), if not
+// set, default Info log level is used.
+func SetupLogging() {
+	level := slog.LevelInfo
+	switch os.Getenv("WTEST_LOG_LEVEL") {
+	case "debug":
+		level = slog.LevelDebug
+	case "error":
+		level = slog.LevelError
+	}
+	slog.SetDefault(slog.New(slog.NewTextHandler(
+		os.Stdout, &slog.HandlerOptions{Level: level})))
+}
+
+// Setup a fresh database for testing.
+// Exits test by using t.Fatalf if something fails.
+func SetupDB(t *testing.T) *gorm.DB {
+	t.Helper()
+
+	// Open our test db.
+	// Note: Could have used inmemory db, but it breaks our WAL migration and
+	// errors out, and I don't wanna mess with prod code simply so I can use
+	// an inmem db for testing, so we make a temporary file db.
+	db, err := gorm.Open(
+		sqlite.Open(filepath.Join(t.TempDir(), "test-watcharr.db")),
+		&gorm.Config{TranslateError: true},
+	)
+	if err != nil {
+		t.Fatalf("failed to open test db: %v", err)
+	}
+
+	// Setup the db same as we do for prod.
+	if err := database.Setup(db); err != nil {
+		t.Fatalf("failed to migrate test db: %v", err)
+	}
+
+	return db
+}

server/util/io.go

@@ -0,0 +1,18 @@
+package util
+
+import (
+	"errors"
+	"io"
+)
+
+func LimitedReadAll(r io.Reader, maxSize int64) ([]byte, error) {
+	// LimitReader: Allow reading one extra byte so our maxSize check later works.
+	b, err := io.ReadAll(io.LimitReader(r, maxSize+1))
+	if err != nil {
+		return b, err
+	}
+	if int64(len(b)) > maxSize {
+		return b, errors.New("file is too big")
+	}
+	return b, nil
+}

server/util/size.go

@@ -0,0 +1,7 @@
+package util
+
+const (
+	Byte     int64 = 1
+	Kibibyte int64 = 1024
+	Mebibyte int64 = 1024 * Kibibyte
+)