From 5b595b829a7d02cc15453f4be1c713f44f73045e Mon Sep 17 00:00:00 2001
From: Xiujuan Xiang <xiujuanx@vmware.com>
Date: Tue, 13 Dec 2022 15:50:22 +0800
Subject: [PATCH] forbid calico mgmt creation

---
 tkg/test/tkgctl/docker/docker_suite_test.go |  2 +
 tkg/tkgctl/init.go                          |  6 ++
 tkg/tkgctl/init_test.go                     | 68 +++++++++++++++++++++
 3 files changed, 76 insertions(+)
 create mode 100644 tkg/tkgctl/init_test.go

diff --git a/tkg/test/tkgctl/docker/docker_suite_test.go b/tkg/test/tkgctl/docker/docker_suite_test.go
index f5c31f7a2b..3148b18394 100644
--- a/tkg/test/tkgctl/docker/docker_suite_test.go
+++ b/tkg/test/tkgctl/docker/docker_suite_test.go
@@ -99,6 +99,7 @@ var _ = SynchronizedBeforeSuite(func() []byte {
 
 	// create management cluster
 	if !e2eConfig.UseExistingCluster {
+		os.Setenv("_ALLOW_CALICO_ON_MANAGEMENT_CLUSTER", "true")
 		err := cli.Init(tkgctl.InitRegionOptions{
 			Plan:                        e2eConfig.ManagementClusterOptions.Plan,
 			ClusterName:                 e2eConfig.ManagementClusterName,
@@ -114,6 +115,7 @@ var _ = SynchronizedBeforeSuite(func() []byte {
 
 		Expect(err).To(BeNil())
 	}
+	os.Unsetenv("_ALLOW_CALICO_ON_MANAGEMENT_CLUSTER")
 
 	// Create initial workload cluster
 	clusterName = e2eConfig.ClusterPrefix + "wc"
diff --git a/tkg/tkgctl/init.go b/tkg/tkgctl/init.go
index dc495fb23c..5285c0f6b0 100644
--- a/tkg/tkgctl/init.go
+++ b/tkg/tkgctl/init.go
@@ -120,6 +120,12 @@ func (t *tkgctl) Init(options InitRegionOptions) error {
 	optionsIR := t.populateClientInitRegionOptions(&options, nodeSizeOptions, ceipOptIn)
 	optionsIR.IsInputFileClusterClassBased = isInputFileClusterClassBased
 
+	// Forbid calico management-cluster creation
+	// TODO: _ALLOW_CALICO_ON_MANAGEMENT_CLUSTER parameter is just used for internal debugging.
+	//  After the migration from calico to antrea management-cluster is done, it will be removed.
+	if optionsIR.CniType == "calico" && os.Getenv("_ALLOW_CALICO_ON_MANAGEMENT_CLUSTER") != "true" {
+		return errors.Errorf("Calico management-cluster creation is forbidden...")
+	}
 	// take the provided hidden flags and enable the related feature flags
 	t.tkgClient.ParseHiddenArgsAsFeatureFlags(&optionsIR)
 
diff --git a/tkg/tkgctl/init_test.go b/tkg/tkgctl/init_test.go
new file mode 100644
index 0000000000..7e5278fd7f
--- /dev/null
+++ b/tkg/tkgctl/init_test.go
@@ -0,0 +1,68 @@
+package tkgctl
+
+import (
+	"os"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+
+	"github.com/vmware-tanzu/tanzu-framework/tkg/fakes"
+)
+
+var _ = Describe("Init", func() {
+	var (
+		tkgClient              *fakes.Client
+		tkgconfigreaderwriter  *fakes.TKGConfigReaderWriter
+		initRegionOptions      InitRegionOptions
+		err                    error
+		tkgctlClient           *tkgctl
+		tkgConfigUpdaterClient *fakes.TKGConfigUpdaterClient
+		tkgBomClient           *fakes.TKGConfigBomClient
+	)
+
+	BeforeEach(func() {
+		tkgClient = &fakes.Client{}
+		tkgconfigreaderwriter = &fakes.TKGConfigReaderWriter{}
+		tkgConfigUpdaterClient = &fakes.TKGConfigUpdaterClient{}
+		tkgBomClient = &fakes.TKGConfigBomClient{}
+
+		tkgctlClient = &tkgctl{
+			tkgClient:              tkgClient,
+			tkgConfigReaderWriter:  tkgconfigreaderwriter,
+			tkgConfigUpdaterClient: tkgConfigUpdaterClient,
+			tkgBomClient:           tkgBomClient,
+		}
+		initRegionOptions = InitRegionOptions{
+			Plan:                   "dev",
+			ClusterName:            "foobar",
+			InfrastructureProvider: "FOOBAR",
+			CniType:                "calico",
+			UseExistingCluster:     true,
+			UI:                     false,
+			ClusterConfigFile:      "../fakes/config/config.yaml",
+		}
+
+	})
+
+	Context("When _ALLOW_CALICO_ON_MANAGEMENT_CLUSTER is not set", func() {
+		It("should return an error", func() {
+			err = tkgctlClient.Init(initRegionOptions)
+			Expect(err.Error()).To(Equal("Calico management-cluster creation is forbidden..."))
+		})
+	})
+
+	Context("When _ALLOW_CALICO_ON_MANAGEMENT_CLUSTER is set", func() {
+		BeforeEach(func() {
+			os.Setenv("_ALLOW_CALICO_ON_MANAGEMENT_CLUSTER", "true")
+		})
+
+		It("should succeed", func() {
+			err = tkgctlClient.Init(initRegionOptions)
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		AfterEach(func() {
+			os.Unsetenv("_ALLOW_CALICO_ON_MANAGEMENT_CLUSTER")
+		})
+	})
+})