repo.me/internal/handler/repo.go

package handler

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/loveuer/nf"
	"github.com/sirupsen/logrus"
	"io"
	"net/http"
	"net/url"
	"nf-repo/internal/interfaces"
	"nf-repo/internal/model"
	"nf-repo/internal/opt"
	"nf-repo/internal/util/r"
	"nf-repo/internal/util/rerr"
	"nf-repo/internal/util/tools"
	"strings"
)

func RepoSettings(c *nf.Ctx) error {
	return r.Resp200(c, nf.Map{
		"base_address": opt.BaseAddress,
	})
}

func RepoList(mh interfaces.ManifestHandler) nf.HandlerFunc {
	return func(c *nf.Ctx) error {
		type Req struct {
			Keyword string `json:"keyword" query:"keyword"`
			N       int    `json:"n" query:"n"`
			Last    int    `json:"last" query:"last"`
		}

		var (
			err     error
			re      *rerr.RepositoryError
			req     = new(Req)
			catalog *model.Catalog
		)

		if err = c.QueryParser(req); err != nil {
			return r.Resp400(c, err.Error())
		}

		if req.N == 0 {
			req.N = 20
		}

		if req.N > 1000 {
			return r.Resp400(c, "limit invalid: too big")
		}

		if catalog, re = mh.Catalog(c.Request.Context(), req.N, req.Last, req.Keyword); re != nil {
			return r.Resp(c, uint32(re.Status), "", re.Code, re.Message)
		}

		return r.Resp200(c, nf.Map{"list": catalog.Repos, "total": 0})
	}
}

func TagList(mh interfaces.ManifestHandler) nf.HandlerFunc {
	return func(c *nf.Ctx) error {
		type Req struct {
			Repo    string `json:"repo" query:"repo"`
			Keyword string `json:"keyword" query:"keyword"`
			N       int    `json:"n" query:"n"`
			Last    int    `json:"last" query:"last"`
		}

		var (
			err error
			req = new(Req)
			tag *model.Tag
			re  *rerr.RepositoryError
		)

		if err = c.QueryParser(req); err != nil {
			return r.Resp400(c, err.Error())
		}

		if req.Repo == "" {
			return r.Resp400(c, "repo invalid")
		}

		if req.N <= 0 {
			req.N = 20
		}

		if req.N > 1000 {
			return r.Resp400(c, "limit invalid: too big")
		}

		if tag, re = mh.Tags(c.Request.Context(), req.Repo, req.N, req.Last, req.Keyword); re != nil {
			return r.Resp(c, uint32(re.Status), "", re.Code, re.Message)
		}

		return r.Resp200(c, nf.Map{"list": tag.RepoTags, "total": tag.Tags})
	}
}

