NianAIGC/backend/internal/httpapi/assets.go

package httpapi

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"strings"

	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/assets"
	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/publicapi"
)

const (
	defaultAssetJSONBytes   int64 = 1 << 20
	defaultAssetUploadBytes int64 = 20 << 20
)

type PublicAssetAuthenticator interface {
	Authenticate(*http.Request) (publicapi.PublicClient, string, error)
}

type AssetsConfig struct {
	MaxJSONBytes   int64
	MaxUploadBytes int64
}

type assetsHandler struct {
	service  *assets.Service
	platform *PlatformAuthorizer
	public   PublicAssetAuthenticator
	config   AssetsConfig
}

func NewAssetsHandler(service *assets.Service, platform *PlatformAuthorizer, public PublicAssetAuthenticator, config AssetsConfig) (http.Handler, error) {
	if service == nil || platform == nil || public == nil {
		return nil, errors.New("assets HTTP dependencies are not configured")
	}
	if config.MaxJSONBytes <= 0 {
		config.MaxJSONBytes = defaultAssetJSONBytes
	}
	if config.MaxUploadBytes <= 0 {
		config.MaxUploadBytes = defaultAssetUploadBytes
	}
	return &assetsHandler{service: service, platform: platform, public: public, config: config}, nil
}

func (h *assetsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	route, values := matchAssetRoute(r.URL.Path)
	if route == "" {
		http.NotFound(w, r)
		return
	}
	allow := assetRouteAllow(route)
	if r.Method == http.MethodOptions {
		w.Header().Set("Allow", allow)
		w.WriteHeader(http.StatusNoContent)
		return
	}
	method := r.Method
	if method == http.MethodHead && strings.Contains(allow, http.MethodHead) {
		method = http.MethodGet
		w = headResponseWriter{ResponseWriter: w}
	}
	if !methodAllowed(allow, method) {
		w.Header().Set("Allow", allow)
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	switch route {
	case "platform-collection":
		h.platformCollection(w, r, method)
	case "platform-upload":
		h.upload(w, r, false)
	case "platform-item":
		h.delete(w, r, values[0])
	case "platform-download":
		h.download(w, r, false, values[0])
	case "public-collection":
		h.publicCollection(w, r, method)
	case "public-item":
		h.publicGet(w, r, values[0])
	case "public-download":
		h.download(w, r, true, values[0])
	case "served-file":
		h.serveStored(w, r, values[0])
	}
}

func (h *assetsHandler) platformCollection(w http.ResponseWriter, r *http.Request, method string) {
	scope, ok := h.platformScope(w, r)
	if !ok {
		return
	}
	if method == http.MethodGet {
		values, err := h.service.List(r.Context(), scope)
		if err != nil {
			writeAssetError(w, err, false, "")
			return
		}
		if values == nil {
			values = []assets.Asset{}
		}
		writeJSON(w, http.StatusOK, map[string]any{"assets": values})
		return
	}
	var input createAssetInput
	if !decodeAssetJSON(w, r, h.config.MaxJSONBytes, &input) {
		return
	}
	command := assets.CreateExternalCommand{URL: input.URL, Name: input.Name, Kind: input.Kind, Source: input.Source, Tags: input.Tags}
	created, err := h.service.CreateExternal(r.Context(), scope, command)
	if err != nil {
		writeAssetError(w, err, false, "")
		return
	}
	writeJSON(w, http.StatusCreated, map[string]any{"asset": created})
}

func (h *assetsHandler) publicCollection(w http.ResponseWriter, r *http.Request, method string) {
	client, scope, ok := h.publicScope(w, r)
	if !ok {
		return
	}
	if method == http.MethodGet {
		values, err := h.service.List(r.Context(), scope)
		if err != nil {
			writeAssetError(w, err, true, "")
			return
		}
		if values == nil {
			values = []assets.Asset{}
		}
		writeJSON(w, http.StatusOK, map[string]any{"assets": values})
		return
	}
	if strings.Contains(strings.ToLower(r.Header.Get("Content-Type")), "multipart/form-data") {
		h.uploadScope(w, r, scope, true)
		return
	}
	_ = client
	var input createAssetInput
	if !decodeAssetJSON(w, r, h.config.MaxJSONBytes, &input) {
		return
	}
	command := assets.CreateExternalCommand{URL: input.URL, Name: input.Name, Kind: input.Kind, Tags: input.Tags}
	created, err := h.service.CreateExternal(r.Context(), scope, command)
	if err != nil {
		writeAssetError(w, err, true, "")
		return
	}
	writeJSON(w, http.StatusCreated, map[string]any{"asset": created})
}

type createAssetInput struct {
	URL    string        `json:"url"`
	Name   string        `json:"name"`
	Kind   assets.Kind   `json:"kind"`
	Source assets.Source `json:"source"`
	Tags   []string      `json:"tags"`
}

func (h *assetsHandler) upload(w http.ResponseWriter, r *http.Request, public bool) {
	scope, ok := h.platformScope(w, r)
	if !ok {
		return
	}
	h.uploadScope(w, r, scope, public)
}

func (h *assetsHandler) uploadScope(w http.ResponseWriter, r *http.Request, scope assets.Scope, public bool) {
	r.Body = http.MaxBytesReader(w, r.Body, h.config.MaxUploadBytes)
	reader, err := r.MultipartReader()
	if err != nil {
		writeAssetError(w, err, public, "")
		return
	}
	type pendingUpload struct {
		data                  []byte
		fileName, contentType string
	}
	pending := make([]pendingUpload, 0)
	for {
		part, nextErr := reader.NextPart()
		if errors.Is(nextErr, io.EOF) {
			break
		}
		if nextErr != nil {
			writeMultipartError(w, nextErr, public)
			return
		}
		if part.FormName() != "files" || part.FileName() == "" {
			_ = part.Close()
			continue
		}
		data, readErr := io.ReadAll(part)
		_ = part.Close()
		if readErr != nil {
			writeMultipartError(w, readErr, public)
			return
		}
		contentType := part.Header.Get("Content-Type")
		if contentType == "" {
			contentType = "application/octet-stream"
		}
		pending = append(pending, pendingUpload{data: data, fileName: part.FileName(), contentType: contentType})
	}
	if len(pending) == 0 {
		writeAssetJSONError(w, http.StatusBadRequest, "No files uploaded.")
		return
	}
	created := make([]assets.Asset, 0, len(pending))
	for _, file := range pending {
		a, createErr := h.service.Upload(r.Context(), scope, assets.UploadCommand{Bytes: file.data, FileName: file.fileName, ContentType: file.contentType, Origin: requestOrigin(r)})
		if createErr != nil {
			writeAssetError(w, createErr, public, "")
			return
		}
		created = append(created, a)
	}
	writeJSON(w, http.StatusCreated, map[string]any{"assets": created})
}

func (h *assetsHandler) delete(w http.ResponseWriter, r *http.Request, id string) {
	scope, ok := h.platformScope(w, r)
	if !ok {
		return
	}
	_, err := h.service.Delete(r.Context(), scope, id)
	if err != nil {
		writeAssetError(w, err, false, "资产不存在")
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"ok": true, "deletedAssetId": id})
}

