feat: root to agent upsert key rpc

Makefile

@@ -82,6 +82,7 @@ dev: postgres root
 .PHONY: proto-gen
 proto-gen:
 	./scripts/proto-gen.sh "api-specs/v1/proto/agents"
+	./scripts/proto-gen.sh "api-specs/v1/proto/agents/keys"
 	./scripts/proto-gen.sh "api-specs/v1/proto/sealer"
 	./scripts/proto-gen.sh "api-specs/v1/proto/admin"
 	./scripts/proto-gen.sh "api-specs/v1/proto/admin/keys"

api-specs/v1/proto/agents/keys/keys.proto

@@ -0,0 +1,22 @@
+syntax = "proto3";
+
+package krypton.v1.agents.keys;
+
+option go_package = "github.com/openkcm/krypton/pkg/api/v1/proto/agents/keys";
+
+service KeyService {
+  rpc UpsertKey(UpsertKeyRequest) returns (UpsertKeyResponse) {}
+}
+
+message UpsertKeyRequest {
+  string tenant_id = 1;
+  string key_id = 2;
+  string kind = 3;
+  string name = 4;
+  string parent_id = 5;
+  string lifecycle_state = 6;
+  string managed_by = 7;
+  map<string, string> labels = 8;
+}
+
+message UpsertKeyResponse {}

internal/handler/announcekey/task.go

@@ -60,7 +60,7 @@ func NewTaskHandler(keyStore store.Key) orbital.HandlerFunc {
 		}
 
 		// Idempotent re-delivery: agent already has the key.
