iop/apps/control-plane/internal/credentiallease/service_test.go

package credentiallease

import (
	"context"
	"crypto/ecdh"
	"crypto/ed25519"
	"crypto/rand"
	"google.golang.org/protobuf/proto"
	"path/filepath"
	"testing"
	"time"

	"iop/apps/control-plane/internal/credentialseal"
	"iop/apps/control-plane/internal/credentialstore"
	"iop/packages/go/config"
	lease "iop/packages/go/credentiallease"
	iop "iop/proto/gen/iop"
)

type acceptingRegistry struct{}

func (acceptingRegistry) HasEnvelopeKey(context.Context, string, uint64) (bool, error) {
	return true, nil
}

type recordingOpener struct {
	calls  int
	secret []byte
}

func (o *recordingOpener) Open(_ context.Context, _ credentialstore.SecretEnvelope, _ credentialseal.Context) ([]byte, error) {
	o.calls++
	return append([]byte(nil), o.secret...), nil
}

func TestAcquireChecksDurableBindingBeforeOpeningAndIssuesConsumableLease(t *testing.T) {
	ctx := context.Background()
	store, err := credentialstore.Open(ctx, filepath.Join(t.TempDir(), "credential.db"), credentialstore.WithEnvelopeKeyRegistry(acceptingRegistry{}))
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	principal, err := store.CreatePrincipalWithToken(ctx, credentialstore.CreatePrincipalInput{Alias: "principal"})
	if err != nil {
		t.Fatal(err)
	}
	slot, err := store.CreateSlot(ctx, credentialstore.CreateSlotInput{
		PrincipalID: principal.Principal.ID, Vendor: "openai", CredentialKind: credentialstore.CredentialKindBearer,
		Envelope: credentialstore.SecretEnvelope{Algorithm: "AES-256-GCM", KeyID: "at-rest", KeyVersion: 1, Nonce: []byte("123456789012"), Ciphertext: []byte("opaque")},
	})
	if err != nil {
		t.Fatal(err)
	}
	route, err := store.CreateRoute(ctx, credentialstore.CreateRouteInput{PrincipalID: principal.Principal.ID, SlotID: slot.ID, ProfileID: "openai", UpstreamModel: "gpt-4o", ResourceSelector: "default"})
	if err != nil {
		t.Fatal(err)
	}
	slot, err = store.GetSlot(ctx, principal.Principal.ID, slot.ID)
	if err != nil {
		t.Fatal(err)
	}
	generation, err := store.ProjectionGeneration(ctx)
	if err != nil {
		t.Fatal(err)
	}
	issuerPublic, issuerPrivate, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	recipient, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	now := time.Unix(1700000000, 0).UTC()
	opener := &recordingOpener{secret: []byte("provider-secret-sentinel")}
	service, err := New(store, opener, "issuer-1", issuerPrivate, 30*time.Second, 8, func() time.Time { return now }, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	originalProfile, ok := config.BuiltInProtocolProfiles["openai"]
	if !ok {
		t.Fatal("missing built-in openai profile")
	}
	delete(config.BuiltInProtocolProfiles, "openai")
	t.Cleanup(func() { config.BuiltInProtocolProfiles["openai"] = originalProfile })
	binding := &iop.CredentialLeaseBinding{
		PrincipalRef: principal.Principal.ID, CredentialSlotRef: slot.ID, RouteId: route.ID, ProfileId: route.ProfileID,
		UpstreamTarget: route.UpstreamModel, NodeId: "node-1", RecipientKeyId: "recipient-1",
		CredentialRevision: uint64(slot.Revision), RouteRevision: uint64(route.Revision), ProjectionGeneration: generation,
	}
	request := &iop.AcquireLeaseRequest{EdgeId: "edge-1", Binding: binding, RecipientPublicKey: recipient.PublicKey().Bytes()}
	stale := proto.Clone(request).(*iop.AcquireLeaseRequest)
	stale.Binding.ProjectionGeneration--
	response, err := service.Acquire(ctx, stale)
	if err != nil || response.GetError() == "" || opener.calls != 0 {
		t.Fatalf("stale acquire response=%v err=%v opener_calls=%d", response, err, opener.calls)
	}
	response, err = service.Acquire(ctx, request)
	if err != nil || response.GetLease() == nil || response.GetError() != "" || opener.calls != 1 {
		t.Fatalf("valid acquire response=%v err=%v opener_calls=%d", response, err, opener.calls)
	}
	consumer, err := lease.NewConsumer("node-1", "recipient-1", recipient.Bytes(), "issuer-1", issuerPublic, 8, func() time.Time { return now })
	if err != nil {
		t.Fatal(err)
	}
	envelope, err := lease.FromProto(response.GetLease())
	if err != nil {
		t.Fatal(err)
	}
	material, err := consumer.Consume(ctx, envelope, lease.ExpectedFromProto(binding))
	if err != nil {
		t.Fatal(err)
	}
	if string(material.Secret) != "provider-secret-sentinel" || material.HeaderName != "Authorization" || material.Scheme != "Bearer" {
		t.Fatal("unexpected lease material")
	}
	material.Zero()
}

// TestAcquireCanonicalizesAPIKeyProfileHeader proves that a built-in profile
// declaring a lowercase API-key header (seulgi_messages -> "x-api-key") is
// issued and consumed as the canonical "X-Api-Key" scope with no scheme, so
// the managed Messages route no longer fails closed at lease issuance.
func TestAcquireCanonicalizesAPIKeyProfileHeader(t *testing.T) {
	ctx := context.Background()
	store, err := credentialstore.Open(ctx, filepath.Join(t.TempDir(), "credential.db"), credentialstore.WithEnvelopeKeyRegistry(acceptingRegistry{}))
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	principal, err := store.CreatePrincipalWithToken(ctx, credentialstore.CreatePrincipalInput{Alias: "principal"})
	if err != nil {
		t.Fatal(err)
	}
	slot, err := store.CreateSlot(ctx, credentialstore.CreateSlotInput{
		PrincipalID: principal.Principal.ID, Vendor: "seulgi", CredentialKind: credentialstore.CredentialKindAPIKey,
		Envelope: credentialstore.SecretEnvelope{Algorithm: "AES-256-GCM", KeyID: "at-rest", KeyVersion: 1, Nonce: []byte("123456789012"), Ciphertext: []byte("opaque")},
	})
	if err != nil {
		t.Fatal(err)
	}
	route, err := store.CreateRoute(ctx, credentialstore.CreateRouteInput{PrincipalID: principal.Principal.ID, SlotID: slot.ID, ProfileID: "seulgi_messages", UpstreamModel: "seulgi-messages-1", ResourceSelector: "default"})
	if err != nil {
		t.Fatal(err)
	}
	slot, err = store.GetSlot(ctx, principal.Principal.ID, slot.ID)
	if err != nil {
		t.Fatal(err)
	}
	generation, err := store.ProjectionGeneration(ctx)
	if err != nil {
		t.Fatal(err)
	}
	issuerPublic, issuerPrivate, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	recipient, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	now := time.Unix(1700000000, 0).UTC()
	opener := &recordingOpener{secret: []byte("api-key-secret-sentinel")}
	service, err := New(store, opener, "issuer-1", issuerPrivate, 30*time.Second, 8, func() time.Time { return now }, rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	binding := &iop.CredentialLeaseBinding{
		PrincipalRef: principal.Principal.ID, CredentialSlotRef: slot.ID, RouteId: route.ID, ProfileId: route.ProfileID,
		UpstreamTarget: route.UpstreamModel, NodeId: "node-1", RecipientKeyId: "recipient-1",
		CredentialRevision: uint64(slot.Revision), RouteRevision: uint64(route.Revision), ProjectionGeneration: generation,
	}
	request := &iop.AcquireLeaseRequest{EdgeId: "edge-1", Binding: binding, RecipientPublicKey: recipient.PublicKey().Bytes()}
	response, err := service.Acquire(ctx, request)
	if err != nil || response.GetLease() == nil || response.GetError() != "" || opener.calls != 1 {
		t.Fatalf("api-key acquire response=%v err=%v opener_calls=%d", response, err, opener.calls)
	}
	consumer, err := lease.NewConsumer("node-1", "recipient-1", recipient.Bytes(), "issuer-1", issuerPublic, 8, func() time.Time { return now })
	if err != nil {
		t.Fatal(err)
	}
	envelope, err := lease.FromProto(response.GetLease())
	if err != nil {
		t.Fatal(err)
	}
	material, err := consumer.Consume(ctx, envelope, lease.ExpectedFromProto(binding))
	if err != nil {
		t.Fatal(err)
	}
	if string(material.Secret) != "api-key-secret-sentinel" || material.HeaderName != "X-Api-Key" || material.Scheme != "" {
		t.Fatalf("unexpected api-key lease material header=%q scheme=%q", material.HeaderName, material.Scheme)
	}
	material.Zero()
}