func (h *assetsHandler) publicGet(w http.ResponseWriter, r *http.Request, id string) {
	_, scope, ok := h.publicScope(w, r)
	if !ok {
		return
	}
	a, err := h.service.Get(r.Context(), scope, id)
	if err != nil {
		writeAssetError(w, err, true, "Asset not found.")
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"asset": a})
}

func (h *assetsHandler) download(w http.ResponseWriter, r *http.Request, public bool, id string) {
	var scope assets.Scope
	var ok bool
	if public {
		_, scope, ok = h.publicScope(w, r)
	} else {
		scope, ok = h.platformScope(w, r)
	}
	if !ok {
		return
	}
	a, err := h.service.Get(r.Context(), scope, id)
	if err != nil {
		if public {
			writeAssetError(w, err, true, "Asset not found.")
		} else {
			writeAssetError(w, err, false, "资产不存在")
		}
		return
	}
	inline := r.URL.Query().Get("inline") == "1"
	if inline {
		signedURL, signed, signErr := h.service.SignedDownloadURL(r.Context(), scope, id, assets.DefaultSignedURLTTL)
		if signErr != nil {
			writeAssetError(w, signErr, public, "")
			return
		}
		if signed {
			w.Header().Set("Cache-Control", "private, no-store")
			w.Header().Set("Location", signedURL)
			w.WriteHeader(http.StatusTemporaryRedirect)
			return
		}
	}
	blob, err := h.service.Download(r.Context(), scope, id)
	if err != nil {
		if public {
			writeAssetError(w, err, true, "Asset file is not downloadable.")
		} else {
			writeAssetError(w, err, false, "资产文件不可下载")
		}
		return
	}
	defer blob.Body.Close()
	disposition := contentDisposition(a.Name)
	if inline {
		disposition = inlineContentDisposition(a.Name)
	}
	writeBlob(w, blob, "private, no-store", disposition)
}

