From 2e96ba95e0c69e47d16456d59c4ce21597e6471e Mon Sep 17 00:00:00 2001
From: Tony Fang <nhfang@amazon.com>
Date: Thu, 2 Mar 2023 02:08:08 +0000
Subject: [PATCH] Create config struct to take user input

Signed-off-by: Tony Fang <nhfang@amazon.com>
---
 pkg/transfer/local/transfer.go | 24 +++----------
 plugins/transfer/plugin.go     | 61 ++++++++++++++++++++++++++++++++--
 2 files changed, 62 insertions(+), 23 deletions(-)

diff --git a/pkg/transfer/local/transfer.go b/pkg/transfer/local/transfer.go
index 3345f1a2a..0189290dd 100644
--- a/pkg/transfer/local/transfer.go
+++ b/pkg/transfer/local/transfer.go
@@ -22,7 +22,6 @@ import (
 	"io"
 	"time"
 
-	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/images"
@@ -30,7 +29,6 @@ import (
 	"github.com/containerd/containerd/pkg/kmutex"
 	"github.com/containerd/containerd/pkg/transfer"
 	"github.com/containerd/containerd/pkg/unpack"
-	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/typeurl/v2"
 	"golang.org/x/sync/semaphore"
 )
@@ -153,10 +151,9 @@ func (ts *localTransferService) withLease(ctx context.Context, opts ...leases.Op
 
 type TransferConfig struct {
 	// MaxConcurrentDownloads is the max concurrent content downloads for pull.
-	MaxConcurrentDownloads int `toml:"max_concurrent_downloads"`
-
+	MaxConcurrentDownloads int
 	// MaxConcurrentUploadedLayers is the max concurrent uploads for push
-	MaxConcurrentUploadedLayers int `toml:"max_concurrent_uploaded_layers"`
+	MaxConcurrentUploadedLayers int
 
 	// DuplicationSuppressor is used to make sure that there is only one
 	// in-flight fetch request or unpack handler for a given descriptor's
@@ -169,21 +166,8 @@ type TransferConfig struct {
 	BaseHandlers []images.Handler
 
 	// UnpackPlatforms are used to specify supported combination of platforms and snapshotters
-	UnpackPlatforms []unpack.Platform `toml:"unpack_platforms"`
+	UnpackPlatforms []unpack.Platform
 
 	// RegistryConfigPath is a path to the root directory containing registry-specific configurations
-	RegistryConfigPath string `toml:"config_path"`
-}
-
-func DefaultConfig() *TransferConfig {
-	return &TransferConfig{
-		MaxConcurrentDownloads:      3,
-		MaxConcurrentUploadedLayers: 3,
-		UnpackPlatforms: []unpack.Platform{
-			{
-				Platform:       platforms.Only(platforms.DefaultSpec()),
-				SnapshotterKey: containerd.DefaultSnapshotter,
-			},
-		},
-	}
+	RegistryConfigPath string
 }
diff --git a/plugins/transfer/plugin.go b/plugins/transfer/plugin.go
index 72d0ece10..d5b25b896 100644
--- a/plugins/transfer/plugin.go
+++ b/plugins/transfer/plugin.go
@@ -17,9 +17,14 @@
 package transfer
 
 import (
+	"fmt"
+
+	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/leases"
 	"github.com/containerd/containerd/metadata"
 	"github.com/containerd/containerd/pkg/transfer/local"
+	"github.com/containerd/containerd/pkg/unpack"
+	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/containerd/plugin"
 
 	// Load packages with type registrations
@@ -36,9 +41,9 @@ func init() {
 			plugin.LeasePlugin,
 			plugin.MetadataPlugin,
 		},
-		Config: local.DefaultConfig(),
+		Config: defaultConfig(),
 		InitFn: func(ic *plugin.InitContext) (interface{}, error) {
-			config := ic.Config.(*local.TransferConfig)
+			config := ic.Config.(*transferConfig)
 			m, err := ic.Get(plugin.MetadataPlugin)
 			if err != nil {
 				return nil, err
@@ -49,7 +54,57 @@ func init() {
 				return nil, err
 			}
 
-			return local.NewTransferService(l.(leases.Manager), ms.ContentStore(), metadata.NewImageStore(ms), config), nil
+			// Set configuration based on default or user input
+			var lc local.TransferConfig
+			lc.MaxConcurrentDownloads = config.maxConcurrentDownloads
+			lc.MaxConcurrentUploadedLayers = config.maxConcurrentUploadedLayers
+			for _, uc := range config.unpackConfiguration {
+				p, err := platforms.Parse(uc.platform)
+				if err != nil {
+					return nil, fmt.Errorf("%s: platform configuration %v invalid", plugin.TransferPlugin, uc.platform)
+				}
+
+				up := unpack.Platform{
+					Platform:       platforms.OnlyStrict(p),
+					SnapshotterKey: uc.snapshotter,
+				}
+				lc.UnpackPlatforms = append(lc.UnpackPlatforms, up)
+			}
+			lc.RegistryConfigPath = config.registryConfigPath
+
+			return local.NewTransferService(l.(leases.Manager), ms.ContentStore(), metadata.NewImageStore(ms), &lc), nil
 		},
 	})
 }
+
+type transferConfig struct {
+	// maxConcurrentDownloads is the max concurrent content downloads for pull.
+	maxConcurrentDownloads int `toml:"max_concurrent_downloads"`
+
+	// maxConcurrentUploadedLayers is the max concurrent uploads for push
+	maxConcurrentUploadedLayers int `toml:"max_concurrent_uploaded_layers"`
+
+	// unpackConfiguration is used to read config from toml
+	unpackConfiguration []unpackConfiguration `toml:"unpack_config"`
+
+	// registryConfigPath is a path to the root directory containing registry-specific configurations
+	registryConfigPath string `toml:"config_path"`
+}
+
+type unpackConfiguration struct {
+	platform    string
+	snapshotter string
+}
+
+func defaultConfig() *transferConfig {
+	return &transferConfig{
+		maxConcurrentDownloads:      3,
+		maxConcurrentUploadedLayers: 3,
+		unpackConfiguration: []unpackConfiguration{
+			{
+				platform:    platforms.Format(platforms.DefaultSpec()),
+				snapshotter: containerd.DefaultSnapshotter,
+			},
+		},
+	}
+}