-		if errors.Is(err, store.ErrKeyAlreadyExists) {
+		if errors.Is(err, store.ErrKeyInsertConflict) {
 			slogctx.Info(ctx, "key already announced (idempotent)", "keyID", key.ID)
 			resp.Complete()
 			return
@@ -76,7 +76,7 @@ func NewTaskHandler(keyStore store.Key) orbital.HandlerFunc {
 				return
 			case pgCodeUniqueViolation:
 				// Belt-and-suspenders: should already be caught above as
-				// ErrKeyAlreadyExists, but be defensive.
+				// ErrKeyInsertConflict, but be defensive.
 				resp.Complete()
 				return
 			}

internal/handler/announcekey/task_test.go

@@ -74,7 +74,7 @@ func TestTaskHandler_CorruptPayload_TerminalFail(t *testing.T) {
 }
 
 func TestTaskHandler_AlreadyExists_Idempotent(t *testing.T) {
-	handler := announcekey.NewTaskHandler(&taskCreateOverride{createErr: store.ErrKeyAlreadyExists})
+	handler := announcekey.NewTaskHandler(&taskCreateOverride{createErr: store.ErrKeyInsertConflict})
 
 	data := announcekey.TaskData{
 		KeyID:    uuid.New().String(),

internal/keylifecycle/keylifecycle.go

@@ -61,6 +61,12 @@ var defaultLifecycle = lifecycle{
 	},
 }
 
+// IsKnown reports whether s is one of the defined lifecycle states.
+func IsKnown(s model.KeyLifeCycleState) bool {
+	_, ok := defaultLifecycle.transitions[s]
+	return ok
+}
+
 // ValidateTransition checks whether transitioning from one state to another is allowed.
 func ValidateTransition(from, to model.KeyLifeCycleState) error {
 	ts, ok := defaultLifecycle.transitions[from]

internal/keylifecycle/keylifecycle_test.go

@@ -206,3 +206,28 @@ func TestKeyLifecycleKeyUsages(t *testing.T) {
 		}
 	})
 }
+
+func TestIsKnown(t *testing.T) {
+	t.Parallel()
+
+	tts := []struct {
+		state model.KeyLifeCycleState
+		want  bool
+	}{
+		{state: "", want: false},
+		{state: "bogus", want: false},
+		{state: model.KeyLifeCyclePreActivation, want: true},
+		{state: model.KeyLifeCycleActive, want: true},
+		{state: model.KeyLifeCycleSuspended, want: true},
+		{state: model.KeyLifeCycleDeactivated, want: true},
+		{state: model.KeyLifeCycleCompromised, want: true},
+		{state: model.KeyLifeCycleDestroyed, want: true},
+	}
+
+	for _, tt := range tts {
+		t.Run(fmt.Sprintf("[%s]=%t", tt.state, tt.want), func(t *testing.T) {
+			t.Parallel()
+			assert.Equal(t, tt.want, keylifecycle.IsKnown(tt.state))
+		})
+	}
+}

internal/keyoperator/key.go

@@ -34,8 +34,54 @@ var (
 
 	// ErrGetKey signals a failed read of the target key.
 	ErrGetKey = errors.New("failed to get key")
+
+	// ErrCreateKey signals a failed key creation.
+	ErrCreateKey = errors.New("failed to create key")
+
+	// ErrKeyConflict signals that an existing key with the same identity
+	// (tenant + id or tenant + name) does not match the upsert request:
+	// one of Name, TenantID, ManagedBy, Kind, or ParentID differs.
+	ErrKeyConflict = errors.New("existing key does not match upsert request")
 )
 
+// UpsertKey inserts newKey; on conflict it reconciles the existing row.
+func UpsertKey(newKey model.Key) store.TransactionFunc {
+	return func(ctx context.Context, stores store.Stores) error {
+		err := stores.Keys.CreateKey(ctx, newKey)
+		if err == nil {
+			return nil
+		}
+		if !errors.Is(err, store.ErrKeyInsertConflict) {
+			return fmt.Errorf("%w: %w", ErrCreateKey, err)
+		}
+
+		existing, err := stores.Keys.GetKeyByID(ctx, newKey.ID, newKey.TenantID)
+		if err != nil {
+			return fmt.Errorf("%w: %w", ErrGetKey, err)
+		}
+		if !existing.IsSame(&newKey) {
+			return ErrKeyConflict
+		}
+
+		err = stores.Keys.UpdateKeyStates(ctx, store.UpdateKeyStatesQuery{
+			ID:         existing.ID,
+			TenantID:   existing.TenantID,
+			ToState:    newKey.LifeCycleState,
+			ToStatus:   newKey.KeyProcessingState.Status,
+			FromState:  []model.KeyLifeCycleState{existing.LifeCycleState},
+			FromStatus: []model.KeyProcessingStatus{model.KeyProcessingPending, model.KeyProcessingFailed},
+		})
+		// compare-and-swap matched zero rows: row is already in the target state (idempotent replay).
+		if errors.Is(err, store.ErrKeyNotFound) {
+			return nil
+		}
+		if err != nil {
+			return fmt.Errorf("%w: %w", ErrUpdateKeyState, err)
+		}
+		return nil
+	}
+}
+
 // UpdateKeyState returns a transaction step that transitions the key's
 // life cycle and processing status.
 func UpdateKeyState(tenantID, keyID string, transition Transition) store.TransactionFunc {

internal/keyoperator/key_test.go

@@ -22,10 +22,15 @@ const (
 type stubKeyStore struct {
 	store.Key
 
+	createKey       func(ctx context.Context, key model.Key) error
 	getKeyByID      func(ctx context.Context, id, tenantID string) (*model.Key, error)
 	updateKeyStates func(ctx context.Context, q store.UpdateKeyStatesQuery) error
 }
 
+func (s *stubKeyStore) CreateKey(ctx context.Context, key model.Key) error {
+	return s.createKey(ctx, key)
+}
+
 func (s *stubKeyStore) GetKeyByID(ctx context.Context, id, tenantID string) (*model.Key, error) {
 	return s.getKeyByID(ctx, id, tenantID)
 }
@@ -94,3 +99,182 @@ func TestUpdateKeyState(t *testing.T) {
 		})
 	}
 }