func (h *assetsHandler) serveStored(w http.ResponseWriter, r *http.Request, key string) {
	scope, ok := h.platformScope(w, r)
	if !ok {
		return
	}
	blob, err := h.service.DownloadPath(r.Context(), scope, key)
	if err != nil {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}
	defer blob.Body.Close()
	writeBlob(w, blob, "public, max-age=31536000, immutable", "")
}

func (h *assetsHandler) platformScope(w http.ResponseWriter, r *http.Request) (assets.Scope, bool) {
	session, err := h.platform.Authorize(r, PlatformApp)
	if err != nil {
		writeAssetError(w, err, false, "")
		return assets.Scope{}, false
	}
	return assets.PlatformScope(session.User.ID), true
}
func (h *assetsHandler) publicScope(w http.ResponseWriter, r *http.Request) (publicapi.PublicClient, assets.Scope, bool) {
	client, _, err := h.public.Authenticate(r)
	if err != nil {
		writeAssetError(w, err, true, "")
		return publicapi.PublicClient{}, assets.Scope{}, false
	}
	return client, assets.PublicScope(client.ID), true
}

func decodeAssetJSON(w http.ResponseWriter, r *http.Request, limit int64, target any) bool {
	r.Body = http.MaxBytesReader(w, r.Body, limit)
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(target); err != nil {
		if isTooLarge(err) {
			writeAssetJSONError(w, http.StatusRequestEntityTooLarge, "Request body is too large.")
		} else {
			writeAssetJSONError(w, http.StatusBadRequest, "Invalid request body.")
		}
		return false
	}
	return true
}
func writeMultipartError(w http.ResponseWriter, err error, public bool) {
	if isTooLarge(err) {
		writeAssetJSONError(w, http.StatusRequestEntityTooLarge, "Request body is too large.")
		return
	}
	writeAssetError(w, err, public, "")
}
func isTooLarge(err error) bool { var max *http.MaxBytesError; return errors.As(err, &max) }
func writeAssetError(w http.ResponseWriter, err error, public bool, notFound string) {
	if errors.Is(err, assets.ErrNotFound) || errors.Is(err, assets.ErrBlobNotFound) {
		if notFound == "" {
			if public {
				notFound = "Asset not found."
			} else {
				notFound = "资产不存在"
			}
		}
		writeAssetJSONError(w, http.StatusNotFound, notFound)
		return
	}
	var platformErr *PlatformAuthError
	if errors.As(err, &platformErr) {
		writeAssetJSONError(w, platformErr.Status, platformErr.Message)
		return
	}
	var publicErr *publicapi.AuthError
	if errors.As(err, &publicErr) {
		writeAssetJSONError(w, publicErr.Status, publicErr.Message)
		return
	}
	if err != nil && (err.Error() == "url is required" || strings.Contains(err.Error(), "multipart")) {
		writeAssetJSONError(w, http.StatusBadRequest, err.Error())
		return
	}
	writeAssetJSONError(w, http.StatusInternalServerError, "Internal server error.")
}
func writeAssetJSONError(w http.ResponseWriter, status int, message string) {
	writeJSON(w, status, map[string]string{"error": message})
}
func writeBlob(w http.ResponseWriter, blob assets.Blob, cache, disposition string) {
	if blob.ContentType == "" {
		blob.ContentType = "application/octet-stream"
	}
	w.Header().Set("Content-Type", blob.ContentType)
	if blob.Size >= 0 {
		w.Header().Set("Content-Length", fmt.Sprint(blob.Size))
	}
	w.Header().Set("Cache-Control", cache)
	if disposition != "" {
		w.Header().Set("Content-Disposition", disposition)
	}
	w.WriteHeader(http.StatusOK)
	_, _ = io.Copy(w, blob.Body)
}
func contentDisposition(name string) string {
	return namedContentDisposition("attachment", name)
}
func inlineContentDisposition(name string) string {
	return namedContentDisposition("inline", name)
}
func namedContentDisposition(kind, name string) string {
	clean := strings.TrimSpace(strings.NewReplacer("\r", "_", "\n", "_", "/", "_", "\\", "_").Replace(name))
	if clean == "" {
		clean = "download"
	}
	var ascii strings.Builder
	for _, r := range clean {
		if r >= 0x20 && r <= 0x7e && r != '"' {
			ascii.WriteRune(r)
		} else {
			ascii.WriteByte('_')
		}
	}
	return kind + `; filename="` + ascii.String() + `"; filename*=UTF-8''` + url.PathEscape(clean)
}
func requestOrigin(r *http.Request) string {
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	if forwarded := r.Header.Get("X-Forwarded-Proto"); forwarded != "" {
		scheme = strings.TrimSpace(strings.Split(forwarded, ",")[0])
	}
	return scheme + "://" + r.Host
}
func methodAllowed(allow, method string) bool {
	for _, v := range strings.Split(allow, ",") {
		if strings.TrimSpace(v) == method {
			return true
		}
	}
	return false
}

