add auth header for cdn

Signed-off-by: yxxhero <[email protected]>

manager/job/preheat.go

@@ -18,12 +18,14 @@ package job
 
 import (
 	"context"
+	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"regexp"
+	"strings"
 	"time"
 
 	machineryv1tasks "github.com/RichardKnop/machinery/v1/tasks"
@@ -50,6 +52,7 @@ type PreheatType string
 const (
 	PreheatImageType PreheatType = "image"
 	PreheatFileType  PreheatType = "file"
+	timeout                      = 1 * time.Minute
 )
 
 var accessURLPattern, _ = regexp.Compile("^(.*)://(.*)/v2/(.*)/manifests/(.*)")
@@ -173,7 +176,7 @@ func (p *preheat) getLayers(ctx context.Context, image *preheatImage, preheatArg
 	if err != nil {
 		return nil, err
 	}
-	layers, err := p.parseLayers(ocispecManifest, preheatArgs.Filter, httputils.MapToHeader(preheatArgs.Headers), image)
+	layers, err := p.parseLayers(ocispecManifest, preheatArgs.URL, preheatArgs.Filter, httputils.MapToHeader(preheatArgs.Headers), image)
 	if err != nil {
 		return nil, err
 	}
@@ -207,7 +210,12 @@ func (p *preheat) getResolver(ctx context.Context, username, password string) re
 	creds := func(string) (string, string, error) {
 		return username, password, nil
 	}
-	client := &http.Client{}
+	client := &http.Client{
+		Timeout: timeout,
+		Transport: &http.Transport{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		},
+	}
 	options := docker.ResolverOptions{}
 	options.Hosts = docker.ConfigureDefaultRegistries(
 		docker.WithClient(client),
@@ -220,6 +228,63 @@ func (p *preheat) getResolver(ctx context.Context, username, password string) re
 
 }
 
+func getAuthToken(ctx context.Context, header http.Header) (string, error) {
+	ctx, span := tracer.Start(ctx, config.SpanAuthWithRegistry, trace.WithSpanKind(trace.SpanKindProducer))
+	defer span.End()
+
+	authURL := authURL(header.Values("WWW-Authenticate"))
+	if len(authURL) == 0 {
+		return "", errors.New("authURL is empty")
+	}
+
+	req, err := http.NewRequestWithContext(ctx, "GET", authURL, nil)
+	if err != nil {
+		return "", err
+	}
+
+	client := &http.Client{
+		Timeout: timeout,
+		Transport: &http.Transport{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		},
+	}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	body, _ := io.ReadAll(resp.Body)
+	var result map[string]interface{}
+	if err := json.Unmarshal(body, &result); err != nil {
+		return "", err
+	}
+
+	if result["token"] == nil {
+		return "", errors.New("token is empty")
+	}
+
+	token := fmt.Sprintf("%v", result["token"])
+	return token, nil
+
+}
+
+func authURL(wwwAuth []string) string {
+	// Bearer realm="<auth-service-url>",service="<service>",scope="repository:<name>:pull"
+	if len(wwwAuth) == 0 {
+		return ""
+	}
+	polished := make([]string, 0)
+	for _, it := range wwwAuth {
+		polished = append(polished, strings.ReplaceAll(it, "\"", ""))
+	}
+	fileds := strings.Split(polished[0], ",")
+	host := strings.Split(fileds[0], "=")[1]
+	query := strings.Join(fileds[1:], "&")
+	return fmt.Sprintf("%s?%s", host, query)
+}
+
 func (p *preheat) getManifests(ctx context.Context, image *preheatImage, username, password string) (ocispec.Manifest, error) {
 
 	resolver := p.getResolver(ctx, username, password)
@@ -283,10 +348,41 @@ func references(om ocispec.Manifest) []ocispec.Descriptor {
 	return references
 }
 
-func (p *preheat) parseLayers(om ocispec.Manifest, filter string, header http.Header, image *preheatImage) ([]*internaljob.PreheatRequest, error) {
+func (p *preheat) parseLayers(om ocispec.Manifest, url, filter string, header http.Header, image *preheatImage) ([]*internaljob.PreheatRequest, error) {
 
 	var layers []*internaljob.PreheatRequest
 
+	req, err := http.NewRequestWithContext(context.Background(), "GET", url, nil)
+	if err != nil {
+		return nil, err
+	}
+	client := &http.Client{
+		Timeout: timeout,
+		Transport: &http.Transport{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		},
+	}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode/100 != 2 && resp.StatusCode != http.StatusUnauthorized {
+		return nil, fmt.Errorf("request registry %d", resp.StatusCode)
+	}
+
+	layerHeader := header.Clone()
+	if resp.StatusCode == http.StatusUnauthorized {
+		token, err := getAuthToken(context.Background(), resp.Header)
+		if err != nil {
+			return nil, err
+		}
+
+		layerHeader.Set("Authorization", fmt.Sprintf("Bearer %s", token))
+	}
+
 	for _, v := range references(om) {
 		digest := v.Digest.String()
 		if len(digest) == 0 {
@@ -297,7 +393,7 @@ func (p *preheat) parseLayers(om ocispec.Manifest, filter string, header http.He
 			Tag:     p.bizTag,
 			Filter:  filter,
 			Digest:  digest,
-			Headers: httputils.HeaderToMap(header),
+			Headers: httputils.HeaderToMap(layerHeader),
 		}
 
 		layers = append(layers, layer)

manager/job/preheat_test.go

@@ -46,8 +46,11 @@ func TestGetLayers(t *testing.T) {
 		domain:   "registry-1.docker.io",
 		name:     "dragonflyoss/busybox",
 		tag:      "1.35.0",
-	}, types.PreheatArgs{})
+	}, types.PreheatArgs{
+		URL:  "https://registry-1.docker.io/v2/dragonflyoss/busybox/manifests/1.35.0",
+		Type: "image",
+	})
 
-	require.NotEmpty(t, ps)
 	require.NoError(t, err)
+	require.NotEmpty(t, ps)
 }