cluster/internal/module/registry/handler.fetch.go

package registry

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"gitea.loveuer.com/loveuer/cluster/internal/model"
	"gitea.loveuer.com/loveuer/cluster/pkg/resp"
	"gitea.loveuer.com/loveuer/cluster/pkg/store"
	"github.com/gofiber/fiber/v3"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/google/go-containerregistry/pkg/name"
	goproxy "golang.org/x/net/proxy"
	"gorm.io/gorm"
)

type FetchImageRequest struct {
	Image string `json:"image"`
	Proxy string `json:"proxy"`
}

func RegistryImageFetch(ctx context.Context, db *gorm.DB, store store.Store) fiber.Handler {
	return func(c fiber.Ctx) error {
		var req FetchImageRequest
		if err := c.Bind().JSON(&req); err != nil {
			return resp.R400(c, "INVALID_REQUEST", nil, "invalid request body")
		}

		if req.Image == "" {
			return resp.R400(c, "MISSING_IMAGE", nil, "image name is required")
		}

		log.Printf("[FetchImage] Start fetching image: %s, proxy: %s", req.Image, req.Proxy)

		// Parse image name to extract repo and tag
		parts := strings.SplitN(req.Image, ":", 2)
		repo := parts[0]
		tag := "latest"
		if len(parts) == 2 {
			tag = parts[1]
		}

		// Pull image
		manifest, err := pullImage(c.Context(), db, store, repo, tag, req.Proxy)
		if err != nil {
			log.Printf("[FetchImage] Failed to pull image: %v", err)
			return resp.R500(c, "PULL_FAILED", nil, fmt.Sprintf("拉取镜像失败: %v", err))
		}

		log.Printf("[FetchImage] Successfully pulled image: %s:%s", repo, tag)

		return resp.R200(c, map[string]interface{}{
			"message":    "镜像拉取成功",
			"image":      req.Image,
			"repository": repo,
			"tag":        tag,
			"digest":     manifest.Config.Digest.String(),
		})
	}
}