type headResponseWriter struct{ http.ResponseWriter }

func (headResponseWriter) Write(p []byte) (int, error) { return len(p), nil }

func matchAssetRoute(value string) (string, []string) {
	if value == "/api/assets" {
		return "platform-collection", nil
	}
	if value == "/api/assets/upload" {
		return "platform-upload", nil
	}
	if value == "/api/v1/assets" {
		return "public-collection", nil
	}
	if strings.HasPrefix(value, "/api/assets/") {
		rest := strings.TrimPrefix(value, "/api/assets/")
		if rest != "" && !strings.Contains(rest, "/") {
			return "platform-item", []string{rest}
		}
		if strings.HasSuffix(rest, "/download") && strings.Count(rest, "/") == 1 {
			return "platform-download", []string{strings.TrimSuffix(rest, "/download")}
		}
	}
	if strings.HasPrefix(value, "/api/v1/assets/") {
		rest := strings.TrimPrefix(value, "/api/v1/assets/")
		if rest != "" && !strings.Contains(rest, "/") {
			return "public-item", []string{rest}
		}
		if strings.HasSuffix(rest, "/download") && strings.Count(rest, "/") == 1 {
			return "public-download", []string{strings.TrimSuffix(rest, "/download")}
		}
	}
	for _, prefix := range []string{"/uploads/", "/generated-results/"} {
		if strings.HasPrefix(value, prefix) {
			rest := strings.TrimPrefix(value, "/")
			if rest != "" && path.Clean(rest) == rest && !strings.Contains(rest, "\\") {
				return "served-file", []string{rest}
			}
		}
	}
	return "", nil
}
func assetRouteAllow(route string) string {
	switch route {
	case "platform-collection", "public-collection":
		return "GET, HEAD, POST, OPTIONS"
	case "platform-upload":
		return "POST, OPTIONS"
	case "platform-item":
		return "DELETE, OPTIONS"
	default:
		return "GET, HEAD, OPTIONS"
	}
}