From 9993cde866a59e163bcfe627ce0ebb8d0b5c682d Mon Sep 17 00:00:00 2001
From: SataQiu <1527062125@qq.com>
Date: Tue, 13 Jul 2021 14:34:05 +0800
Subject: [PATCH] fix the bug that tunnel-agent/tunnel-server crashes when the
 local certificate can not be loaded correctly (#378)

Signed-off-by: SataQiu <1527062125@qq.com>
---
 pkg/yurttunnel/pki/certmanager/certmanager.go |  3 +-
 .../certmanager/store/filestore_wrapper.go    | 54 +++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 pkg/yurttunnel/pki/certmanager/store/filestore_wrapper.go

diff --git a/pkg/yurttunnel/pki/certmanager/certmanager.go b/pkg/yurttunnel/pki/certmanager/certmanager.go
index f912ee728fe..d8d13645df8 100644
--- a/pkg/yurttunnel/pki/certmanager/certmanager.go
+++ b/pkg/yurttunnel/pki/certmanager/certmanager.go
@@ -28,6 +28,7 @@ import (
 
 	"github.com/openyurtio/openyurt/pkg/projectinfo"
 	"github.com/openyurtio/openyurt/pkg/yurttunnel/constants"
+	"github.com/openyurtio/openyurt/pkg/yurttunnel/pki/certmanager/store"
 	"github.com/openyurtio/openyurt/pkg/yurttunnel/server/serveraddr"
 
 	certificates "k8s.io/api/certificates/v1beta1"
@@ -121,7 +122,7 @@ func newCertManager(
 	organizations,
 	dnsNames []string, ipAddrs []net.IP) (certificate.Manager, error) {
 	certificateStore, err :=
-		certificate.NewFileStore(componentName, certDir, certDir, "", "")
+		store.NewFileStoreWrapper(componentName, certDir, certDir, "", "")
 	if err != nil {
 		return nil, fmt.Errorf("failed to initialize the server certificate store: %v", err)
 	}
diff --git a/pkg/yurttunnel/pki/certmanager/store/filestore_wrapper.go b/pkg/yurttunnel/pki/certmanager/store/filestore_wrapper.go
new file mode 100644
index 00000000000..542fddefa77
--- /dev/null
+++ b/pkg/yurttunnel/pki/certmanager/store/filestore_wrapper.go
@@ -0,0 +1,54 @@
+/*
+Copyright 2021 The OpenYurt Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package store
+
+import (
+	"crypto/tls"
+
+	"k8s.io/client-go/util/certificate"
+	"k8s.io/klog/v2"
+)
+
+// fileStoreWrapper is a wrapper for "k8s.io/client-go/util/certificate#FileStore"
+// This wrapper increases tolerance for unexpected situations and is more robust.
+type fileStoreWrapper struct {
+	certificate.FileStore
+}
+
+// NewFileStoreWrapper returns a wrapper for "k8s.io/client-go/util/certificate#FileStore"
+// This wrapper increases tolerance for unexpected situations and is more robust.
+func NewFileStoreWrapper(pairNamePrefix, certDirectory, keyDirectory, certFile, keyFile string) (certificate.FileStore, error) {
+	fileStore, err := certificate.NewFileStore(pairNamePrefix, certDirectory, keyDirectory, certFile, keyFile)
+	if err != nil {
+		return nil, err
+	}
+	return &fileStoreWrapper{
+		FileStore: fileStore,
+	}, nil
+}
+
+func (s *fileStoreWrapper) Current() (*tls.Certificate, error) {
+	cert, err := s.FileStore.Current()
+	// If an error occurs, just return the NoCertKeyError.
+	// The cert-manager will regenerate the related certificates when it receives the NoCertKeyError.
+	if err != nil {
+		klog.Warningf("unexpected error occurred when loading the certificate: %v, will regenerate it", err)
+		noCertKeyErr := certificate.NoCertKeyError("NO_VALID_CERT")
+		return nil, &noCertKeyErr
+	}
+	return cert, nil
+}