kefu/im/backend/internal/app/media_upload_test.go

package app

import (
	"bytes"
	"context"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"image"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

type mediaUploadDB struct {
	directory                                  string
	inserted, finalized, deleted, failFinalize bool
	publicURL, mime                            string
	size                                       int64
}
type mediaUploadConnector struct{ db *mediaUploadDB }

func (c mediaUploadConnector) Connect(context.Context) (driver.Conn, error) { return c.db, nil }
func (c mediaUploadConnector) Driver() driver.Driver                        { return oauthTestDriver{} }
func (*mediaUploadDB) Prepare(string) (driver.Stmt, error) {
	return nil, errors.New("unexpected prepare")
}
func (*mediaUploadDB) Begin() (driver.Tx, error) { return nil, errors.New("unexpected transaction") }
func (*mediaUploadDB) Close() error              { return nil }
func (s *mediaUploadDB) QueryContext(_ context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	if !strings.HasPrefix(query, "SELECT config_value,value_type FROM system_configs") {
		return nil, errors.New("unexpected query")
	}
	values := map[string]string{"storage.provider": "local", "storage.object_prefix": "media", "storage.local.directory": s.directory, "storage.local.public_base_url": "https://media.example.com/uploads"}
	value, ok := values[args[0].Value.(string)]
	if !ok {
		return &oauthTestRows{columns: []string{"config_value", "value_type"}}, nil
	}
	return oauthRow(value, "text"), nil
}
func (s *mediaUploadDB) ExecContext(_ context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	switch {
	case strings.HasPrefix(query, "INSERT INTO media_assets"):
		s.inserted = true
		s.mime = args[5].Value.(string)
		s.size = args[6].Value.(int64)
		return provisioningInsertResult{}, nil
	case strings.HasPrefix(query, "UPDATE media_assets SET public_url="):
		if s.failFinalize {
			return nil, errors.New("finalization unavailable")
		}
		s.publicURL = args[0].Value.(string)
		s.finalized = true
	case strings.HasPrefix(query, "DELETE FROM media_assets"):
		s.deleted = true
	default:
		return nil, errors.New("unexpected exec")
	}
	return driver.RowsAffected(1), nil
}

func uploadTestApp(t *testing.T) (*App, *mediaUploadDB) {
	t.Helper()
	store := &mediaUploadDB{directory: t.TempDir()}
	db := sql.OpenDB(mediaUploadConnector{store})
	t.Cleanup(func() { _ = db.Close() })
	return &App{db: db, config: Config{MediaDir: store.directory, Environment: "development"}}, store
}
func uploadTestRequest(t *testing.T, a *App, source []byte, purpose string) *httptest.ResponseRecorder {
	t.Helper()
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	part, err := writer.CreateFormFile("file", "camera.png")
	if err != nil {
		t.Fatal(err)
	}
	if _, err = part.Write(source); err != nil {
		t.Fatal(err)
	}
	if purpose != "" {
		if err = writer.WriteField("purpose", purpose); err != nil {
			t.Fatal(err)
		}
	}
	if err = writer.Close(); err != nil {
		t.Fatal(err)
	}
	r := httptest.NewRequest(http.MethodPost, "/api/v1/media/upload", &body)
	r.Header.Set("Content-Type", writer.FormDataContentType())
	r = r.WithContext(context.WithValue(r.Context(), identityKey{}, identity{ID: 42, Role: "user"}))
	w := httptest.NewRecorder()
	a.uploadMedia(w, r)
	return w
}

func TestAvatarUploadStoresAndServesAllVariants(t *testing.T) {
	a, store := uploadTestApp(t)
	source := avatarPNG(t, 1000, 800)
	w := uploadTestRequest(t, a, source, "avatar")
	if w.Code != 200 {
		t.Fatalf("upload failed: %s", w.Body.String())
	}
	var payload struct {
		Data struct {
			URL          string `json:"url"`
			ThumbnailURL string `json:"thumbnailUrl"`
			OriginalURL  string `json:"originalUrl"`
		} `json:"data"`
	}
	if err := json.Unmarshal(w.Body.Bytes(), &payload); err != nil {
		t.Fatal(err)
	}
	data := payload.Data
	if !store.inserted || !store.finalized || store.publicURL != data.URL || store.mime != "image/jpeg" {
		t.Fatal("media record does not identify the display image")
	}
	for _, variant := range []struct {
		url string
		max int
	}{{data.URL, 640}, {data.ThumbnailURL, 256}, {data.OriginalURL, 0}} {
		u, err := url.Parse(variant.url)
		if err != nil || !mediaNamePattern.MatchString(filepath.Base(u.Path)) {
			t.Fatal("invalid variant URL", variant.url)
		}
		content, err := os.ReadFile(filepath.Join(store.directory, filepath.Base(u.Path)))
		if err != nil {
			t.Fatal(err)
		}
		if variant.max == 0 {
			if !bytes.Equal(content, source) {
				t.Fatal("original upload was not preserved")
			}
		} else {
			cfg, format, err := image.DecodeConfig(bytes.NewReader(content))
			if err != nil || format != "jpeg" || cfg.Width != variant.max || cfg.Height > variant.max {
				t.Fatal("invalid variant dimensions", err)
			}
			if variant.url == data.URL && store.size != int64(len(content)) {
				t.Fatal("DB file_size does not match uploaded JPEG")
			}
		}
		response := httptest.NewRecorder()
		a.serveMedia(response, httptest.NewRequest(http.MethodGet, u.Path, nil))
		if response.Code != 200 || !bytes.Equal(response.Body.Bytes(), content) || !strings.Contains(response.Header().Get("Cache-Control"), "immutable") {
			t.Fatal("variant is not publicly cacheable")
		}
	}
}

func TestAvatarUploadFinalizationFailureCleansEveryObject(t *testing.T) {
	a, store := uploadTestApp(t)
	store.failFinalize = true
	w := uploadTestRequest(t, a, avatarPNG(t, 300, 200), "avatar")
	files, err := os.ReadDir(store.directory)
	if w.Code != 500 || !store.deleted || store.finalized || err != nil || len(files) != 0 {
		t.Fatal("failed upload retained objects or an active record", w.Code, err, len(files))
	}
}

func TestOrdinaryMediaUploadPreservesOriginal(t *testing.T) {
	a, store := uploadTestApp(t)
	source := avatarPNG(t, 800, 400)
	w := uploadTestRequest(t, a, source, "")
	if w.Code != 200 || store.mime != "image/png" || store.size != int64(len(source)) {
		t.Fatal("ordinary image upload changed", w.Body.String())
	}
	var payload struct {
		Data struct {
			URL          string `json:"url"`
			ThumbnailURL string `json:"thumbnailUrl"`
			OriginalURL  string `json:"originalUrl"`
		} `json:"data"`
	}
	if err := json.Unmarshal(w.Body.Bytes(), &payload); err != nil {
		t.Fatal(err)
	}
	if payload.Data.ThumbnailURL == "" || payload.Data.OriginalURL != "" || !strings.Contains(payload.Data.URL, "-im1.png") {
		t.Fatal("ordinary image did not publish a bounded list thumbnail")
	}
	files, err := os.ReadDir(store.directory)
	if err != nil || len(files) != 2 {
		t.Fatal("unexpected variants on ordinary media")
	}
	mainURL, _ := url.Parse(payload.Data.URL)
	data, err := os.ReadFile(filepath.Join(store.directory, filepath.Base(mainURL.Path)))
	if err != nil || !bytes.Equal(data, source) {
		t.Fatal("ordinary image was recompressed")
	}
	thumbnailURL, _ := url.Parse(payload.Data.ThumbnailURL)
	thumbnail, err := os.ReadFile(filepath.Join(store.directory, filepath.Base(thumbnailURL.Path)))
	config, format, decodeErr := image.DecodeConfig(bytes.NewReader(thumbnail))
	if err != nil || decodeErr != nil || format != "jpeg" || config.Width > 256 || config.Height > 256 {
		t.Fatal("ordinary image thumbnail is invalid", err, decodeErr)
	}
}

func TestRawMP3FrameUploadIsAcceptedAsVoice(t *testing.T) {
	a, store := uploadTestApp(t)
	// A native recorder may begin directly with an MPEG-1 Layer III frame and
	// omit the ID3 prefix recognized by http.DetectContentType.
	source := append([]byte{0xff, 0xfb, 0x90, 0x64}, bytes.Repeat([]byte{0x55}, 1024)...)
	w := uploadTestRequest(t, a, source, "voice")
	if w.Code != 200 || store.mime != "audio/mpeg" || store.size != int64(len(source)) {
		t.Fatal("raw MP3 voice upload failed", w.Code, w.Body.String(), store.mime,
			"detected", http.DetectContentType(source), detectUploadContentType(source))
	}
	if !strings.HasSuffix(store.publicURL, ".mp3") {
		t.Fatal("voice upload was not published as MP3", store.publicURL)
	}
}

func TestInvalidAvatarUploadCreatesNoRecords(t *testing.T) {
	for _, purpose := range []string{"avatar", "unknown"} {
		a, store := uploadTestApp(t)
		w := uploadTestRequest(t, a, []byte("GIF89a\x01\x00\x01\x00broken"), purpose)
		if w.Code != 400 || store.inserted {
			t.Fatal("invalid upload created media records", w.Code)
		}
	}
}

func TestMediaThumbnailURLOnlyRewritesVersionedImages(t *testing.T) {
	tests := map[string]string{
		"https://cdn.example.com/media/42-1-im1.png":             "https://cdn.example.com/media/42-1-im1-thumb.jpg",
		"https://cdn.example.com/media/42-1-im1-thumb.jpg":       "https://cdn.example.com/media/42-1-im1-thumb.jpg",
		"https://cdn.example.com/media/legacy.png":               "https://cdn.example.com/media/legacy.png",
		"https://cdn.example.com/media/42-1-im1.png?signature=x": "https://cdn.example.com/media/42-1-im1.png?signature=x",
	}
	for source, want := range tests {
		if got := mediaThumbnailURL(source); got != want {
			t.Errorf("mediaThumbnailURL(%q) = %q, want %q", source, got, want)
		}
	}
}