Warp/server/internal/mediaapi/mediaapi.go

// Package mediaapi is the media repository: uploads, downloads,
// thumbnails, and the cache of other servers' media.
//
// Content lives in a storage.BlobStore, named by its SHA-256; this package
// keeps the metadata that maps an mxc:// URI onto a blob, in the
// storage.MediaStore. The HTTP surface is elsewhere — the Client-Server
// endpoints in clientapi, the federation endpoints in federationapi — and
// both come here for everything but the response format.
//
// Media from other servers is fetched on first use, through the
// federation client, and kept: an mxc:// URI names immutable content, so a
// copy never goes stale. With federation off the client refuses to fetch,
// and only what was cached before is served.
package mediaapi

import (
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"mime"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"

	"github.com/dariuszjaniak/warp/server/internal/config"
	"github.com/dariuszjaniak/warp/server/internal/federation"
	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// Errors the API returns. The HTTP layers map them onto Matrix errors.
var (
	// ErrNotFound means there is no such media, or it cannot be had.
	ErrNotFound = errors.New("media not found")

	// ErrTooLarge means content is larger than media.max_upload_bytes.
	ErrTooLarge = errors.New("media is too large")

	// ErrQuotaExceeded means the upload would take the user past their
	// media quota.
	ErrQuotaExceeded = errors.New("media quota exceeded")

	// ErrNotYetUploaded means the media ID was reserved and its content
	// has not arrived within the time the request was willing to wait.
	ErrNotYetUploaded = errors.New("media has not been uploaded yet")

	// ErrCannotOverwrite means the content of a reserved ID was already
	// uploaded.
	ErrCannotOverwrite = errors.New("media content has already been uploaded")

	// ErrNotYours means somebody other than the user who reserved an ID
	// tried to upload its content.
	ErrNotYours = errors.New("media ID was reserved by another user")

	// ErrTooManyPending means the user holds as many reserved IDs as they
	// may.
	ErrTooManyPending = errors.New("too many pending uploads")

	// ErrCannotThumbnail means the media is not an image this server can
	// make a thumbnail of.
	ErrCannotThumbnail = errors.New("cannot make a thumbnail of this media")

	// ErrFederationDisabled means the media is on another server and
	// federation is off, or that server is not allowed.
	ErrFederationDisabled = errors.New("media is on another server and federation with it is off")

	// ErrRemote means the server the media is on failed to provide it.
	ErrRemote = errors.New("the server the media is on did not provide it")
)

// Remote fetches media from other servers. *federation.Client satisfies
// it.
type Remote interface {
	DownloadMedia(ctx context.Context, origin, mediaID string, wait time.Duration) (federation.MediaResponse, error)
	MediaThumbnail(ctx context.Context, origin, mediaID string, req federation.ThumbnailRequest, wait time.Duration) (federation.MediaResponse, error)
}

// Options configure the API.
type Options struct {
	// ServerName is this server's name: the origin of every upload.
	ServerName string

	Store storage.MediaStore
	Blobs storage.BlobStore

	// Remote fetches other servers' media. Nil means only local and
	// already cached media can be served.
	Remote Remote

	Config config.Media
	Logger *slog.Logger

	// Now is the clock, so tests can expire reservations without waiting.
	Now func() time.Time
}

// API is the media repository.
type API struct {
	serverName string
	store      storage.MediaStore
	blobs      storage.BlobStore
	cfg        config.Media
	log        *slog.Logger
	now        func() time.Time

	remoteMu sync.RWMutex
	remote   Remote

	// waiters wakes downloads blocked on a reserved ID.
	waiters *waiters

	// fetches makes concurrent requests for the same remote media, or
	// the same thumbnail, share one fetch or one resize.
	fetches singleflight.Group

	// resizing bounds how many thumbnails are made at once. Decoding an
	// image is the most memory a single request can cost this server.
	resizing chan struct{}
}

// New builds the API.
func New(opts Options) (*API, error) {
	switch {
	case opts.ServerName == "":
		return nil, errors.New("mediaapi: ServerName must not be empty")
	case opts.Store == nil, opts.Blobs == nil:
		return nil, errors.New("mediaapi: Store and Blobs are required")
	case opts.Config.MaxUploadBytes < 1:
		return nil, errors.New("mediaapi: Config.MaxUploadBytes must be positive")
	}
	if opts.Logger == nil {
		opts.Logger = slog.Default()
	}
	if opts.Now == nil {
		opts.Now = time.Now
	}
	if opts.Config.PendingUploadExpiry <= 0 {
		opts.Config.PendingUploadExpiry = 24 * time.Hour
	}
	if opts.Config.MaxPendingUploads < 1 {
		opts.Config.MaxPendingUploads = 1
	}
	if opts.Config.MaxThumbnailPixels < 1 {
		opts.Config.MaxThumbnailPixels = config.Default().Media.MaxThumbnailPixels
	}
	return &API{
		serverName: opts.ServerName,
		store:      opts.Store,
		blobs:      opts.Blobs,
		cfg:        opts.Config,
		log:        opts.Logger,
		now:        opts.Now,
		waiters:    newWaiters(),
		resizing:   make(chan struct{}, 4),
		remote:     opts.Remote,
	}, nil
}

// SetRemote hands the API the client it fetches other servers' media
// with. The server builds the media repository before federation, which
// needs a signing key it may not have; until this is called only local
// and already cached media is served.
func (a *API) SetRemote(remote Remote) {
	a.remoteMu.Lock()
	defer a.remoteMu.Unlock()
	a.remote = remote
}

// remoteClient reads the remote, or nil.
func (a *API) remoteClient() Remote {
	a.remoteMu.RLock()
	defer a.remoteMu.RUnlock()
	return a.remote
}

// MaxUploadBytes is the largest upload, for /media/config.
func (a *API) MaxUploadBytes() int64 { return a.cfg.MaxUploadBytes }

// ServerName is the origin of this server's media.
func (a *API) ServerName() string { return a.serverName }

func (a *API) nowMS() int64 { return a.now().UnixMilli() }

// Upload is content arriving from a client.
type Upload struct {
	UserID string

	// ContentType is the request's Content-Type; empty becomes
	// application/octet-stream.
	ContentType string

	// FileName is the ?filename= of the request, or "".
	FileName string

	Body io.Reader
}

// mediaIDLength is long enough that an ID cannot be guessed: 24
// characters of 62 is 142 bits.
const mediaIDLength = 24

const mediaIDAlphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"

// newMediaID makes a random media ID.
func newMediaID() string {
	// Bytes of 248 and over are thrown away rather than folded in with a
	// modulo, which would make the first eight letters more likely than
	// the rest. 248 is the largest multiple of 62 a byte holds.
	const accept = 256 - 256%len(mediaIDAlphabet)
	out := make([]byte, 0, mediaIDLength)
	buf := make([]byte, mediaIDLength)
	for len(out) < mediaIDLength {
		_, _ = rand.Read(buf)
		for _, c := range buf {
			if int(c) < accept && len(out) < mediaIDLength {
				out = append(out, mediaIDAlphabet[int(c)%len(mediaIDAlphabet)])
			}
		}
	}
	return string(out)
}

// MXC builds an mxc:// URI.
func MXC(origin, mediaID string) string { return "mxc://" + origin + "/" + mediaID }

// Upload stores new content and returns its mxc:// URI.
func (a *API) Upload(ctx context.Context, up Upload) (string, error) {
	content, err := a.storeContent(ctx, up)
	if err != nil {
		return "", err
	}
	m := storage.Media{
		Origin:      a.serverName,
		ContentType: content.ContentType,
		UploadName:  content.UploadName,
		Size:        content.Size,
		SHA256:      content.SHA256,
		UserID:      up.UserID,
		CreatedAtMS: a.nowMS(),
	}
	// A collision of 142 random bits does not happen; the loop is for the
	// day the generator is replaced by something that can.
	for range 3 {
		m.MediaID = newMediaID()
		err = a.store.CreateMedia(ctx, m)
		if !errors.Is(err, storage.ErrMediaExists) {
			break
		}
	}
	if err != nil {
		return "", err
	}
	a.log.Debug("media uploaded", "media_id", m.MediaID, "user_id", up.UserID,
		"content_type", m.ContentType, "bytes", m.Size)
	return MXC(m.Origin, m.MediaID), nil
}

// Create reserves a media ID whose content is uploaded later, returning
// its mxc:// URI and when the reservation lapses.
func (a *API) Create(ctx context.Context, userID string) (string, int64, error) {
	now := a.nowMS()
	// The sweep rides on reservations rather than a timer: expired rows
	// only matter to the count below, and nothing else reads them.
	if n, err := a.store.DeleteExpiredUploads(ctx, now); err != nil {
		a.log.Warn("could not delete expired media reservations", "error", err)
	} else if n > 0 {
		a.log.Debug("deleted expired media reservations", "count", n)
	}

	pending, err := a.store.PendingUploads(ctx, userID, now)
	if err != nil {
		return "", 0, err
	}
	if pending >= a.cfg.MaxPendingUploads {
		return "", 0, ErrTooManyPending
	}
	m := storage.Media{
		Origin:      a.serverName,
		MediaID:     newMediaID(),
		UserID:      userID,
		CreatedAtMS: now,
		Pending:     true,
		ExpiresAtMS: now + a.cfg.PendingUploadExpiry.Milliseconds(),
	}
	if err := a.store.CreateMedia(ctx, m); err != nil {
		return "", 0, err
	}
	return MXC(m.Origin, m.MediaID), m.ExpiresAtMS, nil
}

// UploadTo stores the content of a reserved ID.
func (a *API) UploadTo(ctx context.Context, origin, mediaID string, up Upload) error {
	if origin != a.serverName {
		return ErrNotFound
	}
	m, err := a.store.Media(ctx, origin, mediaID)
	switch {
	case errors.Is(err, storage.ErrNoSuchMedia):
		return ErrNotFound
	case err != nil:
		return err
	case m.UserID != up.UserID:
		return ErrNotYours
	case !m.Pending:
		return ErrCannotOverwrite
	case m.ExpiresAtMS <= a.nowMS():
		return ErrNotFound
	}

	content, err := a.storeContent(ctx, up)
	if err != nil {
		return err
	}
	err = a.store.CompleteUpload(ctx, origin, mediaID, content)
	switch {
	case errors.Is(err, storage.ErrMediaUploaded):
		return ErrCannotOverwrite
	case errors.Is(err, storage.ErrNoSuchMedia):
		return ErrNotFound
	case err != nil:
		return err
	}
	a.waiters.wake(origin, mediaID)
	return nil
}

// storeContent writes an upload's body to the blob store.
//
// The user's quota is applied by lowering the limit the content is
// written with to what the quota has left, so an upload that would go
// over is cut off as it arrives rather than stored and then refused. Two
// uploads racing for the last of a quota can both get in; a quota is a
// bound on what a user keeps here, not an accounting system.
func (a *API) storeContent(ctx context.Context, up Upload) (storage.MediaContent, error) {
	limit, quotaBound, err := a.uploadLimit(ctx, up.UserID)
	if err != nil {
		return storage.MediaContent{}, err
	}
	if limit <= 0 {
		return storage.MediaContent{}, ErrQuotaExceeded
	}
	blob, err := a.blobs.Put(ctx, up.Body, limit)
	switch {
	case errors.Is(err, storage.ErrBlobTooLarge) && quotaBound:
		return storage.MediaContent{}, ErrQuotaExceeded
	case errors.Is(err, storage.ErrBlobTooLarge):
		return storage.MediaContent{}, ErrTooLarge
	case err != nil:
		return storage.MediaContent{}, err
	}
	return storage.MediaContent{
		ContentType: normaliseContentType(up.ContentType),
		UploadName:  up.FileName,
		Size:        blob.Size,
		SHA256:      blob.SHA256,
	}, nil
}

// uploadLimit is the most one upload by a user may be: the server's
// upload limit, or what is left of the user's quota if that is less, in
// which case quotaBound is set.
func (a *API) uploadLimit(ctx context.Context, userID string) (limit int64, quotaBound bool, err error) {
	quota, err := a.Quota(ctx, userID)
	if err != nil {
		return 0, false, err
	}
	if quota.MaxBytes == 0 {
		return a.cfg.MaxUploadBytes, false, nil
	}
	left := quota.MaxBytes - quota.Usage.Bytes
	if left < a.cfg.MaxUploadBytes {
		return left, true, nil
	}
	return a.cfg.MaxUploadBytes, false, nil
}

// Quota is a user's media quota and what they are using of it.
type Quota struct {
	Usage storage.MediaUsage

	// MaxBytes is the quota in force; 0 is no limit.
	MaxBytes int64

	// Custom reports a quota set for this user through the admin API,
	// rather than media.max_user_storage_bytes.
	Custom bool
}

// Quota reads a user's quota and usage.
func (a *API) Quota(ctx context.Context, userID string) (Quota, error) {
	usage, err := a.store.MediaUsage(ctx, a.serverName, userID)
	if err != nil {
		return Quota{}, err
	}
	q := Quota{Usage: usage, MaxBytes: a.cfg.MaxUserStorageBytes}
	custom, ok, err := a.store.MediaQuota(ctx, userID)
	if err != nil {
		return Quota{}, err
	}
	if ok {
		q.MaxBytes, q.Custom = custom, true
	}
	return q, nil
}

// SetQuota gives a user a quota of their own; 0 is no limit.
func (a *API) SetQuota(ctx context.Context, userID string, maxBytes int64) error {
	return a.store.SetMediaQuota(ctx, userID, maxBytes)
}

// ResetQuota puts a user back on media.max_user_storage_bytes.
func (a *API) ResetQuota(ctx context.Context, userID string) error {
	return a.store.DeleteMediaQuota(ctx, userID)
}

// normaliseContentType keeps a Content-Type that parses and replaces one
// that does not. What is stored is what downloads are served with, and a
// header that does not parse is one every client would guess at.
func normaliseContentType(raw string) string {
	if strings.TrimSpace(raw) == "" {
		return "application/octet-stream"
	}
	mediaType, params, err := mime.ParseMediaType(raw)
	if err != nil {
		return "application/octet-stream"
	}
	return mime.FormatMediaType(mediaType, params)
}

// Content is media ready to be sent: the body and what to say about it.
type Content struct {
	ContentType string
	FileName    string
	Size        int64

	// SHA256 is the content's hash, which makes a perfect ETag.
	SHA256 string

	Body io.ReadSeekCloser
}

// FetchOptions are the parts of a download request that are not the
// media's name.
type FetchOptions struct {
	// Wait is how long to wait for a reserved ID's content. Zero does not
	// wait.
	Wait time.Duration

	// LocalOnly serves only media uploaded here: what another server is
	// asking for, which must not make this one fetch from a third.
	LocalOnly bool
}

// MaxWait caps what a request may ask to wait for content.
const MaxWait = 2 * time.Minute

// Download returns media's content.
func (a *API) Download(ctx context.Context, origin, mediaID string, opts FetchOptions) (Content, error) {
	if !validServerName(origin) || !ValidMediaID(mediaID) {
		return Content{}, ErrNotFound
	}
	if origin == a.serverName {
		m, err := a.localMedia(ctx, mediaID, opts.Wait)
		if err != nil {
			return Content{}, err
		}
		return a.open(ctx, m.ContentType, m.UploadName, m.Size, m.SHA256)
	}
	if opts.LocalOnly {
		return Content{}, ErrNotFound
	}
	m, err := a.remoteMedia(ctx, origin, mediaID, opts.Wait)
	if err != nil {
		return Content{}, err
	}
	return a.open(ctx, m.ContentType, m.UploadName, m.Size, m.SHA256)
}

// localMedia reads media uploaded here, waiting for a reserved ID's
// content if asked to.
func (a *API) localMedia(ctx context.Context, mediaID string, wait time.Duration) (storage.Media, error) {
	wait = min(max(wait, 0), MaxWait)
	deadline := a.now().Add(wait)
	for {
		// Subscribed before the read, so that an upload completing
		// between the read and the wait is not missed.
		woken, cancel := a.waiters.subscribe(a.serverName, mediaID)
		m, err := a.store.Media(ctx, a.serverName, mediaID)
		switch {
		case errors.Is(err, storage.ErrNoSuchMedia):
			cancel()
			return storage.Media{}, ErrNotFound
		case err != nil:
			cancel()
			return storage.Media{}, err
		case !m.Pending:
			cancel()
			return m, nil
		case m.ExpiresAtMS <= a.nowMS():
			cancel()
			return storage.Media{}, ErrNotFound
		}

		remaining := deadline.Sub(a.now())
		if remaining <= 0 {
			cancel()
			return storage.Media{}, ErrNotYetUploaded
		}
		timer := time.NewTimer(remaining)
		select {
		case <-woken:
		case <-timer.C:
		case <-ctx.Done():
		}
		timer.Stop()
		cancel()
		if ctx.Err() != nil {
			return storage.Media{}, ctx.Err()
		}
	}
}

// open opens a blob as Content.
func (a *API) open(ctx context.Context, contentType, fileName string, size int64, sha string) (Content, error) {
	body, err := a.blobs.Open(ctx, sha)
	if errors.Is(err, storage.ErrNoSuchBlob) {
		// The row is there and the file is not: somebody cleaned the
		// media directory by hand. Say so in the log; the client is told
		// the media is gone, which it is.
		a.log.Error("media content is missing from the blob store", "sha256", sha)
		return Content{}, ErrNotFound
	}
	if err != nil {
		return Content{}, err
	}
	return Content{ContentType: contentType, FileName: fileName, Size: size, SHA256: sha, Body: body}, nil
}

// ValidMediaID reports a media ID of the grammar the spec gives:
// letters, digits, underscore and hyphen. Anything else is refused before
// it reaches a database or another server.
func ValidMediaID(id string) bool {
	if id == "" || len(id) > 255 {
		return false
	}
	for _, c := range id {
		switch {
		case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9', c == '_', c == '-':
		default:
			return false
		}
	}
	return true
}

// validServerName is a cheap sanity check on the server name of an
// mxc:// URI: a host, an optional port, an IPv6 literal in brackets.
func validServerName(name string) bool {
	if name == "" || len(name) > 255 {
		return false
	}
	for _, c := range name {
		switch {
		case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9':
		case c == '.', c == '-', c == ':', c == '[', c == ']':
		default:
			return false
		}
	}
	return true
}

// String names media for logs.
func mediaName(origin, mediaID string) string { return fmt.Sprintf("%s/%s", origin, mediaID) }