func pullImage(ctx context.Context, db *gorm.DB, store store.Store, repo string, tag string, proxy string) (*v1.Manifest, error) {
	if repo == "" || tag == "" {
		return nil, fmt.Errorf("invalid repo or tag")
	}

	log.Printf("[PullImage] Pulling %s:%s with proxy: %s", repo, tag, proxy)

	var (
		err       error
		transport = &http.Transport{
			MaxIdleConns:        10,
			IdleConnTimeout:     30 * time.Second,
			TLSHandshakeTimeout: 30 * time.Second,
			ResponseHeaderTimeout: 30 * time.Second,
		}
		options   []remote.Option
		tn        name.Tag
		des       *remote.Descriptor
		img       v1.Image
		manifest  *v1.Manifest
		target    = fmt.Sprintf("%s:%s", repo, tag)
	)

	// Setup proxy if provided
	if proxy != "" {
		var pu *url.URL
		if pu, err = url.Parse(proxy); err != nil {
			return nil, fmt.Errorf("invalid proxy URL: %w", err)
		}

		// Handle socks5 proxy
		if pu.Scheme == "socks5" {
			log.Printf("[PullImage] Using SOCKS5 proxy: %s", proxy)
			
			// Create SOCKS5 dialer
			dialer, err := goproxy.SOCKS5("tcp", pu.Host, nil, goproxy.Direct)
			if err != nil {
				return nil, fmt.Errorf("failed to create SOCKS5 dialer: %w", err)
			}
			
			// Set custom DialContext for SOCKS5
			transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
				return dialer.Dial(network, addr)
			}
		} else {
			// HTTP/HTTPS proxy
			log.Printf("[PullImage] Using HTTP(S) proxy: %s", proxy)
			transport.Proxy = http.ProxyURL(pu)
		}
	}

	options = append(options, remote.WithTransport(transport))
	options = append(options, remote.WithContext(ctx))

	// Parse image reference
	if tn, err = name.NewTag(target); err != nil {
		return nil, fmt.Errorf("invalid image tag: %w", err)
	}

	log.Printf("[PullImage] Fetching image descriptor for %s", target)

	// Get image descriptor with timeout
	done := make(chan error, 1)
	go func() {
		var e error
		des, e = remote.Get(tn, options...)
		done <- e
	}()

	select {
	case err = <-done:
		if err != nil {
			return nil, fmt.Errorf("failed to get image (network timeout or connection error, try using a proxy): %w", err)
		}
	case <-time.After(60 * time.Second):
		return nil, fmt.Errorf("timeout fetching image descriptor (60s), the registry may be unreachable from your location, try using a proxy")
	}

	// Get image
	if img, err = des.Image(); err != nil {
		return nil, fmt.Errorf("failed to get image from descriptor: %w", err)
	}

	// Get manifest
	if manifest, err = img.Manifest(); err != nil {
		return nil, fmt.Errorf("failed to get manifest: %w", err)
	}

	log.Printf("[PullImage] Got manifest with %d layers", len(manifest.Layers))

	// Create repository
	if err := db.FirstOrCreate(&model.Repository{}, model.Repository{Name: repo}).Error; err != nil {
		return nil, fmt.Errorf("failed to create repository: %w", err)
	}

	if err := store.CreatePartition(ctx, "registry"); err != nil {
		return nil, fmt.Errorf("failed to create partition: %w", err)
	}

	// Pull config blob
	log.Printf("[PullImage] Pulling config blob: %s", manifest.Config.Digest)
	var (
		configLayer  v1.Layer
		configDigest v1.Hash
		configReader io.ReadCloser
	)

	if configLayer, err = img.LayerByDigest(manifest.Config.Digest); err != nil {
		return nil, fmt.Errorf("failed to get config layer: %w", err)
	}

	if configDigest, err = configLayer.Digest(); err != nil {
		return nil, fmt.Errorf("failed to get config digest: %w", err)
	}

	if configReader, err = configLayer.Uncompressed(); err != nil {
		return nil, fmt.Errorf("failed to get config reader: %w", err)
	}
	defer configReader.Close()

	digest := fmt.Sprintf("%s:%s", configDigest.Algorithm, configDigest.Hex)
	if err = store.WriteBlob(ctx, digest, configReader); err != nil {
		return nil, fmt.Errorf("failed to write config blob: %w", err)
	}

	if err := db.Create(&model.Blob{
		Digest:     digest,
		Size:       manifest.Config.Size,
		MediaType:  "application/vnd.docker.container.image.v1+json",
		Repository: repo,
	}).Error; err != nil && !strings.Contains(err.Error(), "UNIQUE constraint failed") {
		return nil, fmt.Errorf("failed to save config blob metadata: %w", err)
	}

	log.Printf("[PullImage] Config blob saved: %s", digest)

	// Pull layer blobs
	for idx, layerDesc := range manifest.Layers {
		log.Printf("[PullImage] Pulling layer %d/%d: %s", idx+1, len(manifest.Layers), layerDesc.Digest)

		var (
			layer       v1.Layer
			layerReader io.ReadCloser
		)

		if layer, err = img.LayerByDigest(layerDesc.Digest); err != nil {
			return nil, fmt.Errorf("failed to get layer %d: %w", idx, err)
		}

		if layerReader, err = layer.Compressed(); err != nil {
			return nil, fmt.Errorf("failed to get layer reader %d: %w", idx, err)
		}
		defer layerReader.Close()

		layerDigest := fmt.Sprintf("%s:%s", layerDesc.Digest.Algorithm, layerDesc.Digest.Hex)
		if err = store.WriteBlob(ctx, layerDigest, layerReader); err != nil {
			return nil, fmt.Errorf("failed to write layer blob %d: %w", idx, err)
		}

		if err := db.Create(&model.Blob{
			Digest:     layerDigest,
			Size:       layerDesc.Size,
			MediaType:  string(layerDesc.MediaType),
			Repository: repo,
		}).Error; err != nil && !strings.Contains(err.Error(), "UNIQUE constraint failed") {
			return nil, fmt.Errorf("failed to save layer blob metadata %d: %w", idx, err)
		}

		log.Printf("[PullImage] Layer %d saved: %s (%d bytes)", idx+1, layerDigest, layerDesc.Size)
	}

	// Convert manifest to Docker v2 format and save
	manifestData := map[string]interface{}{
		"schemaVersion": 2,
		"mediaType":     "application/vnd.docker.distribution.manifest.v2+json",
		"config": map[string]interface{}{
			"mediaType": "application/vnd.docker.container.image.v1+json",
			"size":      manifest.Config.Size,
			"digest":    manifest.Config.Digest.String(),
		},
		"layers": []map[string]interface{}{},
	}

	layers := []map[string]interface{}{}
	for _, layer := range manifest.Layers {
		layers = append(layers, map[string]interface{}{
			"mediaType": "application/vnd.docker.image.rootfs.diff.tar.gzip",
			"size":      layer.Size,
			"digest":    layer.Digest.String(),
		})
	}
	manifestData["layers"] = layers

	manifestJSON, err := json.Marshal(manifestData)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal manifest: %w", err)
	}

	manifestHash := sha256.Sum256(manifestJSON)
	manifestDigest := "sha256:" + hex.EncodeToString(manifestHash[:])

	if err := store.WriteManifest(ctx, manifestDigest, manifestJSON); err != nil {
		return nil, fmt.Errorf("failed to write manifest: %w", err)
	}

	if err := db.Create(&model.Manifest{
		Repository: repo,
		Tag:        tag,
		Digest:     manifestDigest,
		MediaType:  "application/vnd.docker.distribution.manifest.v2+json",
		Size:       int64(len(manifestJSON)),
		Content:    manifestJSON,
	}).Error; err != nil && !strings.Contains(err.Error(), "UNIQUE constraint failed") {
		return nil, fmt.Errorf("failed to save manifest: %w", err)
	}

	if err := db.Create(&model.Tag{
		Repository: repo,
		Tag:        tag,
		Digest:     manifestDigest,
	}).Error; err != nil && !strings.Contains(err.Error(), "UNIQUE constraint failed") {
		return nil, fmt.Errorf("failed to save tag: %w", err)
	}

	log.Printf("[PullImage] Manifest saved: %s", manifestDigest)

	return manifest, nil
}