diff --git a/manager/job/preheat.go b/manager/job/preheat.go index 4f177d7e374..1ba3e917cc8 100644 --- a/manager/job/preheat.go +++ b/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="",service="",scope="repository::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) diff --git a/manager/job/preheat_test.go b/manager/job/preheat_test.go index f3d80e5d3e7..71b2a0ec340 100644 --- a/manager/job/preheat_test.go +++ b/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) }