k8s.io/kubernetes@v1.31.0-alpha.0.0.20240520171757-56147500dadc/cmd/kubeadm/app/util/pkiutil/testing/testing.go (about)

     1  /*
     2  Copyright 2021 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package testing
    18  
    19  import (
    20  	"crypto"
    21  	"fmt"
    22  	"os"
    23  	"path/filepath"
    24  	"regexp"
    25  	"runtime"
    26  	"runtime/debug"
    27  	"strings"
    28  	"sync"
    29  	"testing"
    30  
    31  	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm"
    32  	"k8s.io/kubernetes/cmd/kubeadm/app/util/pkiutil"
    33  )
    34  
    35  // RunWithPrivateKeyFixtureDirectory overrides NewPrivateKey to return private key fixtures
    36  // while executing tests in m.
    37  func RunWithPrivateKeyFixtureDirectory(m *testing.M) {
    38  	defer install()()
    39  	os.Exit(m.Run())
    40  }
    41  
    42  // Reset indicates a new test is starting and previously returned private key fixtures may be reused.
    43  func Reset() {
    44  	lock.Lock()
    45  	defer lock.Unlock()
    46  	ecdsa = 0
    47  	rsa = 0
    48  }
    49  
    50  var (
    51  	testFunction = regexp.MustCompile(`.*\.Test[^./]+(.func\d*)?$`)
    52  
    53  	lock       = sync.Mutex{}
    54  	fixtureDir = ""
    55  	lastTest   = ""
    56  	ecdsa      = 0
    57  	rsa        = 0
    58  )
    59  
    60  func install() (cleanup func()) {
    61  	lock.Lock()
    62  	defer lock.Unlock()
    63  
    64  	_, filename, _, ok := runtime.Caller(1)
    65  	if !ok {
    66  		fmt.Println("Could not determine testdata location, skipping private key fixture installation")
    67  		return func() {}
    68  	}
    69  	fixtureDir = filepath.Join(filepath.Dir(filename), "testdata")
    70  
    71  	originalNewPrivateKey := pkiutil.NewPrivateKey
    72  	pkiutil.NewPrivateKey = newPrivateKey
    73  	return func() {
    74  		pkiutil.NewPrivateKey = originalNewPrivateKey
    75  	}
    76  }
    77  
    78  func newPrivateKey(keyType kubeadmapi.EncryptionAlgorithmType) (crypto.Signer, error) {
    79  	lock.Lock()
    80  	defer lock.Unlock()
    81  
    82  	var pcs [50]uintptr
    83  	nCallers := runtime.Callers(2, pcs[:])
    84  	frames := runtime.CallersFrames(pcs[:nCallers])
    85  	thisTest := ""
    86  	for {
    87  		frame, more := frames.Next()
    88  		if strings.HasSuffix(frame.File, "_test.go") && testFunction.MatchString(frame.Function) {
    89  			thisTest = frame.Function
    90  			break
    91  		}
    92  		if !more {
    93  			break
    94  		}
    95  	}
    96  
    97  	if len(thisTest) == 0 {
    98  		fmt.Println("could not determine test for private key fixture")
    99  		debug.PrintStack()
   100  		return pkiutil.GeneratePrivateKey(keyType)
   101  	}
   102  
   103  	if thisTest != lastTest {
   104  		rsa = 0
   105  		ecdsa = 0
   106  		lastTest = thisTest
   107  	}
   108  
   109  	keyName := ""
   110  	switch keyType {
   111  	case kubeadmapi.EncryptionAlgorithmECDSAP256:
   112  		ecdsa++
   113  		keyName = fmt.Sprintf("%d.ecdsa", ecdsa)
   114  	default:
   115  		rsa++
   116  		keyName = fmt.Sprintf("%d.rsa", rsa)
   117  	}
   118  
   119  	if len(keyName) > 0 {
   120  		privKey, err := pkiutil.TryLoadKeyFromDisk(fixtureDir, keyName)
   121  		if err == nil {
   122  			return privKey, nil
   123  		}
   124  	}
   125  
   126  	fmt.Println("GeneratePrivateKey " + keyName + " for " + thisTest)
   127  
   128  	signer, err := pkiutil.GeneratePrivateKey(keyType)
   129  	if err != nil {
   130  		return signer, err
   131  	}
   132  
   133  	if len(keyName) > 0 {
   134  		pkiutil.WriteKey(fixtureDir, keyName, signer)
   135  	}
   136  
   137  	return signer, err
   138  }