func ProxyDownloadImage(mh interfaces.ManifestHandler, bh interfaces.BlobHandler) nf.HandlerFunc {
	return func(c *nf.Ctx) error {
		type Req struct {
			Source string `json:"source"`
			Target string `json:"target"`
			Proxy  string `json:"proxy"`
		}

		type Msg struct {
			Target string `json:"target"`
			Action string `json:"action"`
			Msg    string `json:"msg"`
			Data   any    `json:"data"`
		}

		var (
			err       error
			req       = new(Req)
			puller    *remote.Puller
			tn        name.Tag
			des       *remote.Descriptor
			transport = &http.Transport{}
			img       v1.Image
			manifest  *v1.Manifest
			bs        []byte
			mhash     model.Hash
			//progressCh = make(chan v1.Update, 16)
		)

		if err = c.BodyParser(req); err != nil {
			return r.Resp400(c, err.Error())
		}

		if req.Target == "" {
			return r.Resp400(c, "target invalid")
		}

		repoTags := strings.Split(req.Target, ":")
		if len(repoTags) != 2 {
			return r.Resp400(c, "target invalid: should only contain one ':'")
		}

		repo, tag := repoTags[0], repoTags[1]
		repo = strings.TrimPrefix(repo, opt.BaseAddress)
		repo = strings.Trim(repo, "/")

		if req.Source == "" {
			return r.Resp400(c, "source invalid")
		}

		imageName := fmt.Sprintf("%s/%s:%s", opt.BaseAddress, repo, tag)
		_ = imageName

		logrus.
			WithField("path", "handler.ProxyDownloadImage").
			WithField("req", *req).
			WithField("repo", repo).
			Debug()

		if req.Proxy != "" {
			var pu *url.URL
			if pu, err = url.Parse(req.Proxy); err != nil {
				return r.Resp400(c, err.Error())
			}

			transport.Proxy = http.ProxyURL(pu)
		}

		if puller, err = remote.NewPuller(
			remote.WithTransport(transport),
		); err != nil {
			return r.Resp500(c, err.Error())
		}

		if tn, err = name.NewTag(req.Source); err != nil {
			return r.Resp400(c, err.Error())
		}

		_ = c.SSEvent("pull", Msg{
			Target: req.Target,
			Action: "get",
			Msg:    "开始获取镜像信息",
		})
		_ = c.Flush()

		ctx := tools.TimeoutCtx(c.Request.Context())
		if des, err = puller.Get(ctx, tn); err != nil {
			return r.Resp500(c, err.Error())
		}

		if img, err = des.Image(); err != nil {
			return r.Resp500(c, err.Error())
		}

		if manifest, err = img.Manifest(); err != nil {
			return r.Resp500(c, err.Error())
		}

		total := model.ManifestCountSize(manifest)
		size := 0

		_ = c.SSEvent("pull", Msg{
			Target: req.Target,
			Action: "layer",
			Msg:    "正在获取镜像数据: 0.00%",
			Data:   map[string]any{"total": total, "size": size, "index_total": len(manifest.Layers) + 1, "index_size": 0},
		})
		_ = c.Flush()

		var (
			tly     v1.Layer
			tdigest v1.Hash
			treader io.ReadCloser
		)

		if tly, err = img.LayerByDigest(manifest.Config.Digest); err != nil {
			return r.Resp500(c, err.Error())
		}

		if tdigest, err = tly.Digest(); err != nil {
			return r.Resp500(c, err.Error())
		}

		if treader, err = tly.Uncompressed(); err != nil {
			return r.Resp500(c, err.Error())
		}
		defer treader.Close()

		if err = bh.Put(
			c.Request.Context(),
			repo,
			model.Hash{Algorithm: tdigest.Algorithm, Hex: tdigest.Hex},
			treader,
		); err != nil {
			return r.Resp500(c, err.Error())
		}

		size = size + int(manifest.Config.Size)
		_ = c.SSEvent("pull", Msg{
			Target: req.Target,
			Action: "layer",
			Msg:    fmt.Sprintf("正在获取镜像数据: %.2f%%", float64(size)/float64(total)*100),
			Data:   map[string]any{"total": model.ManifestCountSize(manifest), "size": size, "index_total": len(manifest.Layers) + 1, "index_size": 1},
		})
		_ = c.Flush()

		if bs, err = json.Marshal(manifest); err != nil {
			return r.Resp500(c, err.Error())
		}

		if mhash, _, err = model.SHA256(bytes.NewReader(bs)); err != nil {
			return r.Resp500(c, err.Error())
		}

		for idx := range manifest.Layers {
			var (
				reader io.ReadCloser
				lyHash v1.Hash
				ly     v1.Layer
			)

			lyHash = manifest.Layers[idx].Digest

			if ly, err = img.LayerByDigest(manifest.Layers[idx].Digest); err != nil {
				return r.Resp500(c, err.Error())
			}

			if reader, err = ly.Compressed(); err != nil {
				return r.Resp500(c, err.Error())
			}
			defer reader.Close()

			if err = bh.Put(c.Request.Context(), repo, model.Hash{Algorithm: lyHash.Algorithm, Hex: lyHash.Hex}, reader); err != nil {
				return r.Resp500(c, err.Error())
			}

			size = size + int(manifest.Layers[idx].Size)
			_ = c.SSEvent("pull", Msg{
				Target: req.Target,
				Action: "layer",
				Msg:    fmt.Sprintf("正在获取镜像数据: %.2f%%", float64(size)/float64(total)*100),
				Data:   map[string]any{"total": model.ManifestCountSize(manifest), "size": size, "index_total": len(manifest.Layers) + 1, "index_size": 2 + idx},
			})
			_ = c.Flush()
		}

		var re *rerr.RepositoryError

		if re = mh.Put(c.Request.Context(), repo, tag, mhash.String(), &model.RepoSimpleManifest{
			ContentType: string(manifest.MediaType),
			Blob:        bs,
		}); re != nil {
			return r.Resp(c, uint32(re.Status), re.Message, re.Code, nil)
		}

		_ = c.SSEvent("pull", Msg{
			Target: req.Target,
			Action: "done",
			Msg:    "获取成功",
			Data:   manifest,
		})

		return c.Flush()
	}
}