hmn/src/assets/assets.go

package assets

import (
	"bytes"
	"context"
	"crypto/sha1"
	"errors"
	"fmt"
	"image"
	"io"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strings"

	"git.handmade.network/hmn/hmn/src/config"
	"git.handmade.network/hmn/hmn/src/db"
	"git.handmade.network/hmn/hmn/src/hmnurl"
	"git.handmade.network/hmn/hmn/src/jobs"
	"git.handmade.network/hmn/hmn/src/logging"
	"git.handmade.network/hmn/hmn/src/models"
	"git.handmade.network/hmn/hmn/src/oops"
	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/aws/smithy-go"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

var client *s3.Client

func init() {
	cfg, err := awsconfig.LoadDefaultConfig(context.Background(),
		awsconfig.WithCredentialsProvider(
			credentials.NewStaticCredentialsProvider(
				config.Config.DigitalOcean.AssetsSpacesKey,
				config.Config.DigitalOcean.AssetsSpacesSecret,
				"",
			),
		),
		awsconfig.WithRegion(config.Config.DigitalOcean.AssetsSpacesRegion),
		awsconfig.WithEndpointResolver(aws.EndpointResolverFunc(func(service, region string) (aws.Endpoint, error) {
			return aws.Endpoint{
				URL: config.Config.DigitalOcean.AssetsSpacesEndpoint,
			}, nil
		})),
	)
	if err != nil {
		panic(err)
	}
	client = s3.NewFromConfig(cfg, func(o *s3.Options) {
		o.UsePathStyle = true
	})
}

type CreateInput struct {
	Content     []byte
	Filename    string
	ContentType string

	// Optional params
	UploaderID    *int // HMN user id
	Width, Height int
}

var REIllegalFilenameChars = regexp.MustCompile(`[^\w\-.]`)

func SanitizeFilename(filename string) string {
	if filename == "" {
		return "unnamed"
	}
	return REIllegalFilenameChars.ReplaceAllString(filename, "_")
}

func AssetKey(id, filename string) string {
	return fmt.Sprintf("%s/%s", id, filename)
}

type InvalidAssetError error

func Create(ctx context.Context, dbConn db.ConnOrTx, in CreateInput) (*models.Asset, error) {
	filename := SanitizeFilename(in.Filename)

	if len(in.Content) == 0 {
		return nil, InvalidAssetError(fmt.Errorf("could not upload asset '%s': no bytes of data were provided", filename))
	}
	if in.ContentType == "" {
		return nil, InvalidAssetError(fmt.Errorf("could not upload asset '%s': no content type provided", filename))
	}

	// Upload the asset to the DO space
	id := uuid.New()
	key := AssetKey(id.String(), filename)
	checksum := fmt.Sprintf("%x", sha1.Sum(in.Content))

	upload := func() error {
		_, err := client.PutObject(ctx, &s3.PutObjectInput{
			Bucket:      &config.Config.DigitalOcean.AssetsSpacesBucket,
			Key:         &key,
			Body:        bytes.NewReader(in.Content),
			ACL:         types.ObjectCannedACLPublicRead,
			ContentType: &in.ContentType,
		})
		return err
	}

	err := upload()
	if err != nil {
		var apiError smithy.APIError
		if errors.As(err, &apiError) && apiError.ErrorCode() == "NoSuchBucket" {
			_, err := client.CreateBucket(ctx, &s3.CreateBucketInput{
				Bucket: &config.Config.DigitalOcean.AssetsSpacesBucket,
			})
			if err != nil {
				return nil, oops.New(err, "failed to create assets bucket")
			}

			err = upload()
			if err != nil {
				return nil, oops.New(err, "failed to upload asset")
			}
		} else {
			return nil, oops.New(err, "failed to upload asset")
		}
	}

	var thumbnailKey *string

	width := in.Width
	height := in.Height
	if previewBytes, thumbWidth, thumbHeight, err := ExtractPreview(ctx, in.ContentType, in.Content); err != nil {
		logging.Error().Err(err).Msg("Failed to generate preview for asset")
	} else if len(previewBytes) > 0 {
		keyStr := AssetKey(id.String(), fmt.Sprintf("%s_thumb.jpg", id.String()))
		thumbnailType := "image/jpeg"
		_, err = client.PutObject(ctx, &s3.PutObjectInput{
			Bucket:      &config.Config.DigitalOcean.AssetsSpacesBucket,
			Key:         &keyStr,
			Body:        bytes.NewReader(previewBytes),
			ACL:         types.ObjectCannedACLPublicRead,
			ContentType: &thumbnailType,
		})
		if err != nil {
			logging.Error().Err(err).Msg("Failed to upload thumbnail for video")
		} else {
			thumbnailKey = &keyStr
		}

		if width == 0 || height == 0 {
			width = thumbWidth
			height = thumbHeight
		}
	}

	// Save a record in our database
	// TODO(db): Would be convient to use RETURNING here...
	_, err = dbConn.Exec(ctx,
		`
		INSERT INTO asset (id, s3_key, thumbnail_s3_key, filename, size, mime_type, sha1sum, width, height, uploader_id)
		VALUES            ($1, $2,     $3,               $4,       $5,   $6,        $7,      $8,    $9,     $10)
		`,
		id,
		key,
		thumbnailKey,
		filename,
		len(in.Content),
		in.ContentType,
		checksum,
		width,
		height,
		in.UploaderID,
	)
	if err != nil {
		return nil, oops.New(err, "failed to save asset record")
	}

	// Fetch and return the new record
	asset, err := db.QueryOne[models.Asset](ctx, dbConn,
		`
		SELECT $columns
		FROM asset
		WHERE id = $1
		`,
		id,
	)
	if err != nil {
		return nil, oops.New(err, "failed to fetch newly-created asset")
	}

	return asset, nil
}

func getFFMpegPath() string {
	path := config.Config.PreviewGeneration.FFMpegPath
	if path != "" {
		return path
	}
	var err error
	path, err = exec.LookPath("ffmpeg")
	if err == nil {
		return path
	}
	return ""
}

func ExtractPreview(ctx context.Context, mimeType string, inBytes []byte) ([]byte, int, int, error) {
	log := logging.ExtractLogger(ctx)

	execPath := getFFMpegPath()
	if execPath == "" {
		return nil, 0, 0, nil
	}

	if !strings.HasPrefix(mimeType, "video") {
		return nil, 0, 0, nil
	}

	file, err := os.CreateTemp("", "hmnasset")
	if err != nil {
		return nil, 0, 0, oops.New(err, "Failed to create temp file for preview generation")
	}
	defer os.Remove(file.Name())
	_, err = file.Write(inBytes)
	if err != nil {
		return nil, 0, 0, oops.New(err, "Failed to write to temp file for preview generation")
	}
	err = file.Close()
	if err != nil {
		return nil, 0, 0, oops.New(err, "Failed to close temp file for preview generation")
	}

	args := fmt.Sprintf("-i %s -filter_complex [0]select=gte(n\\,1)[s0] -map [s0] -c:v mjpeg -f mjpeg -vframes 1 pipe:1", file.Name())
	if config.Config.PreviewGeneration.CPULimitPath != "" {
		args = fmt.Sprintf("-l 10 -- %s %s", execPath, args)
		execPath = config.Config.PreviewGeneration.CPULimitPath
	}
	ffmpegCmd := exec.CommandContext(ctx, execPath, strings.Split(args, " ")...)

	var output bytes.Buffer
	var errorOut bytes.Buffer
	ffmpegCmd.Stdout = &output
	ffmpegCmd.Stderr = &errorOut
	err = ffmpegCmd.Run()
	if err != nil {
		log.Error().Str("ffmpeg output", errorOut.String()).Msg("FFMpeg returned error while generating preview thumbnail")
		return nil, 0, 0, oops.New(err, "FFMpeg failed for preview generation")
	}

	imageBytes := output.Bytes()
	cfg, _, err := image.DecodeConfig(bytes.NewBuffer(imageBytes))
	if err != nil {
		log.Error().Err(err).Msg("failed to get width/height from video thumbnail")
		return nil, 0, 0, oops.New(err, "FFMpeg failed for preview generation")
	}

	return imageBytes, cfg.Width, cfg.Height, nil
}

func BackgroundPreviewGeneration(ctx context.Context, conn *pgxpool.Pool) jobs.Job {
	log := logging.ExtractLogger(ctx).With().Str("module", "preview_gen").Logger()
	job := jobs.New()

	go func() {
		defer job.Done()
		log.Debug().Msg("Starting preview gen job")

		if getFFMpegPath() == "" {
			log.Warn().Msg("Couldn't find ffmpeg! No thumbnails will be generated.")
			return
		}

		assets, err := db.Query[models.Asset](ctx, conn,
			`
			SELECT $columns
			FROM asset
			WHERE
				mime_type LIKE 'video%'
				AND (
					thumbnail_s3_key IS NULL
					OR thumbnail_s3_key = ''
					OR thumbnail_s3_key LIKE '%.png'
					OR width = 0
					OR height = 0
				)
			`,
		)
		if err != nil {
			log.Error().Err(oops.New(err, "Failed to fetch assets for preview generation")).Msg("Asset preview generation job failed")
			return
		}

		log.Debug().Int("Num assets", len(assets)).Msg("Processing...")

		for _, asset := range assets {
			select {
			case <-ctx.Done():
				return
			default:
			}

			log := log.With().Str("AssetID", asset.ID.String()).Logger()
			ctx := logging.AttachLoggerToContext(&log, ctx)

			log.Debug().Msg("Generating preview")
			assetUrl := hmnurl.BuildS3Asset(asset.S3Key)
			resp, err := http.Get(assetUrl)
			if err != nil || resp.StatusCode != 200 {
				log.Error().Err(err).Msg("Failed to fetch asset file for preview generation")
				continue
			}
			defer resp.Body.Close()
			body, err := io.ReadAll(resp.Body)
			if err != nil {
				log.Error().Err(err).Msg("Failed to read asset body for preview generation")
				continue
			}
			thumbBytes, width, height, err := ExtractPreview(ctx, asset.MimeType, body)
			if err != nil {
				log.Error().Err(err).Msg("Failed to run extraction for preview generation")
				continue
			} else if len(thumbBytes) > 0 {
				keyStr := AssetKey(asset.ID.String(), fmt.Sprintf("%s_thumb.jpg", asset.ID.String()))
				thumbnailType := "image/jpeg"
				_, err = client.PutObject(ctx, &s3.PutObjectInput{
					Bucket:      &config.Config.DigitalOcean.AssetsSpacesBucket,
					Key:         &keyStr,
					Body:        bytes.NewReader(thumbBytes),
					ACL:         types.ObjectCannedACLPublicRead,
					ContentType: &thumbnailType,
				})
				if err != nil {
					log.Error().Err(err).Msg("Failed to upload thumbnail for video")
					continue
				}

				_, err = conn.Exec(ctx,
					`
					UPDATE asset
					SET
						thumbnail_s3_key = $1,
						width = $2,
						height = $3
					WHERE asset.id = $4
					`,
					keyStr,
					width,
					height,
					asset.ID,
				)
				if err != nil {
					log.Error().Err(err).Msg("Failed to update asset for preview generation")
					continue
				}
				log.Debug().Msg("Generated preview successfully!")
			} else {
				log.Debug().Msg("No error, but no thumbnail was generated, skipping")
			}
		}
		log.Debug().Msg("No more previews to generate")
	}()

	return job
}