From 48fbaa91fb46246226266e516604ae81c55142bf Mon Sep 17 00:00:00 2001
From: Harsh Thakur <harshthakur9030@gmail.com>
Date: Thu, 7 Sep 2023 23:32:24 +0530
Subject: [PATCH] Add push flag to support pushing to registry directly

Signed-off-by: Harsh Thakur <harshthakur9030@gmail.com>
---
 pkg/buildkit/buildkit.go | 44 ++++++++++++++++++++++++++++++++++++++++
 pkg/patch/cmd.go         |  5 ++++-
 pkg/patch/patch.go       | 16 ++++++++++-----
 3 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/pkg/buildkit/buildkit.go b/pkg/buildkit/buildkit.go
index 8fbe2355..51321412 100644
--- a/pkg/buildkit/buildkit.go
+++ b/pkg/buildkit/buildkit.go
@@ -225,3 +225,47 @@ func SolveToDocker(ctx context.Context, c *client.Client, st *llb.State, configD
 	})
 	return eg.Wait()
 }
+
+func SolveToRegistry(ctx context.Context, c *client.Client, st *llb.State, configData []byte, tag string) error {
+	def, err := st.Marshal(ctx)
+	if err != nil {
+		log.Errorf("st.Marshal failed with %s", err)
+		return err
+	}
+
+	dockerConfig := config.LoadDefaultConfigFile(os.Stderr)
+	attachable := []session.Attachable{authprovider.NewDockerAuthProvider(dockerConfig)}
+	solveOpt := client.SolveOpt{
+		Exports: []client.ExportEntry{
+			{
+				Type: client.ExporterImage,
+				Attrs: map[string]string{
+					"name": tag,
+					"push": "true",
+					// Pass through resolved configData from original image
+					exptypes.ExporterImageConfigKey: string(configData),
+				},
+			},
+		},
+		Frontend: "",         // i.e. we are passing in the llb.Definition directly
+		Session:  attachable, // used for authprovider, sshagentprovider and secretprovider
+	}
+
+	ch := make(chan *client.SolveStatus)
+	eg, ctx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		_, err := c.Solve(ctx, def, solveOpt, ch)
+		return err
+	})
+	eg.Go(func() error {
+		var c console.Console
+		if cn, err := console.ConsoleFromFile(os.Stderr); err == nil {
+			c = cn
+		}
+		// not using shared context to not disrupt display but let us finish reporting errors
+		_, err = progressui.DisplaySolveStatus(context.TODO(), c, os.Stdout, ch)
+		return err
+	})
+
+	return eg.Wait()
+}
diff --git a/pkg/patch/cmd.go b/pkg/patch/cmd.go
index e56b6ff7..c8a3f62b 100644
--- a/pkg/patch/cmd.go
+++ b/pkg/patch/cmd.go
@@ -30,6 +30,7 @@ type patchArgs struct {
 	ignoreError   bool
 	format        string
 	output        string
+	pushDest      string
 }
 
 func NewPatchCmd() *cobra.Command {
@@ -48,7 +49,8 @@ func NewPatchCmd() *cobra.Command {
 				ua.workingFolder,
 				ua.format,
 				ua.output,
-				ua.ignoreError)
+				ua.ignoreError,
+				ua.pushDest)
 		},
 	}
 	flags := patchCmd.Flags()
@@ -61,6 +63,7 @@ func NewPatchCmd() *cobra.Command {
 	flags.BoolVar(&ua.ignoreError, "ignore-errors", false, "Ignore errors and continue patching")
 	flags.StringVarP(&ua.format, "format", "f", "openvex", "Output format, defaults to 'openvex'")
 	flags.StringVarP(&ua.output, "output", "o", "", "Output file path")
+	flags.StringVarP(&ua.pushDest, "push", "p", "", "Push patched image to destination registry. Format: <registry>/<image>:<tag>. Note: this takes precedence over tag flag if set. ")
 
 	if err := patchCmd.MarkFlagRequired("image"); err != nil {
 		panic(err)
diff --git a/pkg/patch/patch.go b/pkg/patch/patch.go
index 2ae07b68..7741f779 100644
--- a/pkg/patch/patch.go
+++ b/pkg/patch/patch.go
@@ -29,13 +29,13 @@ const (
 )
 
 // Patch command applies package updates to an OCI image given a vulnerability report.
-func Patch(ctx context.Context, timeout time.Duration, buildkitAddr, image, reportFile, patchedTag, workingFolder, format, output string, ignoreError bool) error {
+func Patch(ctx context.Context, timeout time.Duration, buildkitAddr, image, reportFile, patchedTag, workingFolder, format, output string, ignoreError bool, pushDest string) error {
 	timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
 	ch := make(chan error)
 	go func() {
-		ch <- patchWithContext(timeoutCtx, buildkitAddr, image, reportFile, patchedTag, workingFolder, format, output, ignoreError)
+		ch <- patchWithContext(timeoutCtx, buildkitAddr, image, reportFile, patchedTag, workingFolder, format, output, ignoreError, pushDest)
 	}()
 
 	select {
@@ -60,7 +60,7 @@ func removeIfNotDebug(workingFolder string) {
 	}
 }
 
-func patchWithContext(ctx context.Context, buildkitAddr, image, reportFile, patchedTag, workingFolder, format, output string, ignoreError bool) error {
+func patchWithContext(ctx context.Context, buildkitAddr, image, reportFile, patchedTag, workingFolder, format, output string, ignoreError bool, pushDest string) error {
 	imageName, err := ref.ParseNamed(image)
 	if err != nil {
 		return err
@@ -138,8 +138,14 @@ func patchWithContext(ctx context.Context, buildkitAddr, image, reportFile, patc
 		return err
 	}
 
-	if err = buildkit.SolveToDocker(ctx, config.Client, patchedImageState, config.ConfigData, patchedImageName); err != nil {
-		return err
+	if pushDest != "" {
+		if err = buildkit.SolveToRegistry(ctx, config.Client, patchedImageState, config.ConfigData, pushDest); err != nil {
+			return err
+		}
+	} else {
+		if err = buildkit.SolveToDocker(ctx, config.Client, patchedImageState, config.ConfigData, patchedImageName); err != nil {
+			return err
+		}
 	}
 
 	// create a new manifest with the successfully patched packages