From 6a8e707cc3221a3298f05e642db428e4aa64a76c Mon Sep 17 00:00:00 2001
From: Andrew Sy Kim <kim.andrewsy@gmail.com>
Date: Thu, 22 Mar 2018 21:01:30 -0400
Subject: [PATCH] allow pod CIDR to be set using node annotations (#345)

---
 utils/node.go          |  1 -
 utils/pod_cidr.go      | 13 +++++++
 utils/pod_cidr_test.go | 80 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/utils/node.go b/utils/node.go
index 78afa4d9..785cd49e 100644
--- a/utils/node.go
+++ b/utils/node.go
@@ -13,7 +13,6 @@ import (
 
 // GetNodeObject returns the node API object for the node
 func GetNodeObject(clientset kubernetes.Interface, hostnameOverride string) (*apiv1.Node, error) {
-
 	// assuming kube-router is running as pod, first check env NODE_NAME
 	nodeName := os.Getenv("NODE_NAME")
 	if nodeName != "" {
diff --git a/utils/pod_cidr.go b/utils/pod_cidr.go
index f5c6d6f9..52c8abfb 100644
--- a/utils/pod_cidr.go
+++ b/utils/pod_cidr.go
@@ -14,6 +14,10 @@ import (
 	"k8s.io/client-go/kubernetes"
 )
 
+const (
+	podCIDRAnnotation = "kube-router.io/pod-cidr"
+)
+
 // GetPodCidrFromCniSpec gets pod CIDR allocated to the node from CNI spec file and returns it
 func GetPodCidrFromCniSpec(cniConfFilePath string) (net.IPNet, error) {
 	var podCidr net.IPNet
@@ -112,6 +116,15 @@ func GetPodCidrFromNodeSpec(clientset kubernetes.Interface, hostnameOverride str
 		return "", fmt.Errorf("Failed to get pod CIDR allocated for the node due to: " + err.Error())
 	}
 
+	if cidr, ok := node.Annotations[podCIDRAnnotation]; ok {
+		_, _, err = net.ParseCIDR(cidr)
+		if err != nil {
+			return "", fmt.Errorf("error parsing pod CIDR in node annotation: %v", err)
+		}
+
+		return cidr, nil
+	}
+
 	if node.Spec.PodCIDR == "" {
 		return "", fmt.Errorf("node.Spec.PodCIDR not set for node: %v", node.Name)
 	}
diff --git a/utils/pod_cidr_test.go b/utils/pod_cidr_test.go
index 7413076e..f59d76a2 100644
--- a/utils/pod_cidr_test.go
+++ b/utils/pod_cidr_test.go
@@ -8,6 +8,10 @@ import (
 	"os"
 	"reflect"
 	"testing"
+
+	apiv1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes/fake"
 )
 
 func Test_GetPodCidrFromCniSpec(t *testing.T) {
@@ -118,6 +122,82 @@ func Test_InsertPodCidrInCniSpec(t *testing.T) {
 	}
 }
 
+func Test_GetPodCidrFromNodeSpec(t *testing.T) {
+	testcases := []struct {
+		name             string
+		hostnameOverride string
+		existingNode     *apiv1.Node
+		podCIDR          string
+		err              error
+	}{
+		{
+			"node with node.Spec.PoodCIDR",
+			"test-node",
+			&apiv1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "test-node",
+				},
+				Spec: apiv1.NodeSpec{
+					PodCIDR: "172.17.0.0/24",
+				},
+			},
+			"172.17.0.0/24",
+			nil,
+		},
+		{
+			"node with node.Annotations['kube-router.io/pod-cidr']",
+			"test-node",
+			&apiv1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "test-node",
+					Annotations: map[string]string{
+						podCIDRAnnotation: "172.17.0.0/24",
+					},
+				},
+			},
+			"172.17.0.0/24",
+			nil,
+		},
+		{
+			"node with invalid pod cidr in node.Annotations['kube-router.io/pod-cidr']",
+			"test-node",
+			&apiv1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "test-node",
+					Annotations: map[string]string{
+						podCIDRAnnotation: "172.17.0.0",
+					},
+				},
+			},
+			"",
+			errors.New("error parsing pod CIDR in node annotation: invalid CIDR address: 172.17.0.0"),
+		},
+	}
+
+	for _, testcase := range testcases {
+		t.Run(testcase.name, func(t *testing.T) {
+			clientset := fake.NewSimpleClientset()
+			_, err := clientset.Core().Nodes().Create(testcase.existingNode)
+			if err != nil {
+				t.Fatalf("failed to create existing nodes for test: %v", err)
+			}
+
+			podCIDR, err := GetPodCidrFromNodeSpec(clientset, testcase.hostnameOverride)
+			if !reflect.DeepEqual(err, testcase.err) {
+				t.Logf("actual error: %v", err)
+				t.Logf("expected error: %v", testcase.err)
+				t.Error("did not get expected error")
+			}
+
+			if podCIDR != testcase.podCIDR {
+				t.Logf("actual podCIDR: %q", podCIDR)
+				t.Logf("expected podCIDR: %q", testcase.podCIDR)
+				t.Error("did not get expected podCIDR")
+			}
+		})
+	}
+}
+
 func createFile(content, filename string) (*os.File, error) {
 	file, err := os.Create(filename)
 	if err != nil {