From da1d9672f6deea7d81e7dcf985acc9a5c76844bd Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcg.dev>
Date: Tue, 28 May 2024 13:50:06 -0700
Subject: [PATCH] Enable imgcrypt in cri pull

Signed-off-by: Derek McGowan <derek@mcg.dev>
---
 internal/cri/server/images/image_pull.go      | 22 ++++----
 internal/cri/server/images/image_pull_test.go | 55 ++++++++++---------
 2 files changed, 40 insertions(+), 37 deletions(-)

diff --git a/internal/cri/server/images/image_pull.go b/internal/cri/server/images/image_pull.go
index 4d4ba2036..488136b43 100644
--- a/internal/cri/server/images/image_pull.go
+++ b/internal/cri/server/images/image_pull.go
@@ -32,6 +32,9 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/containerd/errdefs"
+	"github.com/containerd/imgcrypt"
+	"github.com/containerd/imgcrypt/images/encryption"
 	"github.com/containerd/log"
 	distribution "github.com/distribution/reference"
 	imagedigest "github.com/opencontainers/go-digest"
@@ -48,7 +51,6 @@ import (
 	crilabels "github.com/containerd/containerd/v2/internal/cri/labels"
 	snpkg "github.com/containerd/containerd/v2/pkg/snapshotters"
 	"github.com/containerd/containerd/v2/pkg/tracing"
-	"github.com/containerd/errdefs"
 )
 
 // For image management:
@@ -583,15 +585,15 @@ func newTransport() *http.Transport {
 // encryptedImagesPullOpts returns the necessary list of pull options required
 // for decryption of encrypted images based on the cri decryption configuration.
 // Temporarily removed for v2 upgrade
-//func (c *CRIImageService) encryptedImagesPullOpts() []containerd.RemoteOpt {
-//	if c.config.ImageDecryption.KeyModel == criconfig.KeyModelNode {
-//		ltdd := imgcrypt.Payload{}
-//		decUnpackOpt := encryption.WithUnpackConfigApplyOpts(encryption.WithDecryptedUnpack(&ltdd))
-//		opt := containerd.WithUnpackOpts([]containerd.UnpackOpt{decUnpackOpt})
-//		return []containerd.RemoteOpt{opt}
-//	}
-//	return nil
-//}
+func (c *CRIImageService) encryptedImagesPullOpts() []containerd.RemoteOpt {
+	if c.config.ImageDecryption.KeyModel == criconfig.KeyModelNode {
+		ltdd := imgcrypt.Payload{}
+		decUnpackOpt := encryption.WithUnpackConfigApplyOpts(encryption.WithDecryptedUnpack(&ltdd))
+		opt := containerd.WithUnpackOpts([]containerd.UnpackOpt{decUnpackOpt})
+		return []containerd.RemoteOpt{opt}
+	}
+	return nil
+}
 
 const (
 	// defaultPullProgressReportInterval represents that how often the
diff --git a/internal/cri/server/images/image_pull_test.go b/internal/cri/server/images/image_pull_test.go
index a01e74710..1e6321f37 100644
--- a/internal/cri/server/images/image_pull_test.go
+++ b/internal/cri/server/images/image_pull_test.go
@@ -26,10 +26,11 @@ import (
 	"github.com/stretchr/testify/assert"
 	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
 
+	"github.com/containerd/platforms"
+
 	"github.com/containerd/containerd/v2/internal/cri/annotations"
 	criconfig "github.com/containerd/containerd/v2/internal/cri/config"
 	"github.com/containerd/containerd/v2/internal/cri/labels"
-	"github.com/containerd/platforms"
 )
 
 func TestParseAuth(t *testing.T) {
@@ -350,32 +351,32 @@ func TestDefaultScheme(t *testing.T) {
 }
 
 // Temporarily remove for v2 upgrade
-//func TestEncryptedImagePullOpts(t *testing.T) {
-//	for _, test := range []struct {
-//		desc         string
-//		keyModel     string
-//		expectedOpts int
-//	}{
-//		{
-//			desc:         "node key model should return one unpack opt",
-//			keyModel:     criconfig.KeyModelNode,
-//			expectedOpts: 1,
-//		},
-//		{
-//			desc:         "no key model selected should default to node key model",
-//			keyModel:     "",
-//			expectedOpts: 0,
-//		},
-//	} {
-//		test := test
-//		t.Run(test.desc, func(t *testing.T) {
-//			c, _ := newTestCRIService()
-//			c.config.ImageDecryption.KeyModel = test.keyModel
-//			got := len(c.encryptedImagesPullOpts())
-//			assert.Equal(t, test.expectedOpts, got)
-//		})
-//	}
-//}
+func TestEncryptedImagePullOpts(t *testing.T) {
+	for _, test := range []struct {
+		desc         string
+		keyModel     string
+		expectedOpts int
+	}{
+		{
+			desc:         "node key model should return one unpack opt",
+			keyModel:     criconfig.KeyModelNode,
+			expectedOpts: 1,
+		},
+		{
+			desc:         "no key model selected should default to node key model",
+			keyModel:     "",
+			expectedOpts: 0,
+		},
+	} {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			c, _ := newTestCRIService()
+			c.config.ImageDecryption.KeyModel = test.keyModel
+			got := len(c.encryptedImagesPullOpts())
+			assert.Equal(t, test.expectedOpts, got)
+		})
+	}
+}
 
 func TestSnapshotterFromPodSandboxConfig(t *testing.T) {
 	defaultSnapshotter := "native"