From af2b0bdb05cf61bae8e2a914bc12cf30b440f328 Mon Sep 17 00:00:00 2001
From: Wei Zhang <zhangweikop@gmail.com>
Date: Thu, 6 Jun 2024 03:07:40 +0000
Subject: [PATCH] add test

---
 pkg/kubelet/certificate/kubelet_test.go | 161 ++++++++++++++++++++++++
 1 file changed, 161 insertions(+)

diff --git a/pkg/kubelet/certificate/kubelet_test.go b/pkg/kubelet/certificate/kubelet_test.go
index 755f07a6798..709bf02c582 100644
--- a/pkg/kubelet/certificate/kubelet_test.go
+++ b/pkg/kubelet/certificate/kubelet_test.go
@@ -17,11 +17,19 @@ limitations under the License.
 package certificate
 
 import (
+	"bytes"
+	"context"
+	"fmt"
 	"net"
+	"os"
+	"path/filepath"
 	"reflect"
 	"testing"
+	"time"
 
 	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/client-go/util/cert"
 	netutils "k8s.io/utils/net"
 )
 
@@ -100,3 +108,156 @@ func TestAddressesToHostnamesAndIPs(t *testing.T) {
 		})
 	}
 }
+
+func removeThenCreate(name string, data []byte, perm os.FileMode) error {
+	if err := os.Remove(name); err != nil {
+		if !os.IsNotExist(err) {
+			return err
+		}
+	}
+	return os.WriteFile(name, data, perm)
+}
+
+func createCertAndKeyFiles(certDir string) (string, string, error) {
+	cert, key, err := cert.GenerateSelfSignedCertKey("k8s.io", nil, nil)
+	if err != nil {
+		return "", "", nil
+	}
+
+	certPath := filepath.Join(certDir, "kubelet.cert")
+	keyPath := filepath.Join(certDir, "kubelet.key")
+	if err := removeThenCreate(certPath, cert, os.FileMode(0644)); err != nil {
+		return "", "", err
+	}
+
+	if err := removeThenCreate(keyPath, key, os.FileMode(0600)); err != nil {
+		return "", "", err
+	}
+
+	return certPath, keyPath, nil
+}
+
+// createCertAndKeyFilesUsingRename creates cert and key files under a parent dir `identity` as
+// <certDir>/identity/kubelet.cert, <certDir>/identity/kubelet.key
+func createCertAndKeyFilesUsingRename(certDir string) (string, string, error) {
+	cert, key, err := cert.GenerateSelfSignedCertKey("k8s.io", nil, nil)
+	if err != nil {
+		return "", "", nil
+	}
+
+	var certKeyPathFn = func(dataDir string) (string, string, string) {
+		outputDir := filepath.Join(certDir, dataDir)
+		return outputDir, filepath.Join(outputDir, "kubelet.cert"), filepath.Join(outputDir, "kubelet.key")
+	}
+
+	writeDir, writeCertPath, writeKeyPath := certKeyPathFn("identity.tmp")
+	if err := os.Mkdir(writeDir, 0777); err != nil {
+		return "", "", err
+	}
+
+	if err := removeThenCreate(writeCertPath, cert, os.FileMode(0644)); err != nil {
+		return "", "", err
+	}
+
+	if err := removeThenCreate(writeKeyPath, key, os.FileMode(0600)); err != nil {
+		return "", "", err
+	}
+
+	targetDir, certPath, keyPath := certKeyPathFn("identity")
+	if err := os.RemoveAll(targetDir); err != nil {
+		if !os.IsNotExist(err) {
+			return "", "", err
+		}
+	}
+	if err := os.Rename(writeDir, targetDir); err != nil {
+		return "", "", err
+	}
+
+	return certPath, keyPath, nil
+}
+
+func TestKubeletServerCertificateFromFiles(t *testing.T) {
+	// test two common ways of certificate file updates:
+	// 1. delete and write the cert and key files directly
+	// 2. create the cert and key files under a child dir and perform dir rename during update
+	tests := []struct {
+		name      string
+		useRename bool
+	}{
+		{
+			name:      "remove and create",
+			useRename: false,
+		},
+		{
+			name:      "rename cert dir",
+			useRename: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			createFn := createCertAndKeyFiles
+			if tt.useRename {
+				createFn = createCertAndKeyFilesUsingRename
+			}
+
+			certDir := t.TempDir()
+			certPath, keyPath, err := createFn(certDir)
+			if err != nil {
+				t.Fatalf("Unable to setup cert files: %v", err)
+			}
+
+			m, err := NewKubeletServerCertificateDynamicFileManager(certPath, keyPath)
+			if err != nil {
+				t.Fatalf("Unable to create certificte provider: %v", err)
+			}
+
+			m.Start()
+			defer m.Stop()
+
+			c := m.Current()
+			if c == nil {
+				t.Fatal("failed to provide valid certificate")
+			}
+			time.Sleep(100 * time.Millisecond)
+			c2 := m.Current()
+			if c2 == nil {
+				t.Fatal("failed to provide valid certificate")
+			}
+			if c2 != c {
+				t.Errorf("expected the same loaded certificate object when there is no cert file change, got different")
+			}
+
+			// simulate certificate files updated in the background
+			if _, _, err := createFn(certDir); err != nil {
+				t.Fatalf("got errors when rotating certificate files in the test: %v", err)
+			}
+
+			err = wait.PollUntilContextTimeout(context.Background(),
+				100*time.Millisecond, 10*time.Second, true,
+				func(_ context.Context) (bool, error) {
+					c3 := m.Current()
+					if c3 == nil {
+						return false, fmt.Errorf("expected valid certificate regardless of file changes, but got nil")
+					}
+					if bytes.Equal(c.Certificate[0], c3.Certificate[0]) {
+						t.Logf("loaded certificate is not updated")
+						return false, nil
+					}
+					return true, nil
+				})
+			if err != nil {
+				t.Errorf("failed to provide the updated certificate after file changes: %v", err)
+			}
+
+			if err = os.Remove(certPath); err != nil {
+				t.Errorf("could not delete file in order to perform test")
+			}
+
+			time.Sleep(1 * time.Second)
+			if m.Current() == nil {
+				t.Errorf("expected the manager still provides cached content when certificate file was not available")
+			}
+		})
+	}
+}