feat: accept agent_instance identities

internal/identitymeta/identitymeta.go

@@ -5,8 +5,11 @@ const (
 	IdentityTypeMetadataKey = "x-identity-type"
 	WorkloadIDMetadataKey   = "x-workload-id"
 
-	IdentityTypeUser   = "user"
-	IdentityTypeAgent  = "agent"
-	IdentityTypeApp    = "app"
-	IdentityTypeRunner = "runner"
+	IdentityTypeUser = "user"
+	// An agent workload authenticates as the instance it runs, not as the class.
+	// IdentityTypeAgent remains for identities minted before the migration.
+	IdentityTypeAgent         = "agent"
+	IdentityTypeAgentInstance = "agent_instance"
+	IdentityTypeApp           = "app"
+	IdentityTypeRunner        = "runner"
 )

internal/reconciler/reconciler_test.go

@@ -66,84 +66,69 @@ func (m *mockReconcilerStore) DeleteExposure(ctx context.Context, id uuid.UUID)
 }
 
 type mockZitiMgmt struct {
+	// Embedded so the mock satisfies the client without restating it.
+	// Adding an RPC to the API cannot break this; calling one the test did
+	// not override panics, which is what a test wants.
+	zitimanagementv1.ZitiManagementServiceClient
+
 	deletePolicy func(ctx context.Context, req *zitimanagementv1.DeleteServicePolicyRequest) (*zitimanagementv1.DeleteServicePolicyResponse, error)
 	deleteSvc    func(ctx context.Context, req *zitimanagementv1.DeleteServiceRequest) (*zitimanagementv1.DeleteServiceResponse, error)
 }
 
