cluster/handler/upload.go

package handler

import (
	"bytes"
	"strconv"
	"strings"

	"gitea.loveuer.com/loveuer/cluster/controller"
	"gitea.loveuer.com/loveuer/cluster/internal/registry/storage"

	"github.com/gofiber/fiber/v3"
)

// StartBlobUpload 开始 blob 上传
func StartBlobUpload(store storage.Storage) fiber.Handler {
	return func(c fiber.Ctx) error {
		repo := c.Locals("repo_name")
		repoStr := ""
		if repo != nil {
			repoStr = repo.(string)
		}
		if repoStr == "" {
			repoStr = strings.TrimPrefix(c.Params("name"), "/")
		}

		uuid, err := controller.StartBlobUpload(store, repoStr)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
				"errors": []fiber.Map{
					{
						"code":    "INTERNAL_ERROR",
						"message": err.Error(),
					},
				},
			})
		}

		// 返回上传 URL
		location := c.Path() + "/" + uuid
		c.Set("Location", location)
		c.Set("Docker-Upload-UUID", uuid)
		c.Set("Range", "0-0")
		return c.SendStatus(fiber.StatusAccepted)
	}
}

// PatchBlobUpload 上传 blob 数据块
func PatchBlobUpload(store storage.Storage) fiber.Handler {
	return func(c fiber.Ctx) error {
		uuid := c.Locals("uuid")
		uuidStr := ""
		if uuid != nil {
			uuidStr = uuid.(string)
		}

		// 获取 Range header
		rangeHeader := c.Get("Content-Range")
		if rangeHeader == "" {
			return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
				"errors": []fiber.Map{
					{
						"code":    "BAD_REQUEST",
						"message": "Content-Range header required",
					},
				},
			})
		}

		// 读取数据
		data := bytes.NewReader(c.Body())
		size, err := controller.PatchBlobUpload(store, uuidStr, data)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
				"errors": []fiber.Map{
					{
						"code":    "INTERNAL_ERROR",
						"message": err.Error(),
					},
				},
			})
		}

		location := c.Path()
		c.Set("Location", location)
		c.Set("Docker-Upload-UUID", uuidStr)
		c.Set("Range", "0-"+strconv.FormatInt(size-1, 10))
		return c.SendStatus(fiber.StatusNoContent)
	}
}

// PutBlobUpload 完成 blob 上传
func PutBlobUpload(store storage.Storage) fiber.Handler {
	return func(c fiber.Ctx) error {
		repo := c.Locals("repo_name")
		repoStr := ""
		if repo != nil {
			repoStr = repo.(string)
		}
		if repoStr == "" {
			repoStr = strings.TrimPrefix(c.Params("name"), "/")
		}
		uuid := c.Locals("uuid")
		uuidStr := ""
		if uuid != nil {
			uuidStr = uuid.(string)
		}

		// 获取 digest
		digest := c.Query("digest")
		if digest == "" {
			return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
				"errors": []fiber.Map{
					{
						"code":    "BAD_REQUEST",
						"message": "digest query parameter required",
					},
				},
			})
		}

		// 如果有请求体，先追加数据
		var data interface{} = nil
		if len(c.Body()) > 0 {
			data = bytes.NewReader(c.Body())
		}

		location, err := controller.PutBlobUpload(store, repoStr, uuidStr, digest, data, c.Path())
		if err != nil {
			return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
				"errors": []fiber.Map{
					{
						"code":    "BAD_REQUEST",
						"message": err.Error(),
					},
				},
			})
		}

		c.Set("Location", location)
		c.Set("Content-Length", "0")
		c.Set("Docker-Content-Digest", digest)
		return c.SendStatus(fiber.StatusCreated)
	}
}

// GetBlobUpload 获取上传状态
func GetBlobUpload(store storage.Storage) fiber.Handler {
	return func(c fiber.Ctx) error {
		uuid := c.Locals("uuid")
		uuidStr := ""
		if uuid != nil {
			uuidStr = uuid.(string)
		}

		size, err := controller.GetBlobUpload(store, uuidStr)
		if err != nil {
			return c.Status(fiber.StatusNotFound).JSON(fiber.Map{
				"errors": []fiber.Map{
					{
						"code":    "UPLOAD_UNKNOWN",
						"message": err.Error(),
					},
				},
			})
		}

		location := c.Path()
		c.Set("Location", location)
		c.Set("Docker-Upload-UUID", uuidStr)
		c.Set("Range", "0-"+strconv.FormatInt(size-1, 10))
		return c.SendStatus(fiber.StatusNoContent)
	}
}