summercms/modules/lagoon/attach/store.go

package attach

import (
	"bufio"
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"mime"
	"net/http"
	"path"
	"regexp"
	"slices"
	"strings"

	"gocloud.dev/blob"
	"gorm.io/gorm"
)

// sniffBytes is how much of an upload Store reads ahead for the content
// sniff and the image guard.
const sniffBytes = 1 << 20

var (
	// ErrTooLarge is returned by Store when the body exceeds Limits.MaxBytes.
	ErrTooLarge = errors.New("attach: file is too large")
	// ErrFileType is returned by Store when the file name's extension is
	// missing, malformed or not in the allowed extension list.
	ErrFileType = errors.New("attach: file type is not allowed")
	// ErrMIMEType is returned by Store when the content type matches none
	// of Limits.MIMETypes.
	ErrMIMEType = errors.New("attach: file content type is not allowed")
	// ErrNotImage is returned by Store in image mode when the bytes are not
	// an image IsAllowedImage accepts.
	ErrNotImage = errors.New("attach: file is not an allowed image")
)

// DefaultImageExtensions is the extension list of an image upload when
// Limits.Extensions is empty: jpg, jpeg, png, gif and webp, the formats
// IsAllowedImage and the thumbnailer handle. WinterCMS's image list also
// has avif, bmp and svg; they are left out because nothing here decodes
// them and svg can carry script.
var DefaultImageExtensions = []string{"jpg", "jpeg", "png", "gif", "webp"}

// DefaultFileExtensions is the extension list of a file upload when
// Limits.Extensions is empty: WinterCMS's default list (winter/storm
// Filesystem\Definitions::defaultExtensions) minus the script-capable types
// svg, js, map, css, less, scss, swf and xml. The final list is avi, avif,
// bmp, doc, docx, eot, flv, gif, ico, ics, jpeg, jpg, mkv, mov, mp3, mp4,
// mpeg, ods, odt, ogg, pdf, png, ppt, pptx, rar, ttf, txt, wav, webm, webp,
// wmv, woff, woff2, xls, xlsx and zip.
var DefaultFileExtensions = []string{
	"avi", "avif", "bmp", "doc", "docx", "eot", "flv", "gif", "ico", "ics",
	"jpeg", "jpg", "mkv", "mov", "mp3", "mp4", "mpeg", "ods", "odt", "ogg",
	"pdf", "png", "ppt", "pptx", "rar", "ttf", "txt", "wav", "webm", "webp",
	"wmv", "woff", "woff2", "xls", "xlsx", "zip",
}

var extPattern = regexp.MustCompile(`^[a-z0-9]{1,10}$`)

// Upload is one file to store. FileName is the client's file name: only its
// extension and base name are used (for the allowed-type check and the
// file_name column); no part of it reaches a blob key. Body is read once,
// to the end or to the size limit. Public sets the row's is_public flag.
type Upload struct {
	FileName string
	Body     io.Reader
	Public   bool
}

// Limits restricts what Store accepts.
//
// MaxBytes is the largest body in bytes; 0 means no limit of its own (the
// caller's request body cap still applies). Extensions lists the allowed
// lower-case extensions without the dot; empty means DefaultImageExtensions
// when Image is set, else DefaultFileExtensions. MIMETypes, when not empty,
// must match the stored content type: an entry containing a slash is a MIME
// pattern such as "image/png" or "image/*", an entry without one is an
// extension. Image applies the image guard (IsAllowedImage) to the content.
type Limits struct {
	MaxBytes   int64
	Extensions []string
	MIMETypes  []string
	Image      bool
}

