gravitational/teleport
1
0
Fork 0
mirror of https://github.com/gravitational/teleport.git synced 2026-10-11 22:49:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Respect context in all process storage APIs (#63826)

Browse Source 
Replaces all context.TODOs with a caller supplied context so
that cancellation is respected when interacting with process
storage.
This commit is contained in:
rosstimothy
2026-03-02 22:05:29 +00:00
committed by GitHub
parent e7569cf382
commit 5b98cd713e
8 changed files with 32 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
integration/hsm/hsm_test.go
+3 -2
View File
@@ -157,8 +157,9 @@ func TestHSMRotation(t *testing.T) {
}))
}
func getAdminClient(authDataDir string, authAddr string) (*authclient.Client, error) {
func getAdminClient(ctx context.Context, authDataDir string, authAddr string) (*authclient.Client, error) {
identity, err := storage.ReadLocalIdentityForRole(
ctx,
filepath.Join(authDataDir, teleport.ComponentProcess),
types.RoleAdmin)
if err != nil {
@@ -182,7 +183,7 @@ func getAdminClient(authDataDir string, authAddr string) (*authclient.Client, er
func testAdminClient(t *testing.T, authDataDir string, authAddr string) {
f := func() error {
clt, err := getAdminClient(authDataDir, authAddr)
clt, err := getAdminClient(t.Context(), authDataDir, authAddr)
if err != nil {
return err
}
lib/auth/storage/storage.go
+12 -12
View File
@@ -119,7 +119,7 @@ func (p *ProcessStorage) GetState(ctx context.Context, role types.SystemRole) (*
}
// CreateState creates process state if it does not exist yet.
func (p *ProcessStorage) CreateState(role types.SystemRole, state state.StateV2) error {
func (p *ProcessStorage) CreateState(ctx context.Context, role types.SystemRole, state state.StateV2) error {
if err := state.CheckAndSetDefaults(); err != nil {
return trace.Wrap(err)
}
@@ -131,7 +131,7 @@ func (p *ProcessStorage) CreateState(role types.SystemRole, state state.StateV2)
Key: backend.NewKey(statesPrefix, strings.ToLower(role.String()), stateName),
Value: value,
}
_, err = p.stateStorage.Create(context.TODO(), item)
_, err = p.stateStorage.Create(ctx, item)
if err != nil {
return trace.Wrap(err)
}
@@ -139,7 +139,7 @@ func (p *ProcessStorage) CreateState(role types.SystemRole, state state.StateV2)
}
// WriteState writes local cluster state to the backend.
func (p *ProcessStorage) WriteState(role types.SystemRole, state state.StateV2) error {
func (p *ProcessStorage) WriteState(ctx context.Context, role types.SystemRole, state state.StateV2) error {
if err := state.CheckAndSetDefaults(); err != nil {
return trace.Wrap(err)
}
@@ -151,7 +151,7 @@ func (p *ProcessStorage) WriteState(role types.SystemRole, state state.StateV2)
Key: backend.NewKey(statesPrefix, strings.ToLower(role.String()), stateName),
Value: value,
}
_, err = p.stateStorage.Put(context.TODO(), item)
_, err = p.stateStorage.Put(ctx, item)
if err != nil {
return trace.Wrap(err)
}
@@ -159,11 +159,11 @@ func (p *ProcessStorage) WriteState(role types.SystemRole, state state.StateV2)
}
// ReadIdentity reads identity using identity name and role.
func (p *ProcessStorage) ReadIdentity(name string, role types.SystemRole) (*state.Identity, error) {
func (p *ProcessStorage) ReadIdentity(ctx context.Context, name string, role types.SystemRole) (*state.Identity, error) {
if name == "" {
return nil, trace.BadParameter("missing parameter name")
}
item, err := p.stateStorage.Get(context.TODO(), backend.NewKey(idsPrefix, strings.ToLower(role.String()), name))
item, err := p.stateStorage.Get(ctx, backend.NewKey(idsPrefix, strings.ToLower(role.String()), name))
if err != nil {
return nil, trace.Wrap(err)
}
@@ -188,7 +188,7 @@ func (p *ProcessStorage) ReadIdentity(name string, role types.SystemRole) (*stat
}
// WriteIdentity writes identity to the backend.
func (p *ProcessStorage) WriteIdentity(name string, id state.Identity) error {
func (p *ProcessStorage) WriteIdentity(ctx context.Context, name string, id state.Identity) error {
res := state.IdentityV2{
ResourceHeader: types.ResourceHeader{
Kind: types.KindIdentity,
@@ -217,7 +217,7 @@ func (p *ProcessStorage) WriteIdentity(name string, id state.Identity) error {
Key: backend.NewKey(idsPrefix, strings.ToLower(id.ID.Role.String()), name),
Value: value,
}
_, err = p.stateStorage.Put(context.TODO(), item)
_, err = p.stateStorage.Put(ctx, item)
return trace.Wrap(err)
}
@@ -297,18 +297,18 @@ func (p *ProcessStorage) ReadRDPLicense(ctx context.Context, key *types.RDPLicen
//
// TODO(nklaassen): delete this after ref has been removed from teleport.e
func ReadLocalIdentity(dataDir string, id state.IdentityID) (*state.Identity, error) {
return ReadLocalIdentityForRole(dataDir, id.Role)
return ReadLocalIdentityForRole(context.TODO(), dataDir, id.Role)
}
// ReadLocalIdentityForRole reads, parses and returns the given pub/pri key +
// cert from the key storage (dataDir).
func ReadLocalIdentityForRole(dataDir string, role types.SystemRole) (*state.Identity, error) {
storage, err := NewProcessStorage(context.TODO(), dataDir)
func ReadLocalIdentityForRole(ctx context.Context, dataDir string, role types.SystemRole) (*state.Identity, error) {
storage, err := NewProcessStorage(ctx, dataDir)
if err != nil {
return nil, trace.Wrap(err)
}
defer storage.Close()
return storage.ReadIdentity(state.IdentityCurrent, role)
return storage.ReadIdentity(ctx, state.IdentityCurrent, role)
}
// ReadOrGenerateHostID tries to read the `host_uuid` from Kubernetes storage (if available) or local storage.
lib/auth/storage/storage_test.go
+2 -2
View File
@@ -302,9 +302,9 @@ func TestImmutableLabels(t *testing.T) {
// run in a loop to ensure we can overwrite existing labels
for i := range 2 {
err = storage.WriteIdentity(identName, *ident)
err = storage.WriteIdentity(t.Context(), identName, *ident)
require.NoError(t, err)
ident, err = storage.ReadIdentity(identName, types.RoleNode)
ident, err = storage.ReadIdentity(t.Context(), identName, types.RoleNode)
require.NoError(t, err)
require.Empty(t, cmp.Diff(immutableLabels, ident.ImmutableLabels, protocmp.Transform()))
immutableLabels.Ssh[fmt.Sprintf("label-%d", i)] = fmt.Sprintf("value-%d", i)
lib/service/connect.go
+9 -9
View File
@@ -280,7 +280,7 @@ func (process *TeleportProcess) connect(role types.SystemRole, opts ...certOptio
case types.RotationPhaseUpdateClients:
// Clients should use updated credentials,
// while servers should use old credentials to answer auth requests.
newIdentity, err := process.storage.ReadIdentity(state.IdentityReplacement, role)
newIdentity, err := process.storage.ReadIdentity(process.GracefulExitContext(), state.IdentityReplacement, role)
if err != nil {
if !trace.IsNotFound(err) {
return nil, trace.Wrap(err)
@@ -418,7 +418,7 @@ func (process *TeleportProcess) healInstanceIdentity(currentIdentity *state.Iden
)
}
if err := process.storage.WriteIdentity(state.IdentityCurrent, *newIdentity); err != nil {
if err := process.storage.WriteIdentity(process.GracefulExitContext(), state.IdentityCurrent, *newIdentity); err != nil {
return nil, trace.Wrap(err, "failed to write new identity to storage")
}
return newIdentity, nil
@@ -537,11 +537,11 @@ func (process *TeleportProcess) firstTimeConnect(role types.SystemRole) (*Connec
)
}
if err := process.storage.WriteIdentity(state.IdentityCurrent, *identity); err != nil {
if err := process.storage.WriteIdentity(process.GracefulExitContext(), state.IdentityCurrent, *identity); err != nil {
process.logger.WarnContext(process.ExitContext(), "Failed to write identity to storage.", "identity", role, "error", err)
}
err = process.storage.WriteState(role, state.StateV2{
err = process.storage.WriteState(process.GracefulExitContext(), role, state.StateV2{
Spec: state.StateSpecV2{
Rotation: ca.GetRotation(),
InitialLocalVersion: teleport.Version,
@@ -794,7 +794,7 @@ func (process *TeleportProcess) syncOpenSSHRotationState() error {
return trace.Wrap(err)
}
id, err := process.storage.ReadIdentity(state.IdentityCurrent, types.RoleNode)
id, err := process.storage.ReadIdentity(process.GracefulExitContext(), state.IdentityCurrent, types.RoleNode)
if err != nil {
return trace.Wrap(err)
}
@@ -1181,12 +1181,12 @@ func (process *TeleportProcess) rotate(conn *Connector, localState state.StateV2
const outOfSync = "%v and cluster rotation state (%v) is out of sync with local (%v). Clear local state and re-register this %v."
writeStateAndIdentity := func(name string, identity *state.Identity) error {
err = storage.WriteIdentity(name, *identity)
err = storage.WriteIdentity(process.GracefulExitContext(), name, *identity)
if err != nil {
return trace.Wrap(err)
}
localState.Spec.Rotation = remote
err = storage.WriteState(id.Role, localState)
err = storage.WriteState(process.GracefulExitContext(), id.Role, localState)
process.rotationCache.Remove(id.Role)
if err != nil {
return trace.Wrap(err)
@@ -1211,7 +1211,7 @@ func (process *TeleportProcess) rotate(conn *Connector, localState state.StateV2
if err != nil {
return nil, trace.Wrap(err)
}
err = storage.WriteIdentity(state.IdentityCurrent, *identity)
err = storage.WriteIdentity(process.GracefulExitContext(), state.IdentityCurrent, *identity)
if err != nil {
return nil, trace.Wrap(err)
}
@@ -1257,7 +1257,7 @@ func (process *TeleportProcess) rotate(conn *Connector, localState state.StateV2
}
// only update local phase, there is no need to reload
localState.Spec.Rotation = remote
err = storage.WriteState(id.Role, localState)
err = storage.WriteState(process.GracefulExitContext(), id.Role, localState)
process.rotationCache.Remove(id.Role)
if err != nil {
return nil, trace.Wrap(err)
lib/service/service.go
+3 -3
View File
@@ -983,7 +983,7 @@ func (process *TeleportProcess) getIdentity(role types.SystemRole) (i *state.Ide
process.Lock()
defer process.Unlock()
i, err = process.storage.ReadIdentity(state.IdentityCurrent, role)
i, err = process.storage.ReadIdentity(process.GracefulExitContext(), state.IdentityCurrent, role)
if err == nil {
return i, nil
@@ -1018,7 +1018,7 @@ func (process *TeleportProcess) getIdentity(role types.SystemRole) (i *state.Ide
return nil, trace.Wrap(err)
}
process.logger.InfoContext(process.ExitContext(), "Found static identity, writing to disk.", "role", role)
if err = process.storage.WriteIdentity(state.IdentityCurrent, *i); err != nil {
if err = process.storage.WriteIdentity(process.GracefulExitContext(), state.IdentityCurrent, *i); err != nil {
return nil, trace.Wrap(err)
}
@@ -2214,7 +2214,7 @@ func (process *TeleportProcess) initAuthService() error {
process.backend = b
clusterName := cfg.Auth.ClusterName.GetClusterName()
ident, err := process.storage.ReadIdentity(state.IdentityCurrent, types.RoleAdmin)
ident, err := process.storage.ReadIdentity(process.GracefulExitContext(), state.IdentityCurrent, types.RoleAdmin)
if err != nil && !trace.IsNotFound(err) {
return trace.Wrap(err)
}
tool/tctl/common/admin_action_test.go
+1
View File
@@ -1180,6 +1180,7 @@ func newAdminActionTestSuite(t *testing.T) *adminActionTestSuite {
require.NoError(t, err)
localAdmin, err := storage.ReadLocalIdentityForRole(
ctx,
filepath.Join(process.Config.DataDir, teleport.ComponentProcess),
types.RoleAdmin,
)
tool/tctl/common/config/global.go
+1 -1
View File
@@ -172,7 +172,7 @@ func ApplyConfig(ccf *GlobalCLIFlags, cfg *servicecfg.Config) (*authclient.Confi
}
return nil, trace.Wrap(err)
}
identity, err := storage.ReadLocalIdentityForRole(filepath.Join(cfg.DataDir, teleport.ComponentProcess), types.RoleAdmin)
identity, err := storage.ReadLocalIdentityForRole(ctx, filepath.Join(cfg.DataDir, teleport.ComponentProcess), types.RoleAdmin)
if err != nil {
// The "admin" identity is not present? This means the tctl is running
// NOT on the auth server
tool/teleport/testenv/test_server.go
+1
View File
@@ -543,6 +543,7 @@ func (p *cliModules) SetFeatures(f modules.Features) {
func NewDefaultAuthClient(process *service.TeleportProcess) (*authclient.Client, error) {
cfg := process.Config
identity, err := storage.ReadLocalIdentityForRole(
process.GracefulExitContext(),
filepath.Join(cfg.DataDir, teleport.ComponentProcess),
types.RoleAdmin,
)