Skip to content

Commit b6c6b6d

Browse files
cardyokcodexgyuho
authored
feat(session): manage KAP mTLS agent credentials (#1264)
## Summary - keep only the KAP mTLS status and credential-update session commands; kubeconfig is now fixed during Lepton bootstrap - validate each per-machine certificate, identity, endpoint, and CA fingerprint before writing an immutable credential generation - atomically switch the `current` symlink, restart the agent, and acknowledge the update only after `/readyz` succeeds - remove enable/disable state, kubeconfig mutation, pending markers, nested TLS preflight, and automatic rollback - keep certificate and private-key material redacted from session audit logs ## Rollout - the companion Lepton MR [dgxc-lepton/lepton!7510](https://gitlab-master.nvidia.com/dgxc-lepton/lepton/-/merge_requests/7510) installs the agent by default on new nodes and points kubelet at loopback from bootstrap - existing nodes are not backfilled; Lepton reconciles only Machines that report the `kap-mtls-agent` package - release this gpud change before or with the Lepton node-init rollout; failures leave new nodes blocked before kubelet registration ## Verification - `GOFLAGS=-mod=mod go test -count=1 ./pkg/kapmtls` - `GOFLAGS=-mod=mod go test -count=1 -gcflags='all=-N -l' ./pkg/session` - `GOFLAGS=-mod=mod go test -race ./pkg/kapmtls` - `GOFLAGS=-mod=mod go vet ./pkg/kapmtls ./pkg/session` The repository-wide default `pkg/session` test command requires Mockey's `-gcflags='all=-N -l'` on macOS/arm64; the changed package scope passes with that required flag. Jira: LEP-5336 --------- Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Gyuho Lee <gyuhol@nvidia.com>
1 parent c15cb49 commit b6c6b6d

23 files changed

Lines changed: 2249 additions & 29 deletions

api/v1/login.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,9 @@ type LoginResponse struct {
7070
// Token is the token used to report data from the machine.
7171
Token string `json:"token,omitempty"`
7272

73+
// MachineProof binds privileged session commands to this machine identity.
74+
MachineProof string `json:"machineProof,omitempty"`
75+
7376
// ValidationResults is the validation results for the login request.
7477
// The validation results are done by the control plane.
7578
ValidationResults []ValidationResult `json:"validationResults,omitempty"`

cmd/gpud/metadata/command.go

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,18 @@ type metadataJSONOutput struct {
3838
Updated *metadataUpdatedEntry `json:"updated,omitempty"`
3939
}
4040

41+
func maskMetadataValue(key, value string) string {
42+
switch key {
43+
case pkgmetadata.MetadataKeyToken:
44+
return pkgmetadata.MaskToken(value)
45+
case pkgmetadata.MetadataKeyMachineProof:
46+
// Machine proof authorizes privileged session commands, so reveal none of it.
47+
return "<redacted>"
48+
default:
49+
return value
50+
}
51+
}
52+
4153
func Command(cliContext *cli.Context) error {
4254
outputFormat, err := gpudcommon.ParseOutputFormat(cliContext.String("output-format"))
4355
if err != nil {
@@ -92,9 +104,7 @@ func Command(cliContext *cli.Context) error {
92104

93105
maskedMetadata := make(map[string]string, len(metadata))
94106
for k, v := range metadata {
95-
if k == pkgmetadata.MetadataKeyToken {
96-
v = pkgmetadata.MaskToken(v)
97-
}
107+
v = maskMetadataValue(k, v)
98108
maskedMetadata[k] = v
99109
if outputFormat == gpudcommon.OutputFormatPlain {
100110
fmt.Printf("%s: %s\n", k, v)
@@ -129,13 +139,14 @@ func Command(cliContext *cli.Context) error {
129139
}()
130140
log.Logger.Debugw("successfully opened state file for writing")
131141

132-
log.Logger.Debugw("setting metadata", "key", setKey, "value", setValue)
142+
maskedSetValue := maskMetadataValue(setKey, setValue)
143+
log.Logger.Debugw("setting metadata", "key", setKey, "value", maskedSetValue)
133144
if err := pkgmetadata.SetMetadata(rootCtx, dbRW, setKey, setValue); err != nil {
134145
return wrapErr("failed_to_update_metadata", fmt.Errorf("failed to update metadata: %w", err))
135146
}
136147
log.Logger.Debugw("successfully updated metadata")
137148

138-
updated = &metadataUpdatedEntry{Key: setKey, Value: setValue}
149+
updated = &metadataUpdatedEntry{Key: setKey, Value: maskedSetValue}
139150
if outputFormat == gpudcommon.OutputFormatPlain {
140151
fmt.Printf("%s successfully updated metadata\n", cmdcommon.CheckMark)
141152
}

cmd/gpud/metadata/mock_command_test.go

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -441,12 +441,14 @@ func TestCommand_JSONOutput(t *testing.T) {
441441
stateFile := filepath.Join(dataDir, "gpud.state")
442442

443443
rawToken := "nvapi-stg-1234567890abcdef"
444+
machineProof := "machine-proof-secret"
444445
evTime := time.Date(2025, 1, 2, 3, 4, 5, 0, time.UTC)
445446
evMessage := "test reboot event"
446447

447448
createMetadataDB(t, stateFile, map[string]string{
448-
pkgmetadata.MetadataKeyMachineID: "machine-1",
449-
pkgmetadata.MetadataKeyToken: rawToken,
449+
pkgmetadata.MetadataKeyMachineID: "machine-1",
450+
pkgmetadata.MetadataKeyToken: rawToken,
451+
pkgmetadata.MetadataKeyMachineProof: machineProof,
450452
})
451453
insertRebootEvent(t, stateFile, evTime, evMessage)
452454

@@ -475,13 +477,47 @@ func TestCommand_JSONOutput(t *testing.T) {
475477
require.NotNil(t, out.Metadata)
476478
assert.Equal(t, "machine-1", out.Metadata[pkgmetadata.MetadataKeyMachineID])
477479
assert.Equal(t, pkgmetadata.MaskToken(rawToken), out.Metadata[pkgmetadata.MetadataKeyToken])
480+
assert.Equal(t, "<redacted>", out.Metadata[pkgmetadata.MetadataKeyMachineProof])
481+
assert.NotContains(t, stdout, machineProof)
478482
require.Len(t, out.RebootHistory, 1)
479483
assert.Equal(t, evTime.Format(time.RFC3339), out.RebootHistory[0].TimeUTC)
480484
assert.Equal(t, evMessage, out.RebootHistory[0].Message)
481485
assert.NotContains(t, stdout, "Reboot History:")
482486
})
483487
}
484488

489+
func TestCommand_JSONOutputRedactsUpdatedMachineProof(t *testing.T) {
490+
mockey.PatchConvey("json output redacts updated machine proof", t, func() {
491+
mockey.Mock(osutil.RequireRoot).To(func() error { return nil }).Build()
492+
493+
dataDir := t.TempDir()
494+
stateFile := filepath.Join(dataDir, "gpud.state")
495+
createMetadataDB(t, stateFile, map[string]string{})
496+
497+
cliContext := newCLIContext(t, []string{
498+
"--data-dir", dataDir,
499+
"--set-key", pkgmetadata.MetadataKeyMachineProof,
500+
"--set-value", "machine-proof-secret",
501+
"--output-format", "json",
502+
})
503+
stdout, stderr := captureOutput(t, func() { require.NoError(t, Command(cliContext)) })
504+
assert.Empty(t, stderr)
505+
assert.NotContains(t, stdout, "machine-proof-secret")
506+
507+
var out metadataJSONOutput
508+
require.NoError(t, json.Unmarshal([]byte(stdout), &out))
509+
require.NotNil(t, out.Updated)
510+
assert.Equal(t, "<redacted>", out.Updated.Value)
511+
512+
dbRO, err := sqlite.Open(stateFile, sqlite.WithReadOnly(true))
513+
require.NoError(t, err)
514+
defer func() { _ = dbRO.Close() }()
515+
stored, err := pkgmetadata.ReadMetadata(context.Background(), dbRO, pkgmetadata.MetadataKeyMachineProof)
516+
require.NoError(t, err)
517+
assert.Equal(t, "machine-proof-secret", stored)
518+
})
519+
}
520+
485521
func TestCommand_JSONOutputSetKeyValue(t *testing.T) {
486522
mockey.PatchConvey("json output set key value", t, func() {
487523
mockey.Mock(osutil.RequireRoot).To(func() error {

cmd/gpud/run/mock_command_test.go

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -260,15 +260,53 @@ func TestGetSessionCredentialsOptions_AllCredentialsPresent(t *testing.T) {
260260
}).Build()
261261

262262
mockey.Mock(pkgmetadata.ReadMetadata).To(func(ctx context.Context, db *sql.DB, key string) (string, error) {
263-
if key == pkgmetadata.MetadataKeyEndpoint {
263+
switch key {
264+
case pkgmetadata.MetadataKeyEndpoint:
264265
return "https://stored.endpoint.com", nil
266+
case pkgmetadata.MetadataKeyMachineProof:
267+
return "machine-proof", nil
265268
}
266269
return "", nil
267270
}).Build()
268271

269272
opts := getSessionCredentialsOptions(true, tmpDir, "")
270273
assert.NotEmpty(t, opts)
271-
assert.Len(t, opts, 3) // token, machine ID, endpoint
274+
assert.Len(t, opts, 4) // token, machine ID, endpoint, machine proof
275+
276+
op := &config.Op{}
277+
require.NoError(t, op.ApplyOpts(opts))
278+
assert.Equal(t, "machine-proof", op.SessionMachineProof)
279+
})
280+
}
281+
282+
func TestGetSessionCredentialsOptions_MachineProofReadError(t *testing.T) {
283+
tmpDir := t.TempDir()
284+
285+
mockey.PatchConvey("getSessionCredentialsOptions machine proof read error", t, func() {
286+
createTestStateFile(t, tmpDir)
287+
288+
mockey.Mock(config.ResolveDataDir).To(func(string) (string, error) {
289+
return tmpDir, nil
290+
}).Build()
291+
mockey.Mock(pkgsqlite.Open).To(func(string, ...pkgsqlite.OpOption) (*sql.DB, error) {
292+
return &sql.DB{}, nil
293+
}).Build()
294+
mockey.Mock((*sql.DB).Close).To(func(*sql.DB) error { return nil }).Build()
295+
mockey.Mock(pkgmetadata.ReadToken).To(func(context.Context, *sql.DB) (string, error) {
296+
return "session-token", nil
297+
}).Build()
298+
mockey.Mock(pkgmetadata.ReadMachineID).To(func(context.Context, *sql.DB) (string, error) {
299+
return "machine-id-123", nil
300+
}).Build()
301+
mockey.Mock(pkgmetadata.ReadMetadata).To(func(_ context.Context, _ *sql.DB, key string) (string, error) {
302+
if key == pkgmetadata.MetadataKeyEndpoint {
303+
return "https://stored.endpoint.com", nil
304+
}
305+
return "", errors.New("machine proof unavailable")
306+
}).Build()
307+
308+
opts := getSessionCredentialsOptions(true, tmpDir, "")
309+
require.Len(t, opts, 3)
272310
})
273311
}
274312

cmd/gpud/run/mock_session_test.go

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"database/sql"
66
"errors"
77
"os"
8+
"path/filepath"
89
"testing"
910

1011
"github.com/bytedance/mockey"
@@ -38,13 +39,39 @@ func TestReadSessionCredentialsFromPersistentFile(t *testing.T) {
3839
require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyToken, "session-token"))
3940
require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyMachineID, "assigned-machine-id"))
4041
require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyEndpoint, "gpud-manager.example.com"))
42+
require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyMachineProof, "machine-proof"))
4143

4244
gotToken, gotMachineID, gotEndpoint, err := readSessionCredentialsFromPersistentFile(ctx, tmpDir)
4345
require.NoError(t, err)
4446

4547
assert.Equal(t, "session-token", gotToken)
4648
assert.Equal(t, "assigned-machine-id", gotMachineID)
4749
assert.Equal(t, "gpud-manager.example.com", gotEndpoint)
50+
51+
gotMachineProof, err := readMachineProofFromPersistentFile(ctx, tmpDir)
52+
require.NoError(t, err)
53+
assert.Equal(t, "machine-proof", gotMachineProof)
54+
}
55+
56+
func TestReadMachineProofFromPersistentFileErrors(t *testing.T) {
57+
t.Run("resolve data directory", func(t *testing.T) {
58+
parentFile := filepath.Join(t.TempDir(), "file")
59+
require.NoError(t, os.WriteFile(parentFile, []byte("data"), 0600))
60+
61+
_, err := readMachineProofFromPersistentFile(context.Background(), filepath.Join(parentFile, "child"))
62+
require.ErrorContains(t, err, "failed to resolve data dir")
63+
})
64+
65+
t.Run("open state file", func(t *testing.T) {
66+
mockey.PatchConvey("open state file error", t, func() {
67+
mockey.Mock(pkgsqlite.Open).To(func(string, ...pkgsqlite.OpOption) (*sql.DB, error) {
68+
return nil, errors.New("open failed")
69+
}).Build()
70+
71+
_, err := readMachineProofFromPersistentFile(context.Background(), t.TempDir())
72+
require.ErrorContains(t, err, "failed to open state file")
73+
})
74+
})
4875
}
4976

