feat: 添加了 imager 工具 package

refactor: 将 images 的获取分散到各个组件里面
2025-12-29 23:01:42 +08:00
parent c53c15fa8c
commit 0bcb138fd5
20 changed files with 1519 additions and 56 deletions
--- a/pkg/imager/pull.go
+++ b/pkg/imager/pull.go
@@ -0,0 +1,648 @@
+package imager
+
+import (
+	"archive/tar"
+	"compress/gzip"
+	"context"
+	"crypto/tls"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+
+	"gitea.loveuer.com/yizhisec/pkg3/logger"
+)
+
+// PullOption is a functional option for configuring image pull
+type PullOption func(*pullOption)
+
+type pullOption struct {
+	Proxy         string        // http or socks5 proxy
+	Rename        string        // pull image and rename it
+	PlainHTTP     bool          // use http instead of https
+	SkipTLSVerify bool          // skip TLS certificate verification
+	Username      string        // registry username for authentication
+	Password      string        // registry password for authentication
+	RetryTimes    int           // retry times for sync/proxy registries, default is 3
+	RetryDelay    time.Duration // delay between retries, default is 2 seconds
+}
+
+// WithProxy sets the proxy for pulling images
+func WithProxy(proxy string) PullOption {
+	return func(o *pullOption) {
+		o.Proxy = proxy
+	}
+}
+
+// WithRename sets a new name for the pulled image
+func WithRename(name string) PullOption {
+	return func(o *pullOption) {
+		o.Rename = name
+	}
+}
+
+// WithPlainHTTP enables plain HTTP (no TLS)
+func WithPlainHTTP() PullOption {
+	return func(o *pullOption) {
+		o.PlainHTTP = true
+	}
+}
+
+// WithSkipTLSVerify skips TLS certificate verification
+func WithSkipTLSVerify() PullOption {
+	return func(o *pullOption) {
+		o.SkipTLSVerify = true
+	}
+}
+
+// WithAuth sets authentication credentials
+func WithAuth(username, password string) PullOption {
+	return func(o *pullOption) {
+		o.Username = username
+		o.Password = password
+	}
+}
+
+// WithRetry sets the retry times for pulling images
+func WithRetry(times int) PullOption {
+	return func(o *pullOption) {
+		if times < 0 {
+			times = 0
+		}
+		o.RetryTimes = times
+	}
+}
+
+// WithRetryDelay sets the delay between retries
+func WithRetryDelay(delay time.Duration) PullOption {
+	return func(o *pullOption) {
+		o.RetryDelay = delay
+	}
+}
+
+// manifestV2 represents Docker manifest v2 schema
+type manifestV2 struct {
+	SchemaVersion int    `json:"schemaVersion"`
+	MediaType     string `json:"mediaType"`
+	Config        struct {
+		MediaType string `json:"mediaType"`
+		Size      int64  `json:"size"`
+		Digest    string `json:"digest"`
+	} `json:"config"`
+	Layers []struct {
+		MediaType string `json:"mediaType"`
+		Size      int64  `json:"size"`
+		Digest    string `json:"digest"`
+	} `json:"layers"`
+}
+
+// imageReference parses image reference into registry, repository, and tag
+type imageReference struct {
+	Registry   string
+	Repository string
+	Tag        string
+}
+
+// parseImageReference parses an image name into its components
+func parseImageReference(name string) (*imageReference, error) {
+	ref := &imageReference{
+		Tag: "latest",
+	}
+
+	// First, split by / to separate registry from repository
+	// This handles cases like localhost:5000/myimage:tag
+	slashParts := strings.SplitN(name, "/", 2)
+
+	var registryPart, repoPart string
+
+	if len(slashParts) == 1 {
+		// No slash, it's just image:tag (e.g., alpine:3.19)
+		repoPart = slashParts[0]
+		ref.Registry = "registry-1.docker.io"
+	} else {
+		// Has slash, check if first part is a registry
+		// Registry contains . or : (for host:port)
+		if strings.Contains(slashParts[0], ".") || strings.Contains(slashParts[0], ":") {
+			// First part is a registry
+			registryPart = slashParts[0]
+			repoPart = slashParts[1]
+			ref.Registry = registryPart
+		} else {
+			// First part is namespace (e.g., library/alpine)
+			repoPart = name
+			ref.Registry = "registry-1.docker.io"
+		}
+	}
+
+	// Now parse the repository and tag from repoPart
+	// Split on last : to get tag (in case repository name contains /)
+	lastColon := strings.LastIndex(repoPart, ":")
+	if lastColon > 0 {
+		ref.Repository = repoPart[:lastColon]
+		ref.Tag = repoPart[lastColon+1:]
+	} else {
+		ref.Repository = repoPart
+	}
+
+	// For docker.io, add library/ prefix if no namespace
+	if ref.Registry == "registry-1.docker.io" && !strings.Contains(ref.Repository, "/") {
+		ref.Repository = "library/" + ref.Repository
+	}
+
+	return ref, nil
+}
+
+// PullImage pulls an OCI image using HTTP requests and saves it as a tar archive
+func PullImage(ctx context.Context, name string, store io.Writer, opts ...PullOption) error {
+	options := &pullOption{
+		RetryTimes: 3,               // default retry times
+		RetryDelay: 2 * time.Second, // default retry delay
+	}
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	logger.Debug("Pulling image: %s with options: %+v", name, options)
+
+	// Parse image reference
+	ref, err := parseImageReference(name)
+	if err != nil {
+		return fmt.Errorf("failed to parse image reference: %w", err)
+	}
+
+	logger.Debug("Parsed image reference: registry=%s, repository=%s, tag=%s", ref.Registry, ref.Repository, ref.Tag)
+
+	// Create HTTP client
+	client := createHTTPClient(options)
+
+	// Retry logic for sync/proxy registries
+	var lastErr error
+	for attempt := 0; attempt <= options.RetryTimes; attempt++ {
+		if attempt > 0 {
+			logger.Debug("Retry attempt %d/%d for image: %s", attempt, options.RetryTimes, name)
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-time.After(options.RetryDelay):
+				// Continue after delay
+			}
+		}
+
+		// Get auth token if needed
+		token, err := getAuthToken(ctx, client, ref, options)
+		if err != nil {
+			lastErr = fmt.Errorf("failed to get auth token: %w", err)
+			logger.Debug("Attempt %d failed: %v", attempt+1, lastErr)
+			continue
+		}
+
+		// Get manifest
+		manifest, err := getManifest(ctx, client, ref, token, options)
+		if err != nil {
+			lastErr = fmt.Errorf("failed to get manifest: %w", err)
+			logger.Debug("Attempt %d failed: %v", attempt+1, lastErr)
+			continue
+		}
+
+		logger.Debug("Got manifest with %d layers", len(manifest.Layers))
+
+		// Download and create tar archive
+		if err := createImageTar(ctx, client, ref, manifest, token, store, name, options); err != nil {
+			lastErr = fmt.Errorf("failed to create image tar: %w", err)
+			logger.Debug("Attempt %d failed: %v", attempt+1, lastErr)
+			continue
+		}
+
+		// Success!
+		logger.Info("Successfully pulled image: %s", name)
+		return nil
+	}
+
+	// All retries exhausted
+	if options.RetryTimes > 0 {
+		return fmt.Errorf("failed to pull image after %d retries: %w", options.RetryTimes+1, lastErr)
+	}
+	return lastErr
+}
+
+// createHTTPClient creates an HTTP client with the given options
+func createHTTPClient(options *pullOption) *http.Client {
+	client := &http.Client{}
+
+	// Configure TLS
+	if options.SkipTLSVerify {
+		transport := &http.Transport{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		}
+		client.Transport = transport
+	}
+
+	// Configure proxy
+	if options.Proxy != "" {
+		proxyURL, err := url.Parse(options.Proxy)
+		if err == nil {
+			if client.Transport == nil {
+				client.Transport = &http.Transport{}
+			}
+			if transport, ok := client.Transport.(*http.Transport); ok {
+				transport.Proxy = http.ProxyURL(proxyURL)
+			}
+		}
+	}
+
+	return client
+}
+
+// getAuthToken gets authentication token from registry
+func getAuthToken(ctx context.Context, client *http.Client, ref *imageReference, options *pullOption) (string, error) {
+	// Try to access registry API v2 to get auth challenge
+	scheme := "https"
+	if options.PlainHTTP {
+		scheme = "http"
+	}
+
+	apiURL := fmt.Sprintf("%s://%s/v2/", scheme, ref.Registry)
+	req, err := http.NewRequestWithContext(ctx, "GET", apiURL, nil)
+	if err != nil {
+		return "", err
+	}
+
+	// Add basic auth if provided
+	if options.Username != "" && options.Password != "" {
+		req.SetBasicAuth(options.Username, options.Password)
+	}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	// If 200 OK, no auth needed
+	if resp.StatusCode == http.StatusOK {
+		return "", nil
+	}
+
+	// Check for WWW-Authenticate header
+	authHeader := resp.Header.Get("WWW-Authenticate")
+	if authHeader == "" {
+		return "", nil
+	}
+
+	// Parse auth challenge
+	token, err := fetchToken(ctx, client, authHeader, ref, options)
+	if err != nil {
+		return "", err
+	}
+
+	return token, nil
+}
+
+// fetchToken fetches authentication token from auth server
+func fetchToken(ctx context.Context, client *http.Client, authHeader string, ref *imageReference, options *pullOption) (string, error) {
+	// Parse WWW-Authenticate header
+	// Format: Bearer realm="https://auth.docker.io/token",service="registry.docker.io",scope="repository:library/alpine:pull"
+	if !strings.HasPrefix(authHeader, "Bearer ") {
+		return "", nil
+	}
+
+	params := make(map[string]string)
+	parts := strings.Split(authHeader[7:], ",")
+	for _, part := range parts {
+		part = strings.TrimSpace(part)
+		idx := strings.Index(part, "=")
+		if idx > 0 {
+			key := part[:idx]
+			value := strings.Trim(part[idx+1:], "\"")
+			params[key] = value
+		}
+	}
+
+	realm, ok := params["realm"]
+	if !ok {
+		return "", fmt.Errorf("no realm in auth header")
+	}
+
+	// Build token URL
+	tokenURL := realm
+	if service, ok := params["service"]; ok {
+		tokenURL += "?service=" + url.QueryEscape(service)
+	}
+	scope := fmt.Sprintf("repository:%s:pull", ref.Repository)
+	if params["scope"] != "" {
+		scope = params["scope"]
+	}
+	tokenURL += "&scope=" + url.QueryEscape(scope)
+
+	req, err := http.NewRequestWithContext(ctx, "GET", tokenURL, nil)
+	if err != nil {
+		return "", err
+	}
+
+	// Add basic auth if provided
+	if options.Username != "" && options.Password != "" {
+		req.SetBasicAuth(options.Username, options.Password)
+	}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("failed to get token: %s", resp.Status)
+	}
+
+	var tokenResp struct {
+		Token string `json:"token"`
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
+		return "", err
+	}
+
+	return tokenResp.Token, nil
+}
+
+// getManifest fetches the image manifest
+func getManifest(ctx context.Context, client *http.Client, ref *imageReference, token string, options *pullOption) (*manifestV2, error) {
+	scheme := "https"
+	if options.PlainHTTP {
+		scheme = "http"
+	}
+
+	manifestURL := fmt.Sprintf("%s://%s/v2/%s/manifests/%s", scheme, ref.Registry, ref.Repository, ref.Tag)
+	req, err := http.NewRequestWithContext(ctx, "GET", manifestURL, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	// Set accept headers for manifest v2 (add both types)
+	req.Header.Add("Accept", "application/vnd.docker.distribution.manifest.v2+json")
+	req.Header.Add("Accept", "application/vnd.oci.image.manifest.v1+json")
+	req.Header.Add("Accept", "application/vnd.docker.distribution.manifest.list.v2+json")
+
+	if token != "" {
+		req.Header.Set("Authorization", "Bearer "+token)
+	}
+
+	logger.Debug("Fetching manifest from: %s", manifestURL)
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		logger.Debug("Manifest request failed: %s, body: %s", resp.Status, string(body))
+		return nil, fmt.Errorf("failed to get manifest: %s", resp.Status)
+	}
+
+	contentType := resp.Header.Get("Content-Type")
+	logger.Debug("Manifest Content-Type: %s", contentType)
+
+	// Read response body
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	// Check if it's a manifest list (multi-arch)
+	if strings.Contains(contentType, "manifest.list") || strings.Contains(contentType, "image.index") {
+		// Parse manifest list and select amd64/linux
+		var manifestList struct {
+			SchemaVersion int `json:"schemaVersion"`
+			Manifests     []struct {
+				Digest   string `json:"digest"`
+				Platform struct {
+					Architecture string `json:"architecture"`
+					OS           string `json:"os"`
+				} `json:"platform"`
+			} `json:"manifests"`
+		}
+		if err := json.Unmarshal(body, &manifestList); err != nil {
+			return nil, fmt.Errorf("failed to parse manifest list: %w", err)
+		}
+
+		logger.Debug("Found manifest list with %d manifests", len(manifestList.Manifests))
+
+		// Find amd64/linux manifest
+		var targetDigest string
+		for _, m := range manifestList.Manifests {
+			logger.Debug("Checking manifest: arch=%s, os=%s, digest=%s",
+				m.Platform.Architecture, m.Platform.OS, m.Digest)
+			if m.Platform.Architecture == "amd64" && m.Platform.OS == "linux" {
+				targetDigest = m.Digest
+				break
+			}
+		}
+		if targetDigest == "" && len(manifestList.Manifests) > 0 {
+			// Fallback to first manifest with known platform
+			for _, m := range manifestList.Manifests {
+				if m.Platform.Architecture != "unknown" && m.Platform.OS != "unknown" {
+					targetDigest = m.Digest
+					logger.Debug("Using fallback manifest: arch=%s, os=%s",
+						m.Platform.Architecture, m.Platform.OS)
+					break
+				}
+			}
+		}
+		if targetDigest == "" {
+			return nil, fmt.Errorf("no suitable manifest found in manifest list")
+		}
+
+		logger.Debug("Selected manifest digest: %s", targetDigest)
+		// Fetch the actual manifest by digest
+		return getManifestByDigest(ctx, client, ref, targetDigest, token, options)
+	}
+
+	var manifest manifestV2
+	if err := json.Unmarshal(body, &manifest); err != nil {
+		return nil, fmt.Errorf("failed to parse manifest: %w", err)
+	}
+
+	logger.Debug("Parsed manifest: schemaVersion=%d, config.digest=%s, layers=%d",
+		manifest.SchemaVersion, manifest.Config.Digest, len(manifest.Layers))
+
+	return &manifest, nil
+}
+
+// getManifestByDigest fetches a manifest by its digest
+func getManifestByDigest(ctx context.Context, client *http.Client, ref *imageReference, digest, token string, options *pullOption) (*manifestV2, error) {
+	scheme := "https"
+	if options.PlainHTTP {
+		scheme = "http"
+	}
+
+	manifestURL := fmt.Sprintf("%s://%s/v2/%s/manifests/%s", scheme, ref.Registry, ref.Repository, digest)
+	req, err := http.NewRequestWithContext(ctx, "GET", manifestURL, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Add("Accept", "application/vnd.docker.distribution.manifest.v2+json")
+	req.Header.Add("Accept", "application/vnd.oci.image.manifest.v1+json")
+
+	if token != "" {
+		req.Header.Set("Authorization", "Bearer "+token)
+	}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("failed to get manifest by digest: %s", resp.Status)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	var manifest manifestV2
+	if err := json.Unmarshal(body, &manifest); err != nil {
+		return nil, fmt.Errorf("failed to parse manifest: %w", err)
+	}
+
+	return &manifest, nil
+}
+
+// createImageTar downloads layers and creates a Docker-compatible tar archive
+func createImageTar(ctx context.Context, client *http.Client, ref *imageReference, manifest *manifestV2, token string, store io.Writer, imageName string, options *pullOption) error {
+	// Create tar writer
+	tw := tar.NewWriter(store)
+	defer tw.Close()
+
+	scheme := "https"
+	if options.PlainHTTP {
+		scheme = "http"
+	}
+
+	// Download config blob
+	logger.Debug("Downloading config blob: %s", manifest.Config.Digest)
+	configBlob, err := downloadBlob(ctx, client, scheme, ref.Registry, ref.Repository, manifest.Config.Digest, token)
+	if err != nil {
+		return fmt.Errorf("failed to download config: %w", err)
+	}
+
+	// Write config as json file
+	configFileName := strings.TrimPrefix(manifest.Config.Digest, "sha256:") + ".json"
+	if err := writeTarEntry(tw, configFileName, configBlob); err != nil {
+		return err
+	}
+
+	// Download and write each layer
+	var layerFiles []string
+	for i, layer := range manifest.Layers {
+		logger.Debug("Downloading layer %d/%d: %s", i+1, len(manifest.Layers), layer.Digest)
+		layerBlob, err := downloadBlob(ctx, client, scheme, ref.Registry, ref.Repository, layer.Digest, token)
+		if err != nil {
+			return fmt.Errorf("failed to download layer %s: %w", layer.Digest, err)
+		}
+
+		layerFileName := strings.TrimPrefix(layer.Digest, "sha256:") + "/layer.tar"
+		layerFiles = append(layerFiles, layerFileName)
+
+		if err := writeTarEntry(tw, layerFileName, layerBlob); err != nil {
+			return err
+		}
+	}
+
+	// Create manifest.json for Docker compatibility
+	if options.Rename != "" {
+		imageName = options.Rename
+	}
+
+	manifestJSON := []map[string]interface{}{
+		{
+			"Config":   configFileName,
+			"RepoTags": []string{imageName},
+			"Layers":   layerFiles,
+		},
+	}
+
+	manifestData, err := json.Marshal(manifestJSON)
+	if err != nil {
+		return err
+	}
+
+	if err := writeTarEntry(tw, "manifest.json", manifestData); err != nil {
+		return err
+	}
+
+	// Create repositories file (legacy)
+	repositories := make(map[string]map[string]string)
+	repositories[imageName] = map[string]string{ref.Tag: strings.TrimPrefix(manifest.Config.Digest, "sha256:")[:12]}
+	reposData, err := json.Marshal(repositories)
+	if err != nil {
+		return err
+	}
+
+	if err := writeTarEntry(tw, "repositories", reposData); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// downloadBlob downloads a blob from registry
+func downloadBlob(ctx context.Context, client *http.Client, scheme, registry, repository, digest, token string) ([]byte, error) {
+	blobURL := fmt.Sprintf("%s://%s/v2/%s/blobs/%s", scheme, registry, repository, digest)
+	req, err := http.NewRequestWithContext(ctx, "GET", blobURL, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	if token != "" {
+		req.Header.Set("Authorization", "Bearer "+token)
+	}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("failed to download blob: %s", resp.Status)
+	}
+
+	// Check if content is gzipped
+	var reader io.Reader = resp.Body
+	if resp.Header.Get("Content-Type") == "application/vnd.docker.image.rootfs.diff.tar.gzip" ||
+		strings.Contains(resp.Header.Get("Content-Encoding"), "gzip") {
+		gzReader, err := gzip.NewReader(resp.Body)
+		if err == nil {
+			defer gzReader.Close()
+			reader = gzReader
+		}
+	}
+
+	return io.ReadAll(reader)
+}
+
+// writeTarEntry writes a file entry to tar archive
+func writeTarEntry(tw *tar.Writer, name string, data []byte) error {
+	hdr := &tar.Header{
+		Name: name,
+		Mode: 0644,
+		Size: int64(len(data)),
+	}
+
+	if err := tw.WriteHeader(hdr); err != nil {
+		return err
+	}
+
+	if _, err := tw.Write(data); err != nil {
+		return err
+	}
+
+	return nil
+}