From 93f27975f22e46ad1fa42c2caa0bab27812d2eb2 Mon Sep 17 00:00:00 2001 From: razzle Date: Fri, 3 Mar 2023 18:38:27 -0600 Subject: [PATCH] http roundtripper progress bar w/ oras Signed-off-by: Wayne Starr --- src/pkg/packager/publish.go | 12 ++++++- src/pkg/utils/oras.go | 65 ++++++++++++++++++++++++++++--------- 2 files changed, 61 insertions(+), 16 deletions(-) diff --git a/src/pkg/packager/publish.go b/src/pkg/packager/publish.go index fcdcb41db9..1eb02d9306 100644 --- a/src/pkg/packager/publish.go +++ b/src/pkg/packager/publish.go @@ -177,7 +177,13 @@ func (p *Packager) publish(ref registry.Reference, paths []string) error { copyOpts := oras.DefaultCopyOptions copyOpts.Concurrency = p.cfg.PublishOpts.CopyOptions.Concurrency - spinner.Updatef("Pushing %d layers concurrently", copyOpts.Concurrency) + spinner.Stop() + var total int64 + for _, desc := range descs { + total += desc.Size + } + dst.ProgressBar = message.NewProgressBar(total, "Publish layers") + defer dst.ProgressBar.Stop() copyOpts.OnCopySkipped = func(ctx context.Context, desc ocispec.Descriptor) error { title := desc.Annotations[ocispec.AnnotationTitle] var format string @@ -213,6 +219,7 @@ func (p *Packager) publish(ref registry.Reference, paths []string) error { // attempt to push the artifact manifest _, err = oras.Copy(ctx, store, root.Digest.String(), dst, dst.Reference.Reference, copyOpts) + dst.ProgressBar.Stop() if err == nil { message.Infof("Published: %s [%s]", ref, root.MediaType) @@ -273,11 +280,14 @@ func (p *Packager) publish(ref registry.Reference, paths []string) error { return nil, nil } + dst.ProgressBar = message.NewProgressBar(total, "Publish layers") + defer dst.ProgressBar.Stop() // attempt to push the image manifest _, err = oras.Copy(ctx, store, root.Digest.String(), dst, dst.Reference.Reference, copyOpts) if err != nil { return err } + dst.ProgressBar.Stop() message.Infof("Published: %s [%s]", ref, root.MediaType) message.Infof("Digest: %s", root.Digest) return nil diff --git a/src/pkg/utils/oras.go b/src/pkg/utils/oras.go index ca85ad1910..f675bd10ad 100644 --- a/src/pkg/utils/oras.go +++ b/src/pkg/utils/oras.go @@ -6,12 +6,13 @@ package utils import ( "context" - "crypto/tls" "errors" "fmt" + "io" "net/http" zarfconfig "github.com/defenseunicorns/zarf/src/config" + "github.com/defenseunicorns/zarf/src/pkg/message" "github.com/docker/cli/cli/config" "github.com/docker/cli/cli/config/configfile" "oras.land/oras-go/v2/registry" @@ -22,6 +23,7 @@ import ( type OrasRemote struct { *remote.Repository context.Context + *message.ProgressBar } // withScopes returns a context with the given scopes. @@ -38,7 +40,7 @@ func withScopes(ref registry.Reference) context.Context { // withAuthClient returns an auth client for the given reference. // // The credentials are pulled using Docker's default credential store. -func withAuthClient(ref registry.Reference) (*auth.Client, error) { +func (o *OrasRemote) withAuthClient(ref registry.Reference) (*auth.Client, error) { cfg, err := config.Load(config.Dir()) if err != nil { return &auth.Client{}, err @@ -68,12 +70,10 @@ func withAuthClient(ref registry.Reference) (*auth.Client, error) { } transport := http.DefaultTransport.(*http.Transport).Clone() - transport.TLSClientConfig = &tls.Config{ - InsecureSkipVerify: zarfconfig.CommonOptions.Insecure, - } + transport.TLSClientConfig.InsecureSkipVerify = zarfconfig.CommonOptions.Insecure // TODO:(@RAZZLE) https://github.com/oras-project/oras/blob/e8bc5acd9b7be47f2f9f387af6a963b14ae49eda/cmd/oras/internal/option/remote.go#L183 - return &auth.Client{ + client := &auth.Client{ Credential: auth.StaticCredential(ref.Registry, cred), Cache: auth.NewCache(), // Gitlab auth fails if ForceAttemptOAuth2 is set to true @@ -81,13 +81,48 @@ func withAuthClient(ref registry.Reference) (*auth.Client, error) { Client: &http.Client{ Transport: transport, }, - }, nil + } + + client.Client.Transport = NewTransport(client.Client.Transport, o) + + return client, nil +} + +// Transport is an http.RoundTripper that keeps track of the in-flight +// request and add hooks to report HTTP tracing events. +type Transport struct { + http.RoundTripper + orasRemote *OrasRemote +} + +func NewTransport(base http.RoundTripper, o *OrasRemote) *Transport { + return &Transport{base, o} +} + +type readCloser struct { + io.Reader + io.Closer +} + +// RoundTrip calls base roundtrip while keeping track of the current request. +func (t *Transport) RoundTrip(req *http.Request) (resp *http.Response, err error) { + if req.Body != nil && t.orasRemote.ProgressBar != nil { + tee := io.TeeReader(req.Body, t.orasRemote.ProgressBar) + teeCloser := readCloser{tee, req.Body} + req.Body = teeCloser + } + + resp, err = t.RoundTripper.RoundTrip(req) + + // do response things here + + return resp, err } // OrasRemote returns an oras remote repository client and context for the given reference. -func NewOrasRemote(ref registry.Reference) (OrasRemote, error) { - r := &OrasRemote{} - r.Context = withScopes(ref) +func NewOrasRemote(ref registry.Reference) (*OrasRemote, error) { + o := &OrasRemote{} + o.Context = withScopes(ref) // patch docker.io to registry-1.docker.io // this allows end users to use docker.io as an alias for registry-1.docker.io if ref.Registry == "docker.io" { @@ -95,14 +130,14 @@ func NewOrasRemote(ref registry.Reference) (OrasRemote, error) { } repo, err := remote.NewRepository(ref.String()) if err != nil { - return OrasRemote{}, err + return &OrasRemote{}, err } repo.PlainHTTP = zarfconfig.CommonOptions.Insecure - authClient, err := withAuthClient(ref) + authClient, err := o.withAuthClient(ref) if err != nil { - return OrasRemote{}, err + return &OrasRemote{}, err } repo.Client = authClient - r.Repository = repo - return *r, nil + o.Repository = repo + return o, nil }