feat: accept agent_instance identities

internal/identity/identity.go

@@ -31,6 +31,19 @@ type IdentityChain struct {
 	OrganizationID string
 }
 
+// isAgentIdentityType 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 resolve to an agent and org.
+func isAgentIdentityType(identityType identityv1.IdentityType) bool {
+	switch identityType {
+	case identityv1.IdentityType_IDENTITY_TYPE_AGENT,
+		identityv1.IdentityType_IDENTITY_TYPE_AGENT_INSTANCE:
+		return true
+	default:
+		return false
+	}
+}
+
 type ZitiResolver interface {
 	ResolveIdentity(ctx context.Context, sourceIdentity string) (ResolvedIdentity, error)
 }
@@ -76,8 +89,8 @@ func (r *Resolver) Resolve(ctx context.Context, sourceIdentity string) (Identity
 	if err != nil {
 		return IdentityChain{}, fmt.Errorf("resolve ziti identity: %w", err)
 	}
-	if resolved.IdentityType != identityv1.IdentityType_IDENTITY_TYPE_AGENT {
-		return IdentityChain{}, fmt.Errorf("identity type is not agent: %s", resolved.IdentityType.String())
+	if !isAgentIdentityType(resolved.IdentityType) {
+		return IdentityChain{}, fmt.Errorf("identity type is not an agent: %s", resolved.IdentityType.String())
 	}
 
 	agentIdentity, err := r.agentsResolver.ResolveAgentIdentity(ctx, resolved.IdentityID)

internal/identity/identity_test.go

@@ -0,0 +1,94 @@
+package identity
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	identityv1 "github.com/agynio/tracing/.gen/go/agynio/api/identity/v1"
+	"github.com/agynio/tracing/internal/cache"
+)
+
+type stubZitiResolver struct {
+	resolved ResolvedIdentity
+}
+
+func (s stubZitiResolver) ResolveIdentity(context.Context, string) (ResolvedIdentity, error) {
+	return s.resolved, nil
+}
+
+type stubAgentsResolver struct {
+	identityID string
+}
+
+func (s *stubAgentsResolver) ResolveAgentIdentity(_ context.Context, identityID string) (AgentIdentity, error) {
+	s.identityID = identityID
+	return AgentIdentity{AgentID: "agent-1", OrganizationID: "org-1"}, nil
+}
+
+func newTestResolver(t *testing.T, identityType identityv1.IdentityType) (*Resolver, *stubAgentsResolver) {
+	t.Helper()
+	lru, err := cache.NewLRU[string, IdentityChain](8, time.Minute)
+	if err != nil {
+		t.Fatalf("new cache: %v", err)
+	}
+	agents := &stubAgentsResolver{}
+	resolver, err := NewResolver(
+		stubZitiResolver{resolved: ResolvedIdentity{IdentityID: "identity-1", IdentityType: identityType}},
+		agents,
+		lru,
+	)
+	if err != nil {
+		t.Fatalf("new resolver: %v", err)
+	}
+	return resolver, agents
+}
+
+// Agent workloads authenticate as their instance, so an instance identity has
+// to resolve; identities minted before instances existed still present the
+// class type and must keep working.
+func TestResolveAcceptsAgentAndInstanceIdentities(t *testing.T) {
+	for _, identityType := range []identityv1.IdentityType{
+		identityv1.IdentityType_IDENTITY_TYPE_AGENT,
+		identityv1.IdentityType_IDENTITY_TYPE_AGENT_INSTANCE,
+	} {
+		t.Run(identityType.String(), func(t *testing.T) {
+			resolver, agents := newTestResolver(t, identityType)
+			chain, err := resolver.Resolve(context.Background(), "source-1")
+			if err != nil {
+				t.Fatalf("expected no error, got %v", err)
+			}
+			if chain.AgentID != "agent-1" || chain.OrganizationID != "org-1" {
+				t.Fatalf("unexpected chain: %+v", chain)
+			}
+			if chain.IdentityType != identityType {
+				t.Fatalf("expected identity type %s, got %s", identityType, chain.IdentityType)
+			}
+			if agents.identityID != "identity-1" {
+				t.Fatalf("expected the resolved identity id, got %q", agents.identityID)
+			}
+		})
+	}
+}
+
+func TestResolveRejectsNonAgentIdentities(t *testing.T) {
+	for _, identityType := range []identityv1.IdentityType{
+		identityv1.IdentityType_IDENTITY_TYPE_RUNNER,
+		identityv1.IdentityType_IDENTITY_TYPE_USER,
+		identityv1.IdentityType_IDENTITY_TYPE_APP,
+		identityv1.IdentityType_IDENTITY_TYPE_SANDBOX,
+		identityv1.IdentityType_IDENTITY_TYPE_UNSPECIFIED,
+	} {
+		t.Run(identityType.String(), func(t *testing.T) {
+			resolver, _ := newTestResolver(t, identityType)
+			_, err := resolver.Resolve(context.Background(), "source-1")
+			if err == nil {
+				t.Fatal("expected an error")
+			}
+			if !strings.Contains(err.Error(), "not an agent") {
+				t.Fatalf("unexpected error: %v", err)
+			}
+		})
+	}
+}

internal/identity/metadata.go

@@ -7,16 +7,19 @@ import (
 )
 
 const (
-	identityTypeAgent  = "agent"
-	identityTypeRunner = "runner"
-	identityTypeUser   = "user"
-	identityTypeApp    = "app"
+	identityTypeAgent         = "agent"
+	identityTypeAgentInstance = "agent_instance"
+	identityTypeRunner        = "runner"
+	identityTypeUser          = "user"
+	identityTypeApp           = "app"
 )
 
 func IdentityTypeMetadataValue(identityType identityv1.IdentityType) (string, error) {
 	switch identityType {
 	case identityv1.IdentityType_IDENTITY_TYPE_AGENT:
 		return identityTypeAgent, nil
+	case identityv1.IdentityType_IDENTITY_TYPE_AGENT_INSTANCE:
+		return identityTypeAgentInstance, nil
 	case identityv1.IdentityType_IDENTITY_TYPE_RUNNER:
 		return identityTypeRunner, nil
 	case identityv1.IdentityType_IDENTITY_TYPE_USER:

internal/identity/metadata_test.go

@@ -0,0 +1,41 @@
+package identity
+
+import (
+	"testing"
+
+	identityv1 "github.com/agynio/tracing/.gen/go/agynio/api/identity/v1"
+)
+
+func TestIdentityTypeMetadataValue(t *testing.T) {
+	tests := map[identityv1.IdentityType]string{
+		identityv1.IdentityType_IDENTITY_TYPE_AGENT:          "agent",
+		identityv1.IdentityType_IDENTITY_TYPE_AGENT_INSTANCE: "agent_instance",
+		identityv1.IdentityType_IDENTITY_TYPE_RUNNER:         "runner",
+		identityv1.IdentityType_IDENTITY_TYPE_USER:           "user",
+		identityv1.IdentityType_IDENTITY_TYPE_APP:            "app",
+	}
+	for identityType, want := range tests {
+		t.Run(identityType.String(), func(t *testing.T) {
+			got, err := IdentityTypeMetadataValue(identityType)
+			if err != nil {
+				t.Fatalf("expected no error, got %v", err)
+			}
+			if got != want {
+				t.Fatalf("expected %q, got %q", want, got)
+			}
+		})
+	}
+}
+
+func TestIdentityTypeMetadataValueRejectsUnsupported(t *testing.T) {
+	for _, identityType := range []identityv1.IdentityType{
+		identityv1.IdentityType_IDENTITY_TYPE_UNSPECIFIED,
+		identityv1.IdentityType_IDENTITY_TYPE_SANDBOX,
+	} {
+		t.Run(identityType.String(), func(t *testing.T) {
+			if _, err := IdentityTypeMetadataValue(identityType); err == nil {
+				t.Fatal("expected an error")
+			}
+		})
+	}
+}