-func (m *mockZitiMgmt) CreateAgentIdentity(context.Context, *zitimanagementv1.CreateAgentIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.CreateAgentIdentityResponse, error) {
-	return nil, status.Error(codes.Unimplemented, "not implemented")
-}
-
-func (m *mockZitiMgmt) CreateAppIdentity(context.Context, *zitimanagementv1.CreateAppIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.CreateAppIdentityResponse, error) {
-	return nil, status.Error(codes.Unimplemented, "not implemented")
-}
-
-func (m *mockZitiMgmt) CreateService(context.Context, *zitimanagementv1.CreateServiceRequest, ...grpc.CallOption) (*zitimanagementv1.CreateServiceResponse, error) {
-	return nil, status.Error(codes.Unimplemented, "not implemented")
+func (m *mockZitiMgmt) DeleteServicePolicy(ctx context.Context, req *zitimanagementv1.DeleteServicePolicyRequest, _ ...grpc.CallOption) (*zitimanagementv1.DeleteServicePolicyResponse, error) {
+	if m.deletePolicy == nil {
+		return nil, errors.New("not implemented")
+	}
+	return m.deletePolicy(ctx, req)
 }
 
-func (m *mockZitiMgmt) DeleteIdentity(context.Context, *zitimanagementv1.DeleteIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.DeleteIdentityResponse, error) {
-	return nil, status.Error(codes.Unimplemented, "not implemented")
+func (m *mockZitiMgmt) DeleteService(ctx context.Context, req *zitimanagementv1.DeleteServiceRequest, _ ...grpc.CallOption) (*zitimanagementv1.DeleteServiceResponse, error) {
+	if m.deleteSvc == nil {
+		return nil, errors.New("not implemented")
+	}
+	return m.deleteSvc(ctx, req)
 }
 
-func (m *mockZitiMgmt) DeleteAppIdentity(context.Context, *zitimanagementv1.DeleteAppIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.DeleteAppIdentityResponse, error) {
-	return nil, status.Error(codes.Unimplemented, "not implemented")
+type mockRunners struct {
+	getWorkload func(ctx context.Context, req *runnersv1.GetWorkloadRequest) (*runnersv1.GetWorkloadResponse, error)
 }
 
-func (m *mockZitiMgmt) CreateRunnerIdentity(context.Context, *zitimanagementv1.CreateRunnerIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.CreateRunnerIdentityResponse, error) {
+func (m *mockRunners) CreateFlavor(context.Context, *runnersv1.CreateFlavorRequest, ...grpc.CallOption) (*runnersv1.CreateFlavorResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-func (m *mockZitiMgmt) DeleteRunnerIdentity(context.Context, *zitimanagementv1.DeleteRunnerIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.DeleteRunnerIdentityResponse, error) {
+func (m *mockRunners) GetFlavor(context.Context, *runnersv1.GetFlavorRequest, ...grpc.CallOption) (*runnersv1.GetFlavorResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-func (m *mockZitiMgmt) ListManagedIdentities(context.Context, *zitimanagementv1.ListManagedIdentitiesRequest, ...grpc.CallOption) (*zitimanagementv1.ListManagedIdentitiesResponse, error) {
+func (m *mockRunners) UpdateFlavor(context.Context, *runnersv1.UpdateFlavorRequest, ...grpc.CallOption) (*runnersv1.UpdateFlavorResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-func (m *mockZitiMgmt) ResolveIdentity(context.Context, *zitimanagementv1.ResolveIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.ResolveIdentityResponse, error) {
+func (m *mockRunners) DeleteFlavor(context.Context, *runnersv1.DeleteFlavorRequest, ...grpc.CallOption) (*runnersv1.DeleteFlavorResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-func (m *mockZitiMgmt) RequestServiceIdentity(context.Context, *zitimanagementv1.RequestServiceIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.RequestServiceIdentityResponse, error) {
+func (m *mockRunners) ListFlavors(context.Context, *runnersv1.ListFlavorsRequest, ...grpc.CallOption) (*runnersv1.ListFlavorsResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-func (m *mockZitiMgmt) ExtendIdentityLease(context.Context, *zitimanagementv1.ExtendIdentityLeaseRequest, ...grpc.CallOption) (*zitimanagementv1.ExtendIdentityLeaseResponse, error) {
+func (m *mockRunners) ReportRunnerCatalog(context.Context, *runnersv1.ReportRunnerCatalogRequest, ...grpc.CallOption) (*runnersv1.ReportRunnerCatalogResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-func (m *mockZitiMgmt) CreateServicePolicy(context.Context, *zitimanagementv1.CreateServicePolicyRequest, ...grpc.CallOption) (*zitimanagementv1.CreateServicePolicyResponse, error) {
+func (m *mockRunners) ListStorageClasses(context.Context, *runnersv1.ListStorageClassesRequest, ...grpc.CallOption) (*runnersv1.ListStorageClassesResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-func (m *mockZitiMgmt) DeleteServicePolicy(ctx context.Context, req *zitimanagementv1.DeleteServicePolicyRequest, _ ...grpc.CallOption) (*zitimanagementv1.DeleteServicePolicyResponse, error) {
-	if m.deletePolicy == nil {
-		return nil, errors.New("not implemented")
-	}
-	return m.deletePolicy(ctx, req)
-}
-
-func (m *mockZitiMgmt) DeleteService(ctx context.Context, req *zitimanagementv1.DeleteServiceRequest, _ ...grpc.CallOption) (*zitimanagementv1.DeleteServiceResponse, error) {
-	if m.deleteSvc == nil {
-		return nil, errors.New("not implemented")
-	}
-	return m.deleteSvc(ctx, req)
-}
-
-func (m *mockZitiMgmt) CreateDeviceIdentity(context.Context, *zitimanagementv1.CreateDeviceIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.CreateDeviceIdentityResponse, error) {
+func (m *mockRunners) ListWorkloadsByAgentInstance(context.Context, *runnersv1.ListWorkloadsByAgentInstanceRequest, ...grpc.CallOption) (*runnersv1.ListWorkloadsByAgentInstanceResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-func (m *mockZitiMgmt) DeleteDeviceIdentity(context.Context, *zitimanagementv1.DeleteDeviceIdentityRequest, ...grpc.CallOption) (*zitimanagementv1.DeleteDeviceIdentityResponse, error) {
+func (m *mockRunners) ListVolumesByAgentInstance(context.Context, *runnersv1.ListVolumesByAgentInstanceRequest, ...grpc.CallOption) (*runnersv1.ListVolumesByAgentInstanceResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-type mockRunners struct {
-	getWorkload func(ctx context.Context, req *runnersv1.GetWorkloadRequest) (*runnersv1.GetWorkloadResponse, error)
-}
-
 func (m *mockRunners) RegisterRunner(context.Context, *runnersv1.RegisterRunnerRequest, ...grpc.CallOption) (*runnersv1.RegisterRunnerResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }

internal/server/identity.go

@@ -25,12 +25,20 @@ const (
 type identityType string
 
 const (
-	identityTypeUser   identityType = identitymeta.IdentityTypeUser
-	identityTypeAgent  identityType = identitymeta.IdentityTypeAgent
-	identityTypeApp    identityType = identitymeta.IdentityTypeApp
-	identityTypeRunner identityType = identitymeta.IdentityTypeRunner
+	identityTypeUser          identityType = identitymeta.IdentityTypeUser
+	identityTypeAgent         identityType = identitymeta.IdentityTypeAgent
+	identityTypeAgentInstance identityType = identitymeta.IdentityTypeAgentInstance
+	identityTypeApp           identityType = identitymeta.IdentityTypeApp
+	identityTypeRunner        identityType = identitymeta.IdentityTypeRunner
 )
 
+// isAgentWorkload reports whether the identity belongs to a running agent.
+// Workloads authenticate as their instance; identities minted before instances
+// existed still present the class type, and both own exposures.
+func (t identityType) isAgentWorkload() bool {
+	return t == identityTypeAgent || t == identityTypeAgentInstance
+}
+
 type resolvedIdentity struct {
 	identityID   string
 	identityType identityType
@@ -82,7 +90,7 @@ func resolveWorkloadIDFromRequest(caller exposureCaller, workloadID string) (str
 	if trimmed != "" {
 		return trimmed, nil
 	}
-	if caller.identity.identityType != identityTypeAgent {
+	if !caller.identity.identityType.isAgentWorkload() {
 		return "", status.Error(codes.InvalidArgument, "workload id is required")
 	}
 	callerWorkloadID := strings.TrimSpace(caller.identity.workloadID)
@@ -155,6 +163,8 @@ func parseIdentityType(value string) (identityType, error) {
 		return identityTypeUser, nil
 	case string(identityTypeAgent):
 		return identityTypeAgent, nil
+	case string(identityTypeAgentInstance):
+		return identityTypeAgentInstance, nil
 	case string(identityTypeApp):
 		return identityTypeApp, nil
 	case string(identityTypeRunner):

internal/server/identity_test.go

@@ -138,3 +138,25 @@ func TestRequireOrgRelation(t *testing.T) {
 		t.Fatalf("unexpected error: %v", err)
 	}
 }
+
+// An agent workload authenticates as its instance. If the type does not parse
+// or does not count as an agent workload, the workload cannot manage the
+// exposures it owns.
+func TestAgentInstanceIdentityIsAnAgentWorkload(t *testing.T) {
+	parsed, err := parseIdentityType("agent_instance")
+	if err != nil {
+		t.Fatalf("parse agent_instance: %v", err)
+	}
+	if parsed != identityTypeAgentInstance {
+		t.Fatalf("parsed = %q, want agent_instance", parsed)
+	}
+	if !parsed.isAgentWorkload() {
+		t.Fatal("agent_instance must count as an agent workload")
+	}
+	if !identityTypeAgent.isAgentWorkload() {
+		t.Fatal("pre-migration agent identities must keep working")
+	}
+	if identityTypeUser.isAgentWorkload() {
+		t.Fatal("user identities are not agent workloads")
+	}
+}

internal/server/server.go

@@ -66,7 +66,7 @@ func (s *Server) AddExposure(ctx context.Context, req *exposev1.AddExposureReque
 		workloadID = parsedWorkloadID
 		agentID = parsedAgentID
 	} else {
-		if caller.identity.identityType != identityTypeAgent {
+		if !caller.identity.identityType.isAgentWorkload() {
 			return nil, status.Error(codes.PermissionDenied, "permission denied")
 		}
 		resolvedWorkloadID, err := resolveWorkloadIDFromRequest(caller, "")
@@ -452,7 +452,7 @@ func workloadOrganizationID(workload *runnersv1.Workload) (uuid.UUID, error) {
 }
 
 func agentMatchesWorkload(caller exposureCaller, agentID uuid.UUID) (bool, error) {
-	if caller.identity.identityType != identityTypeAgent {
+	if !caller.identity.identityType.isAgentWorkload() {
 		return false, nil
 	}
 	callerID, err := parseIdentityUUID(caller.identity.identityID)