From cddd82d39ef522bbe08a709880d9f62c656ae2fe Mon Sep 17 00:00:00 2001
From: Yuki Iwai <yuki.iwai.tz@gmail.com>
Date: Sun, 6 Aug 2023 07:46:29 +0900
Subject: [PATCH] Wait for the certs to be mounted inside the container

Signed-off-by: Yuki Iwai <yuki.iwai.tz@gmail.com>
---
 pkg/certgenerator/v1beta1/generator.go      | 26 +++++++++
 pkg/certgenerator/v1beta1/generator_test.go | 60 +++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/pkg/certgenerator/v1beta1/generator.go b/pkg/certgenerator/v1beta1/generator.go
index f8c1d9520fa..73718ffe519 100644
--- a/pkg/certgenerator/v1beta1/generator.go
+++ b/pkg/certgenerator/v1beta1/generator.go
@@ -26,12 +26,15 @@ import (
 	"errors"
 	"fmt"
 	"math/big"
+	"os"
+	"path/filepath"
 	"strings"
 	"time"
 
 	admissionregistrationv1 "k8s.io/api/admissionregistration/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/wait"
 	"k8s.io/klog"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
@@ -67,11 +70,34 @@ func (c *CertGenerator) Start(ctx context.Context) error {
 	if err := c.generate(ctx); err != nil {
 		return err
 	}
+	if err := wait.ExponentialBackoffWithContext(ctx, wait.Backoff{
+		Duration: time.Second,
+		Factor:   2,
+		Jitter:   1,
+		Steps:    10,
+		Cap:      time.Minute * 5,
+	}, ensureCertMounted); err != nil {
+		return err
+	}
 	// Sending an empty data to a certsReady means it starts to register controllers to the manager.
 	c.certsReady <- struct{}{}
 	return nil
 }
 
+// ensureCertMounted ensures that the generated certs are mounted inside the container.
+func ensureCertMounted(context.Context) (bool, error) {
+	certFile := filepath.Join(consts.CertDir, serverCertName)
+	if _, err := os.Stat(certFile); err != nil {
+		return false, nil
+	}
+	keyFile := filepath.Join(consts.CertDir, serverKeyName)
+	if _, err := os.Stat(keyFile); err != nil {
+		return false, nil
+	}
+	klog.Info("Succeeded to be mounted certs inside the container.")
+	return true, nil
+}
+
 func (c *CertGenerator) NeedLeaderElection() bool {
 	return false
 }
diff --git a/pkg/certgenerator/v1beta1/generator_test.go b/pkg/certgenerator/v1beta1/generator_test.go
index 2e4f4a01105..39c886adb87 100644
--- a/pkg/certgenerator/v1beta1/generator_test.go
+++ b/pkg/certgenerator/v1beta1/generator_test.go
@@ -18,6 +18,8 @@ package certgenerator
 
 import (
 	"context"
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
@@ -31,6 +33,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
 	configv1beta1 "github.com/kubeflow/katib/pkg/apis/config/v1beta1"
+	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
 )
 
 func TestGenerate(t *testing.T) {
@@ -210,3 +213,60 @@ func buildFakeClient(kubeResources []client.Object) client.Client {
 	}
 	return fakeClientBuilder.Build()
 }
+
+func TestEnsureCertMounted(t *testing.T) {
+	tests := map[string]struct {
+		keyExist  bool
+		certExist bool
+		wantExist bool
+	}{
+		"key and cert exist": {
+			keyExist:  true,
+			certExist: true,
+			wantExist: true,
+		},
+		"key doesn't exist": {
+			keyExist:  false,
+			certExist: true,
+			wantExist: false,
+		},
+		"cert doesn't exist": {
+			keyExist:  true,
+			certExist: false,
+			wantExist: false,
+		},
+		"all files doesn't exist": {
+			keyExist:  false,
+			certExist: false,
+			wantExist: false,
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			if tc.keyExist || tc.certExist {
+				if err := os.MkdirAll(consts.CertDir, 0760); err != nil {
+					t.Fatalf("Failed to set up directory: %v", err)
+				}
+				defer func() {
+					if err := os.RemoveAll(consts.CertDir); err != nil {
+						t.Fatalf("Failed to clean up directory: %v", err)
+					}
+				}()
+			}
+			if tc.keyExist {
+				if _, err := os.Create(filepath.Join(consts.CertDir, serverKeyName)); err != nil {
+					t.Fatalf("Failed to create tls.key: %v", err)
+				}
+			}
+			if tc.certExist {
+				if _, err := os.Create(filepath.Join(consts.CertDir, serverCertName)); err != nil {
+					t.Fatalf("Failed to create tls.crt: %v", err)
+				}
+			}
+			got, _ := ensureCertMounted(context.Background())
+			if tc.wantExist != got {
+				t.Errorf("Unexpected value from ensureCertMounted: \n(want: %v, got: %v)\n", tc.wantExist, got)
+			}
+		})
+	}
+}