// Store saves an upload as an unattached system_files row.
//
// It accepts the client extension, lower-cased, only when it matches
// [a-z0-9]{1,10} and is allowed by Limits (else ErrFileType). It reads up to
// 1 MiB ahead to sniff the content type from the bytes; in image mode those
// bytes must pass IsAllowedImage (else ErrNotImage), and Limits.MIMETypes is
// checked against the sniffed type, or the extension's registered type when
// the sniff only says application/octet-stream (else ErrMIMEType). The body
// is then streamed into bucket at BlobKey of a server-generated disk name (22
// random lowercase hex characters, a dot and the extension); a body longer
// than Limits.MaxBytes aborts the write, deletes the key and returns
// ErrTooLarge. Finally it inserts the row with empty attachment columns,
// is_public from Upload.Public, the byte size and the content type, and sets
// sort_order to the new id as WinterCMS's Sortable trait does. When the row
// cannot be written the blob is deleted again.
//
// db may be a transaction. The blob is written before the row, so a caller
// whose transaction rolls back after Store returned must delete the
// returned file's BlobKeys itself.
func Store(ctx context.Context, db *gorm.DB, bucket *blob.Bucket, in Upload, lim Limits) (*File, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if db == nil {
		return nil, fmt.Errorf("attach: store db is nil")
	}
	if bucket == nil {
		return nil, fmt.Errorf("attach: bucket is nil")
	}
	if in.Body == nil {
		return nil, fmt.Errorf("attach: upload body is nil")
	}
	if lim.MaxBytes < 0 {
		return nil, fmt.Errorf("attach: negative size limit %d", lim.MaxBytes)
	}
	name := clientBaseName(in.FileName)
	ext := strings.ToLower(strings.TrimPrefix(path.Ext(name), "."))
	if !extPattern.MatchString(ext) || !slices.Contains(allowedExtensions(lim), ext) {
		return nil, fmt.Errorf("%w: %q", ErrFileType, ext)
	}

	br := bufio.NewReaderSize(in.Body, sniffBytes)
	head, err := br.Peek(sniffBytes)
	if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, bufio.ErrBufferFull) {
		return nil, fmt.Errorf("attach: read upload: %w", err)
	}
	if lim.MaxBytes > 0 && int64(len(head)) > lim.MaxBytes {
		return nil, ErrTooLarge
	}
	if lim.Image && !IsAllowedImage(head) {
		return nil, ErrNotImage
	}
	contentType := baseMediaType(http.DetectContentType(head))
	if contentType == "application/octet-stream" {
		if byExt := baseMediaType(mime.TypeByExtension("." + ext)); byExt != "" {
			contentType = byExt
		}
	}
	if len(lim.MIMETypes) > 0 && !mimeAllowed(lim.MIMETypes, contentType, ext) {
		return nil, fmt.Errorf("%w: %s", ErrMIMEType, contentType)
	}

	diskName, err := newDiskName(ext)
	if err != nil {
		return nil, err
	}
	key := BlobKey(diskName)
	size, err := writeBlob(ctx, bucket, key, br, contentType, lim.MaxBytes)
	if err != nil {
		return nil, err
	}

	public := in.Public
	f := &File{
		DiskName:    diskName,
		FileName:    name,
		FileSize:    size,
		ContentType: contentType,
		IsPublic:    &public,
	}
	q := db.Session(&gorm.Session{NewDB: true, Context: ctx})
	err = q.Transaction(func(tx *gorm.DB) error {
		if err := tx.Create(f).Error; err != nil {
			return err
		}
		f.SortOrder = int(f.ID)
		return tx.Model(&File{}).Where("id = ?", f.ID).Update("sort_order", f.SortOrder).Error
	})
	if err != nil {
		_ = deleteKey(context.WithoutCancel(ctx), bucket, key)
		return nil, fmt.Errorf("attach: store row: %w", err)
	}
	return f, nil
}

// writeBlob streams r into key and returns the byte count. With limit > 0 a
// body of more than limit bytes aborts the write and deletes the key.
func writeBlob(ctx context.Context, bucket *blob.Bucket, key string, r io.Reader, contentType string, limit int64) (int64, error) {
	writeCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	w, err := bucket.NewWriter(writeCtx, key, &blob.WriterOptions{ContentType: contentType})
	if err != nil {
		return 0, fmt.Errorf("attach: blob writer: %w", err)
	}
	src := r
	if limit > 0 {
		src = io.LimitReader(r, limit+1)
	}
	n, copyErr := io.Copy(w, src)
	if copyErr == nil && limit > 0 && n > limit {
		copyErr = ErrTooLarge
	}
	if copyErr != nil {
		// Cancelling the writer's context before Close discards the write.
		cancel()
		_ = w.Close()
		_ = deleteKey(context.WithoutCancel(ctx), bucket, key)
		if errors.Is(copyErr, ErrTooLarge) {
			return 0, ErrTooLarge
		}
		return 0, fmt.Errorf("attach: write upload: %w", copyErr)
	}
	if err := w.Close(); err != nil {
		_ = deleteKey(context.WithoutCancel(ctx), bucket, key)
		return 0, fmt.Errorf("attach: write upload: %w", err)
	}
	return n, nil
}

// clientBaseName is the last element of a client file name, with either
// slash style treated as a separator.
func clientBaseName(name string) string {
	name = strings.ReplaceAll(name, `\`, "/")
	if i := strings.LastIndex(name, "/"); i >= 0 {
		name = name[i+1:]
	}
	return strings.TrimSpace(name)
}

func allowedExtensions(lim Limits) []string {
	if len(lim.Extensions) == 0 {
		if lim.Image {
			return DefaultImageExtensions
		}
		return DefaultFileExtensions
	}
	out := make([]string, 0, len(lim.Extensions))
	for _, e := range lim.Extensions {
		out = append(out, strings.ToLower(strings.TrimPrefix(strings.TrimSpace(e), ".")))
	}
	return out
}

func baseMediaType(ct string) string {
	if ct == "" {
		return ""
	}
	mt, _, err := mime.ParseMediaType(ct)
	if err != nil {
		return strings.ToLower(strings.TrimSpace(strings.SplitN(ct, ";", 2)[0]))
	}
	return mt
}

// mimeAllowed reports whether contentType or ext matches one of patterns.
func mimeAllowed(patterns []string, contentType, ext string) bool {
	for _, p := range patterns {
		p = strings.ToLower(strings.TrimSpace(p))
		if p == "" {
			continue
		}
		if !strings.Contains(p, "/") {
			if strings.TrimPrefix(p, ".") == ext {
				return true
			}
			continue
		}
		pType, pSub, _ := strings.Cut(p, "/")
		cType, cSub, _ := strings.Cut(contentType, "/")
		if (pType == "*" || pType == cType) && (pSub == "*" || pSub == cSub) {
			return true
		}
	}
	return false
}

func newDiskName(ext string) (string, error) {
	raw := make([]byte, 11)
	if _, err := rand.Read(raw); err != nil {
		return "", fmt.Errorf("attach: disk name: %w", err)
	}
	return hex.EncodeToString(raw) + "." + ext, nil
}