+
+func TestUpsertKey(t *testing.T) {
+	errBoom := errors.New("boom")
+
+	parentID := "parent-1"
+	otherParent := "other-parent"
+
+	newKey := model.Key{
+		ID:                 testKeyID,
+		Name:               "some-name",
+		TenantID:           testTenantID,
+		Kind:               "K1",
+		ParentID:           &parentID,
+		ManagedBy:          "agent-aws",
+		LifeCycleState:     model.KeyLifeCyclePreActivation,
+		KeyProcessingState: model.KeyProcessingState{Status: model.KeyProcessingCompleted},
+	}
+
+	tests := []struct {
+		name string
+
+		createErr error
+		existing  *model.Key
+		getKeyErr error
+		updateErr error
+
+		wantErrIs    []error
+		wantErrIsNot []error
+		wantNil      bool
+	}{
+		{
+			name:         "generic store error on insert",
+			createErr:    errBoom,
+			wantErrIs:    []error{keyoperator.ErrCreateKey, errBoom},
+			wantErrIsNot: []error{keyoperator.ErrKeyConflict, store.ErrKeyInsertConflict},
+		},
+		{
+			name:      "conflict then GetKeyByID fails",
+			createErr: store.ErrKeyInsertConflict,
+			getKeyErr: errBoom,
+			wantErrIs: []error{keyoperator.ErrGetKey, errBoom},
+		},
+		{
+			name:      "conflict with different name",
+			createErr: store.ErrKeyInsertConflict,
+			existing: &model.Key{
+				ID:                 testKeyID,
+				Name:               "different-name",
+				TenantID:           testTenantID,
+				Kind:               "K1",
+				ParentID:           &parentID,
+				ManagedBy:          "agent-aws",
+				LifeCycleState:     model.KeyLifeCyclePreActivation,
+				KeyProcessingState: model.KeyProcessingState{Status: model.KeyProcessingPending},
+			},
+			wantErrIs: []error{keyoperator.ErrKeyConflict},
+		},
+		{
+			name:      "conflict with different managed_by",
+			createErr: store.ErrKeyInsertConflict,
+			existing: &model.Key{
+				ID:                 testKeyID,
+				Name:               "some-name",
+				TenantID:           testTenantID,
+				Kind:               "K1",
+				ParentID:           &parentID,
+				ManagedBy:          "other-agent",
+				LifeCycleState:     model.KeyLifeCyclePreActivation,
+				KeyProcessingState: model.KeyProcessingState{Status: model.KeyProcessingPending},
+			},
+			wantErrIs: []error{keyoperator.ErrKeyConflict},
+		},
+		{
+			name:      "conflict with different parent_id",
+			createErr: store.ErrKeyInsertConflict,
+			existing: &model.Key{
+				ID:                 testKeyID,
+				Name:               "some-name",
+				TenantID:           testTenantID,
+				Kind:               "K1",
+				ParentID:           &otherParent,
+				ManagedBy:          "agent-aws",
+				LifeCycleState:     model.KeyLifeCyclePreActivation,
+				KeyProcessingState: model.KeyProcessingState{Status: model.KeyProcessingPending},
+			},
+			wantErrIs: []error{keyoperator.ErrKeyConflict},
+		},
+		{
+			name:      "conflict same identity then update store error",
+			createErr: store.ErrKeyInsertConflict,
+			existing: &model.Key{
+				ID:                 testKeyID,
+				Name:               "some-name",
+				TenantID:           testTenantID,
+				Kind:               "K1",
+				ParentID:           &parentID,
+				ManagedBy:          "agent-aws",
+				LifeCycleState:     model.KeyLifeCyclePreActivation,
+				KeyProcessingState: model.KeyProcessingState{Status: model.KeyProcessingPending},
+			},
+			updateErr:    errBoom,
+			wantErrIs:    []error{keyoperator.ErrUpdateKeyState, errBoom},
+			wantErrIsNot: []error{keyoperator.ErrKeyTransitionRejected},
+		},
+		{
+			name:      "happy path insert succeeds",
+			createErr: nil,
+			wantNil:   true,
+		},
+		{
+			name:      "conflict same identity CAS updates row",
+			createErr: store.ErrKeyInsertConflict,
+			existing: &model.Key{
+				ID:                 testKeyID,
+				Name:               "some-name",
+				TenantID:           testTenantID,
+				Kind:               "K1",
+				ParentID:           &parentID,
+				ManagedBy:          "agent-aws",
+				LifeCycleState:     model.KeyLifeCyclePreActivation,
+				KeyProcessingState: model.KeyProcessingState{Status: model.KeyProcessingPending},
+			},
+			updateErr: nil,
+			wantNil:   true,
+		},
+		{
+			name:      "conflict same identity idempotent replay (CAS matched zero rows)",
+			createErr: store.ErrKeyInsertConflict,
+			existing: &model.Key{
+				ID:                 testKeyID,
+				Name:               "some-name",
+				TenantID:           testTenantID,
+				Kind:               "K1",
+				ParentID:           &parentID,
+				ManagedBy:          "agent-aws",
+				LifeCycleState:     model.KeyLifeCyclePreActivation,
+				KeyProcessingState: model.KeyProcessingState{Status: model.KeyProcessingCompleted},
+			},
+			updateErr: store.ErrKeyNotFound,
+			wantNil:   true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			keys := &stubKeyStore{
+				createKey: func(_ context.Context, _ model.Key) error {
+					return tc.createErr
+				},
+				getKeyByID: func(_ context.Context, _, _ string) (*model.Key, error) {
+					if tc.getKeyErr != nil {
+						return nil, tc.getKeyErr
+					}
+					return tc.existing, nil
+				},
+				updateKeyStates: func(_ context.Context, _ store.UpdateKeyStatesQuery) error {
+					return tc.updateErr
+				},
+			}
+
+			step := keyoperator.UpsertKey(newKey)
+			err := step(t.Context(), store.Stores{Keys: keys})
+
+			if tc.wantNil {
+				assert.NoError(t, err)
+				return
+			}
+			if !assert.Error(t, err) {
+				return
+			}
+			for _, s := range tc.wantErrIs {
+				assert.ErrorIs(t, err, s)
+			}
+			for _, s := range tc.wantErrIsNot {
+				assert.NotErrorIs(t, err, s, "unexpected: err matches %v", s)
+			}
+		})
+	}
+}