5077
func TestReadSessionCredentialsFromPersistentFile_StateFileNotFound(t *testing.T) {

cmd/gpud/run/session.go

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,20 @@ func readSessionCredentialsFromPersistentFile(ctx context.Context, dataDir strin
9595
return sessionToken, assignedMachineID, endpoint, nil
9696
}
9797

98+
func readMachineProofFromPersistentFile(ctx context.Context, dataDir string) (string, error) {
99+
resolvedDataDir, err := config.ResolveDataDir(dataDir)
100+
if err != nil {
101+
return "", fmt.Errorf("failed to resolve data dir: %w", err)
102+
}
103+
stateFile := config.StateFilePath(resolvedDataDir)
104+
dbRO, err := pkgsqlite.Open(stateFile, pkgsqlite.WithReadOnly(true))
105+
if err != nil {
106+
return "", fmt.Errorf("failed to open state file: %w", err)
107+
}
108+
defer func() { _ = dbRO.Close() }()
109+
return pkgmetadata.ReadMetadata(ctx, dbRO, pkgmetadata.MetadataKeyMachineProof)
110+
}
111+
98112
func getSessionCredentialsOptions(dbInMemory bool, dataDir string, controlPlaneEndpoint string) []config.OpOption {
99113
// When --db-in-memory is enabled, read session credentials from the persistent state file
100114
// and pass them to the config so the server can seed them into the in-memory database.
@@ -144,11 +158,18 @@ func getSessionCredentialsOptions(dbInMemory bool, dataDir string, controlPlaneE
144158
"machineID", assignedMachineID,
145159
"endpoint", endpoint,
146160
)
147-
return []config.OpOption{
161+
options := []config.OpOption{
148162
config.WithSessionToken(sessionToken),
149163
config.WithSessionMachineID(assignedMachineID),
150164
config.WithSessionEndpoint(endpoint),
151165
}
166+
machineProof, proofErr := readMachineProofFromPersistentFile(readCtx, dataDir)
167+
if proofErr != nil {
168+
log.Logger.Warnw("failed to read machine proof from persistent state", "error", proofErr)
169+
} else if machineProof != "" {
170+
options = append(options, config.WithSessionMachineProof(machineProof))
171+
}
172+
return options
152173
}
153174

154175
// Credentials were read but are incomplete - this may indicate a partial state

pkg/config/config.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,9 @@ type Config struct {
119119
// This allows gpud up to pass the assigned machine ID from login to gpud run.
120120
SessionMachineID string `json:"-"`
121121

122+
// SessionMachineProof is the per-machine proof returned by login.
123+
SessionMachineProof string `json:"-"`
124+
122125
// SessionEndpoint is the control plane endpoint.
123126
// Used when DBInMemory is true and session credentials are passed via CLI flags.
124127
// This allows gpud up to pass the endpoint from login to gpud run.

pkg/config/default.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,9 +66,10 @@ func DefaultConfig(ctx context.Context, opts ...OpOption) (*Config, error) {
6666

6767
ContainerdServiceActiveCommands: options.ContainerdServiceActiveCommands,
6868

69-
SessionToken: options.SessionToken,
70-
SessionMachineID: options.SessionMachineID,
71-
SessionEndpoint: options.SessionEndpoint,
69+
SessionToken: options.SessionToken,
70+
SessionMachineID: options.SessionMachineID,
71+
SessionMachineProof: options.SessionMachineProof,
72+
SessionEndpoint: options.SessionEndpoint,
7273
}
7374

7475
if cfg.State == "" {

pkg/config/options.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,9 @@ type Op struct {
4545
// this machine ID into the in-memory database.
4646
SessionMachineID string
4747

48+
// SessionMachineProof is the per-machine proof for db-in-memory mode.
49+
SessionMachineProof string
50+
4851
// SessionEndpoint is the control plane endpoint for db-in-memory mode.
4952
// When DBInMemory is true and this is set, the server will seed
5053
// this endpoint into the in-memory database.
@@ -170,6 +173,13 @@ func WithSessionMachineID(machineID string) OpOption {
170173
}
171174
}
172175

176+
// WithSessionMachineProof sets the per-machine proof for db-in-memory mode.
177+
func WithSessionMachineProof(machineProof string) OpOption {
178+
return func(op *Op) {
179+
op.SessionMachineProof = machineProof
180+
}
181+
}
182+
173183
// WithSessionEndpoint sets the control plane endpoint for db-in-memory mode.
174184
// When DBInMemory is true and this is set, the server will seed
175185
// this endpoint into the in-memory database.

pkg/config/options_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ func TestDBInMemoryWithSessionCredentials(t *testing.T) {
147147
WithDBInMemory(true),
148148
WithSessionToken("session-token-from-login-response"),
149149
WithSessionMachineID("assigned-machine-id-from-login-response"),
150+
WithSessionMachineProof("machine-proof-from-login-response"),
150151
WithSessionEndpoint("https://api.example.com"),
151152
}
152153

@@ -156,6 +157,7 @@ func TestDBInMemoryWithSessionCredentials(t *testing.T) {
156157
assert.True(t, op.DBInMemory)
157158
assert.Equal(t, "session-token-from-login-response", op.SessionToken)
158159
assert.Equal(t, "assigned-machine-id-from-login-response", op.SessionMachineID)
160+
assert.Equal(t, "machine-proof-from-login-response", op.SessionMachineProof)
159161
assert.Equal(t, "https://api.example.com", op.SessionEndpoint)
160162
}
161163

0 commit comments

Comments
 (0)