diff --git a/cmd/haulerctl/app/copy.go b/cmd/haulerctl/app/copy.go index 73fad58..8c841fc 100644 --- a/cmd/haulerctl/app/copy.go +++ b/cmd/haulerctl/app/copy.go @@ -3,14 +3,16 @@ package app import ( "context" + "github.com/oras-project/oras-go/pkg/content" "github.com/rancherfederal/hauler/pkg/copy" "github.com/spf13/cobra" ) type copyOpts struct { - dir string - mediatype string - src string + dir string + allowedMediaTypes []string + allowAllMediaTypes bool + sourceRef string } // NewCopyCommand creates a new sub command under @@ -24,25 +26,35 @@ func NewCopyCommand() *cobra.Command { Aliases: []string{"c", "cp"}, //Args: cobra.MinimumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - return opts.Run() + opts.sourceRef = args[0] + return opts.Run(opts.sourceRef) }, } f := cmd.Flags() + f.StringArrayVarP(&opts.allowedMediaTypes, "media-type", "t", nil, "allowed media types to be pulled") + f.BoolVarP(&opts.allowAllMediaTypes, "allow-all", "a", false, "allow all media types to be pulled") f.StringVarP(&opts.dir, "dir", "d", ".", "Target directory for file copy") - f.StringVarP(&opts.src, "registry", "r", "localhost:5000/file:test", "URI for object in the registry") return cmd } // Run performs the operation. -func (o *copyOpts) Run() error { +func (o *copyOpts) Run(src string) error { ctx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() - cp := copy.NewCopier(o.dir, o.mediatype) + store := content.NewFileStore(o.dir) + defer store.Close() - if err := cp.Get(ctx, o.src); err != nil { + if o.allowAllMediaTypes { + o.allowedMediaTypes = nil + } else if len(o.allowedMediaTypes) == 0 { + o.allowedMediaTypes = []string{content.DefaultBlobMediaType, content.DefaultBlobDirMediaType} + } + + cp := copy.NewCopier(o.dir, o.allowedMediaTypes, store) + if err := cp.Get(ctx, src); err != nil { return err } diff --git a/pkg/copy/copy.go b/pkg/copy/copy.go index 1e840ef..b6153cb 100644 --- a/pkg/copy/copy.go +++ b/pkg/copy/copy.go @@ -3,8 +3,6 @@ package copy import ( "context" "fmt" - "os" - "time" "github.com/containerd/containerd/remotes/docker" "github.com/oras-project/oras-go/pkg/content" @@ -15,51 +13,32 @@ import ( type Copier struct { Dir string fileStore *content.FileStore - mediaType string + mediaOpts []string logger *logrus.Entry } -func NewCopier(path string, media string) *Copier { - - fs := createFileStoreIfNotExists(path) +func NewCopier(dir string, media []string, fileStore *content.FileStore) *Copier { return &Copier{ - Dir: path, - fileStore: *fs, - mediaType: media, + Dir: dir, + fileStore: fileStore, + mediaOpts: media, logger: logrus.WithFields(logrus.Fields{ "store": "hauler", }), } } -func createFileStoreIfNotExists(path string) *content.FileStore { - if _, err := os.Stat(path); err != nil { - if !os.IsNotExist(err) { - logrus.Error(err) - } - if err := os.MkdirAll(path, 0755); err != nil { - logrus.Error(err) - } - } - - fs := content.NewFileStore(path) - defer fs.Close() - - return fs -} - func (c Copier) Get(ctx context.Context, src string) error { - ctx, cancel := context.WithTimeout(ctx, 2*time.Minute) + ctx, cancel := context.WithCancel(context.Background()) defer cancel() - resolver := docker.NewResolver(docker.ResolverOptions{}) + resolver := docker.NewResolver(docker.ResolverOptions{PlainHTTP: true}) // Pull file(s) from registry and save to disk fmt.Printf("Pulling from %s and saving to %s\n", src, c.Dir) - allowedMediaTypes := []string{(c.mediaType)} - desc, _, err := oras.Pull(ctx, resolver, src, &c.fileStore, oras.WithAllowedMediaTypes(allowedMediaTypes)) + desc, _, err := oras.Pull(ctx, resolver, src, c.fileStore, oras.WithAllowedMediaTypes(c.mediaOpts)) if err != nil { return err