Skip to content

Commit

Permalink
feat: generate RSA key lazily for lazy refresh (#826)
Browse files Browse the repository at this point in the history
When lazy refresh is enabled, the RSA key generation will not occur
until a client calls Dial. This change helps cold start times in
serverless environments.

Prior to this commit, the RSA key generation was done at the package
level using globals. This commit removes the use of globals and moves
the key generation into the dialer. This means multiple dialers will now
have multiple RSA keys. Before this commit, multiple dialers all shared
the same RSA key. If this is a concern, callers may still configure
their own RSA key across multiple dialers.

Fixes #823
  • Loading branch information
enocom authored Jun 6, 2024
1 parent 294d77f commit bf293e2
Show file tree
Hide file tree
Showing 2 changed files with 202 additions and 28 deletions.
95 changes: 67 additions & 28 deletions dialer.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,18 +60,47 @@ var (
//go:embed version.txt
versionString string
userAgent = "cloud-sql-go-connector/" + strings.TrimSpace(versionString)

// defaultKey is the default RSA public/private keypair used by the clients.
defaultKey *rsa.PrivateKey
defaultKeyErr error
keyOnce sync.Once
)

func getDefaultKeys() (*rsa.PrivateKey, error) {
keyOnce.Do(func() {
defaultKey, defaultKeyErr = rsa.GenerateKey(rand.Reader, 2048)
})
return defaultKey, defaultKeyErr
// keyGenerator encapsulates the details of RSA key generation to provide lazy
// generation, custom keys, or a default RSA generator.
type keyGenerator struct {
once sync.Once
key *rsa.PrivateKey
err error
genFunc func() (*rsa.PrivateKey, error)
}

// newKeyGenerator initializes a keyGenerator that will (in order):
// - always return the RSA key if one is provided, or
// - generate an RSA key lazily when it's requested, or
// - (default) immediately generate an RSA key as part of the initializer.
func newKeyGenerator(
k *rsa.PrivateKey, lazy bool, genFunc func() (*rsa.PrivateKey, error),
) (*keyGenerator, error) {
g := &keyGenerator{genFunc: genFunc}
switch {
case k != nil:
// If the caller has provided a key, initialize the key and consume the
// sync.Once now.
g.once.Do(func() { g.key, g.err = k, nil })
case lazy:
// If lazy refresh is enabled, do nothing and wait for the call to
// rsaKey.
default:
// If no key has been provided and lazy refresh isn't enabled, generate
// the key and consume the sync.Once now.
g.once.Do(func() { g.key, g.err = g.genFunc() })
}
return g, g.err
}

// rsaKey will generate an RSA key if one is not already cached. Otherwise, it
// will return the cached key.
func (g *keyGenerator) rsaKey() (*rsa.PrivateKey, error) {
g.once.Do(func() { g.key, g.err = g.genFunc() })

return g.key, g.err
}

type connectionInfoCache interface {
Expand All @@ -95,7 +124,7 @@ type monitoredCache struct {
type Dialer struct {
lock sync.RWMutex
cache map[instance.ConnName]monitoredCache
key *rsa.PrivateKey
keyGenerator *keyGenerator
refreshTimeout time.Duration
// closed reports if the dialer has been closed.
closed chan struct{}
Expand Down Expand Up @@ -184,14 +213,6 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
cfg.iamLoginTokenSource = scoped
}

if cfg.rsaKey == nil {
key, err := getDefaultKeys()
if err != nil {
return nil, fmt.Errorf("failed to generate RSA keys: %v", err)
}
cfg.rsaKey = key
}

if cfg.setUniverseDomain && cfg.setAdminAPIEndpoint {
return nil, errors.New(
"can not use WithAdminAPIEndpoint and WithUniverseDomain Options together, " +
Expand Down Expand Up @@ -225,11 +246,18 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
if err := trace.InitMetrics(); err != nil {
return nil, err
}
g, err := newKeyGenerator(cfg.rsaKey, cfg.lazyRefresh,
func() (*rsa.PrivateKey, error) {
return rsa.GenerateKey(rand.Reader, 2048)
})
if err != nil {
return nil, err
}
d := &Dialer{
closed: make(chan struct{}),
cache: make(map[instance.ConnName]monitoredCache),
lazyRefresh: cfg.lazyRefresh,
key: cfg.rsaKey,
keyGenerator: g,
refreshTimeout: cfg.refreshTimeout,
sqladmin: client,
logger: cfg.logger,
Expand Down Expand Up @@ -272,7 +300,11 @@ func (d *Dialer) Dial(ctx context.Context, icn string, opts ...DialOption) (conn

var endInfo trace.EndSpanFunc
ctx, endInfo = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn/internal.InstanceInfo")
c := d.connectionInfoCache(ctx, cn, &cfg.useIAMAuthN)
c, err := d.connectionInfoCache(ctx, cn, &cfg.useIAMAuthN)
if err != nil {
endInfo(err)
return nil, err
}
ci, err := c.ConnectionInfo(ctx)
if err != nil {
d.removeCached(ctx, cn, c, err)
Expand Down Expand Up @@ -401,7 +433,10 @@ func (d *Dialer) EngineVersion(ctx context.Context, icn string) (string, error)
if err != nil {
return "", err
}
i := d.connectionInfoCache(ctx, cn, &d.defaultDialConfig.useIAMAuthN)
i, err := d.connectionInfoCache(ctx, cn, &d.defaultDialConfig.useIAMAuthN)
if err != nil {
return "", err
}
ci, err := i.ConnectionInfo(ctx)
if err != nil {
return "", err
Expand All @@ -421,8 +456,8 @@ func (d *Dialer) Warmup(ctx context.Context, icn string, opts ...DialOption) err
for _, opt := range opts {
opt(&cfg)
}
_ = d.connectionInfoCache(ctx, cn, &cfg.useIAMAuthN)
return nil
_, err = d.connectionInfoCache(ctx, cn, &cfg.useIAMAuthN)
return err
}

// newInstrumentedConn initializes an instrumentedConn that on closing will
Expand Down Expand Up @@ -475,7 +510,7 @@ func (d *Dialer) Close() error {
// modify the existing one, or leave it unchanged as needed.
func (d *Dialer) connectionInfoCache(
ctx context.Context, cn instance.ConnName, useIAMAuthN *bool,
) monitoredCache {
) (monitoredCache, error) {
d.lock.RLock()
c, ok := d.cache[cn]
d.lock.RUnlock()
Expand All @@ -490,20 +525,24 @@ func (d *Dialer) connectionInfoCache(
useIAMAuthNDial = *useIAMAuthN
}
d.logger.Debugf(ctx, "[%v] Connection info added to cache", cn.String())
k, err := d.keyGenerator.rsaKey()
if err != nil {
return monitoredCache{}, err
}
var cache connectionInfoCache
if d.lazyRefresh {
cache = cloudsql.NewLazyRefreshCache(
cn,
d.logger,
d.sqladmin, d.key,
d.sqladmin, k,
d.refreshTimeout, d.iamTokenSource,
d.dialerID, useIAMAuthNDial,
)
} else {
cache = cloudsql.NewRefreshAheadCache(
cn,
d.logger,
d.sqladmin, d.key,
d.sqladmin, k,
d.refreshTimeout, d.iamTokenSource,
d.dialerID, useIAMAuthNDial,
)
Expand All @@ -516,5 +555,5 @@ func (d *Dialer) connectionInfoCache(

c.UpdateRefresh(useIAMAuthN)

return c
return c, nil
}
135 changes: 135 additions & 0 deletions key_gen_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,135 @@
// Copyright 2024 Google LLC
//
// 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
//
// https://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 cloudsqlconn

import (
"crypto/rsa"
"errors"
"testing"
)

func TestKeyGenerator(t *testing.T) {
custom := &rsa.PrivateKey{}
generated := &rsa.PrivateKey{}

tcs := []struct {
desc string
key *rsa.PrivateKey
lazy bool
genFunc func() (*rsa.PrivateKey, error)
wantKey *rsa.PrivateKey
// whether key generation should happen in the initializer or the call
// to rsaKey
wantLazy bool
}{
{
desc: "by default a key is generated",
genFunc: func() (*rsa.PrivateKey, error) {
return generated, nil
},
wantKey: generated,
},
{
desc: "a custom key skips the generator",
key: custom,
genFunc: func() (*rsa.PrivateKey, error) {
return nil, errors.New("generator should not be called")
},
wantKey: custom,
},
{
desc: "lazy generates keys on first request",
lazy: true,
genFunc: func() (*rsa.PrivateKey, error) {
return generated, nil
},
wantKey: generated,
wantLazy: true,
},
}

for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
g, err := newKeyGenerator(tc.key, tc.lazy, tc.genFunc)
if err != nil {
t.Fatal(err)
}
if tc.wantLazy && g.key != nil {
t.Fatal("want RSA key to be lazily generated, but it wasn't")
}
k, err := g.rsaKey()
if err != nil {
t.Fatal(err)
}
if tc.wantKey != k {
t.Fatalf("want = %v, got = %v", tc.wantKey, k)
}
})
}
}

func TestKeyGeneratorErrors(t *testing.T) {
sentinel := errors.New("sentinel error")
tcs := []struct {
desc string
key *rsa.PrivateKey
lazy bool
genFunc func() (*rsa.PrivateKey, error)
wantInitError error
wantKeyError error
}{
{
desc: "generator returns errors",
genFunc: func() (*rsa.PrivateKey, error) {
return nil, sentinel
},
wantInitError: sentinel,
wantKeyError: sentinel,
},
{
desc: "custom keys never error",
key: &rsa.PrivateKey{},
genFunc: func() (*rsa.PrivateKey, error) {
return nil, errors.New("generator should not be called")
},
wantInitError: nil,
wantKeyError: nil,
},
{
desc: "lazy generation returns errors",
lazy: true,
genFunc: func() (*rsa.PrivateKey, error) {
return nil, sentinel
},
// initialization should succeed
wantInitError: nil,
// but requesting the key later should fail
wantKeyError: sentinel,
},
}

for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
g, err := newKeyGenerator(tc.key, tc.lazy, tc.genFunc)
if err != tc.wantInitError {
t.Fatal("initialization should fail, but did not")
}
_, err = g.rsaKey()
if err != tc.wantKeyError {
t.Fatal("rsaKey should fail but didn't")
}
})
}
}

0 comments on commit bf293e2

Please sign in to comment.