From 72246d83b7b3377cb9f8c665152c45e51697452a Mon Sep 17 00:00:00 2001
From: Alan Donham <alan.donham@datadoghq.com>
Date: Tue, 3 Dec 2024 11:53:55 -0500
Subject: [PATCH] Updating go implementation for zstd, removing the defer call,
 which causes invalid memory access when the writer closes. We already have
 calls to Close(), so we really don't need to defer this.

Updating oci_image_layer rule to define outputs inside the impl, so that it can choose the correct extension for the archive.
---
 go/cmd/ocitool/createlayer_cmd.go |  2 --
 oci/image.bzl                     | 16 ++++++++++------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/go/cmd/ocitool/createlayer_cmd.go b/go/cmd/ocitool/createlayer_cmd.go
index 00a7ade..42546a5 100644
--- a/go/cmd/ocitool/createlayer_cmd.go
+++ b/go/cmd/ocitool/createlayer_cmd.go
@@ -44,13 +44,11 @@ func CreateLayerCmd(c *cli.Context) error {
     if(config.UseZstd){
         zstdWriter = zstd.NewWriter(wc)
         mediaType = ocispec.MediaTypeImageLayerZstd
-        defer zstdWriter.Close()
         tw = tar.NewWriter(zstdWriter)
     } else {
         gzipWriter = gzip.NewWriter(wc)
         gzipWriter.Name = path.Base(out.Name())
         mediaType = ocispec.MediaTypeImageLayerGzip
-        defer gzipWriter.Close()
         tw = tar.NewWriter(gzipWriter)
     }
 
diff --git a/oci/image.bzl b/oci/image.bzl
index 91fcbc0..f2e4cf5 100755
--- a/oci/image.bzl
+++ b/oci/image.bzl
@@ -26,11 +26,17 @@ def _oci_image_layer_impl(ctx):
 
     descriptor_file = ctx.actions.declare_file("{}.descriptor.json".format(ctx.label.name))
 
+    archive = None
+    if ctx.attr.zstd_compression:
+        archive = ctx.actions.declare_file(ctx.label.name + ".tar.zst")
+    else:
+        archive = ctx.actions.declare_file(ctx.label.name + ".tar.gz")
+
     ctx.actions.run(
         executable = toolchain.sdk.ocitool,
         arguments = [
                         "create-layer",
-                        "--out={}".format(ctx.outputs.layer.path),
+                        "--out={}".format(archive.path),
                         "--outd={}".format(descriptor_file.path),
                         "--dir={}".format(ctx.attr.directory),
                         "--bazel-label={}".format(ctx.label),
@@ -42,15 +48,16 @@ def _oci_image_layer_impl(ctx):
         inputs = ctx.files.files + ctx.files.file_map,
         outputs = [
             descriptor_file,
-            ctx.outputs.layer,
+            archive,
         ],
     )
 
     return [
         OCIDescriptor(
             descriptor_file = descriptor_file,
-            file = ctx.outputs.layer,
+            file = archive,
         ),
+        DefaultInfo(files = depset([archive])),
     ]
 
 oci_image_layer = rule(
@@ -77,9 +84,6 @@ oci_image_layer = rule(
         ),
     },
     toolchains = ["@com_github_datadog_rules_oci//oci:toolchain"],
-    outputs = {
-        "layer": "%{name}-layer.tar.gz",
-    },
 )
 
 def _oci_image_index_impl(ctx):