pkg/api/v1/proto/admin/keys/errmap.go

@@ -74,7 +74,7 @@ func mapToProtoErr(err error) error {
 		errors.Is(err, keyoperator.ErrGetKey),
 		errors.Is(err, keyoperator.ErrGetParentKeyVersion):
 		return proto.ErrDetailsWithCode(
-			status.New(codes.Internal, err.Error()),
+			status.New(codes.Internal, "internal error"),
 			proto.Code_ERROR_CODE_RETRY,
 		)
 	}

pkg/api/v1/proto/agents/keys/errmap.go

@@ -0,0 +1,33 @@
+package keys
+
+import (
+	"errors"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/openkcm/krypton/internal/keyoperator"
+	"github.com/openkcm/krypton/pkg/api/v1/proto"
+	"github.com/openkcm/krypton/pkg/store"
+	"github.com/openkcm/krypton/pkg/validator"
+)
+
+func mapToProtoErr(err error) error {
+	switch {
+	case errors.Is(err, store.ErrTenantNotFound):
+		return proto.ErrDetailsWithCode(
+			status.New(codes.FailedPrecondition, validator.ErrInvalidTenantID.Error()),
+			proto.Code_ERROR_CODE_ABORT,
+		)
+	case errors.Is(err, keyoperator.ErrKeyConflict):
+		return proto.ErrDetailsWithCode(
+			status.New(codes.FailedPrecondition, keyoperator.ErrKeyConflict.Error()),
+			proto.Code_ERROR_CODE_ABORT,
+		)
+	}
+
+	return proto.ErrDetailsWithCode(
+		status.New(codes.Internal, "internal error"),
+		proto.Code_ERROR_CODE_RETRY,
+	)
+}