iop/apps/edge/internal/service/provider_tunnel_credential_test.go

package service

import (
	"context"
	"errors"
	"strings"
	"testing"

	edgenode "iop/apps/edge/internal/node"
	iop "iop/proto/gen/iop"
)

type fakeCredentialLeaseProvider struct {
	valid        bool
	acquireCalls int
	gotBinding   *iop.CredentialLeaseBinding
	mutateScope  func(*iop.CredentialLeaseScope)
}

func (p *fakeCredentialLeaseProvider) ValidateCredentialBinding(binding *iop.CredentialLeaseBinding) error {
	if !p.valid {
		return errors.New("stale")
	}
	return nil
}

func (p *fakeCredentialLeaseProvider) AcquireCredentialLease(_ context.Context, binding *iop.CredentialLeaseBinding, publicKey []byte) (*iop.SignedCredentialLease, error) {
	p.acquireCalls++
	p.gotBinding = binding
	if len(publicKey) != 32 {
		return nil, errors.New("bad recipient")
	}
	scope := &iop.CredentialLeaseScope{
		LeaseId: "lease-test", PrincipalRef: binding.GetPrincipalRef(), CredentialSlotRef: binding.GetCredentialSlotRef(),
		RouteId: binding.GetRouteId(), ProfileId: binding.GetProfileId(), UpstreamTarget: binding.GetUpstreamTarget(),
		NodeId: binding.GetNodeId(), RecipientKeyId: binding.GetRecipientKeyId(), HeaderName: "Authorization", Scheme: "Bearer",
		CredentialRevision: binding.GetCredentialRevision(), RouteRevision: binding.GetRouteRevision(),
		ProjectionGeneration: binding.GetProjectionGeneration(),
	}
	if p.mutateScope != nil {
		p.mutateScope(scope)
	}
	return &iop.SignedCredentialLease{Scope: scope}, nil
}

func TestAttachCredentialLeaseBindsSelectedCandidateAndUsesDedicatedFields(t *testing.T) {
	provider := &fakeCredentialLeaseProvider{valid: true}
	service := New(nil, nil)
	service.SetCredentialLeaseProvider(provider)
	service.SetCredentialLeaseLimit(1)
	req := SubmitProviderTunnelRequest{CredentialBinding: &CredentialBinding{
		PrincipalRef: "principal", CredentialSlotRef: "slot", RouteID: "route", ProfileID: "openai",
		CredentialRevision: 7, RouteRevision: 8, ProjectionGeneration: 9,
	}}
	entry := &edgenode.NodeEntry{NodeID: "selected-node", CredentialRecipientKeyID: "recipient-key", CredentialRecipientPublicKey: make([]byte, 32)}
	tunnel := &iop.ProviderTunnelRequest{Headers: map[string]string{"Content-Type": "application/json"}}
	if err := service.attachCredentialLease(context.Background(), req, entry, "selected-target", tunnel); err != nil {
		t.Fatal(err)
	}
	defer service.releaseCredentialLease()
	if provider.acquireCalls != 1 || provider.gotBinding.GetNodeId() != "selected-node" || provider.gotBinding.GetUpstreamTarget() != "selected-target" || provider.gotBinding.GetRecipientKeyId() != "recipient-key" {
		t.Fatalf("lease was not bound to selected candidate: %+v", provider.gotBinding)
	}
	if tunnel.GetCredentialLease() == nil || tunnel.GetCredentialBinding() == nil {
		t.Fatal("dedicated credential fields were not attached")
	}
	for key, value := range tunnel.GetHeaders() {
		if strings.Contains(strings.ToLower(key+value), "secret") || strings.EqualFold(key, "Authorization") {
			t.Fatalf("generic header contains provider credential material: %q", key)
		}
	}
}

func TestAttachCredentialLeaseFailsBeforeAcquireOnStaleBindingOrCollision(t *testing.T) {
	provider := &fakeCredentialLeaseProvider{valid: false}
	service := New(nil, nil)
	service.SetCredentialLeaseProvider(provider)
	service.SetCredentialLeaseLimit(1)
	req := SubmitProviderTunnelRequest{CredentialBinding: &CredentialBinding{PrincipalRef: "p", CredentialSlotRef: "s", RouteID: "r", ProfileID: "openai", ProjectionGeneration: 1}}
	entry := &edgenode.NodeEntry{NodeID: "node", CredentialRecipientKeyID: "key", CredentialRecipientPublicKey: make([]byte, 32)}
	if err := service.attachCredentialLease(context.Background(), req, entry, "target", &iop.ProviderTunnelRequest{}); err == nil || provider.acquireCalls != 0 {
		t.Fatalf("stale binding err=%v acquire_calls=%d", err, provider.acquireCalls)
	}
	provider.valid = true
	if err := service.attachCredentialLease(context.Background(), req, entry, "target", &iop.ProviderTunnelRequest{Headers: map[string]string{"authorization": "caller-value"}}); err == nil || !strings.Contains(err.Error(), "collision") {
		t.Fatalf("collision err=%v", err)
	}
	if provider.acquireCalls != 1 {
		t.Fatalf("collision acquisition calls=%d, want 1 signed-header check", provider.acquireCalls)
	}
}

func TestAttachCredentialLeaseRejectsMismatchedReturnedScope(t *testing.T) {
	provider := &fakeCredentialLeaseProvider{valid: true, mutateScope: func(scope *iop.CredentialLeaseScope) {
		scope.CredentialSlotRef = "different-slot"
	}}
	service := New(nil, nil)
	service.SetCredentialLeaseProvider(provider)
	service.SetCredentialLeaseLimit(1)
	req := SubmitProviderTunnelRequest{CredentialBinding: &CredentialBinding{
		PrincipalRef: "principal", CredentialSlotRef: "slot", RouteID: "route", ProfileID: "openai",
		CredentialRevision: 7, RouteRevision: 8, ProjectionGeneration: 9,
	}}
	entry := &edgenode.NodeEntry{NodeID: "node", CredentialRecipientKeyID: "key", CredentialRecipientPublicKey: make([]byte, 32)}
	err := service.attachCredentialLease(context.Background(), req, entry, "target", &iop.ProviderTunnelRequest{})
	if err == nil || !strings.Contains(err.Error(), "scope mismatch") {
		t.Fatalf("mismatched lease scope error = %v", err)
	}
}

func TestSubmitProviderPoolTunnelAttachesCredentialLeaseAfterCandidateSelection(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	provider := &fakeCredentialLeaseProvider{valid: true}
	env.svc.SetCredentialLeaseProvider(provider)
	env.svc.SetCredentialLeaseLimit(1)

	result, err := env.svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
		Run: SubmitRunRequest{ModelGroupKey: "qwen3.6:35b", ProviderPool: true},
		Tunnel: SubmitProviderTunnelRequest{
			CredentialBinding: &CredentialBinding{
				PrincipalRef: "principal", CredentialSlotRef: "slot", RouteID: "route", ProfileID: "openai",
				CredentialRevision: 7, RouteRevision: 8, ProjectionGeneration: 9,
			},
			Method: "POST", Path: "/v1/chat/completions", Operation: "chat_completions",
			BuildBody: func(string) ([]byte, error) { return []byte(`{"model":"served-qwen"}`), nil },
		},
	})
	if err != nil {
		t.Fatalf("SubmitProviderPool: %v", err)
	}
	defer result.Tunnel.Close()
	waitForCondition(t, func() bool { return env.capturedRequest() != nil }, "fake node did not receive managed provider-pool tunnel")
	captured := env.capturedRequest()
	if captured.GetCredentialLease() == nil || captured.GetCredentialBinding() == nil {
		t.Fatal("managed provider-pool tunnel omitted dedicated credential fields")
	}
	if provider.acquireCalls != 1 || provider.gotBinding.GetNodeId() != "node-pool" || provider.gotBinding.GetRecipientKeyId() != "recipient-pool" || provider.gotBinding.GetUpstreamTarget() != "served-qwen" {
		t.Fatalf("lease acquisition did not use the selected candidate: %+v", provider.gotBinding)
	}
	if got := result.DispatchInfo; got.CredentialSlotRef != "slot" || got.CredentialRevision != 7 {
		t.Fatalf("dispatch credential attribution = %q/%d, want slot/7", got.CredentialSlotRef, got.CredentialRevision)
	}
	if got := result.Tunnel.Dispatch(); got.CredentialSlotRef != "slot" || got.CredentialRevision != 7 {
		t.Fatalf("tunnel credential attribution = %q/%d, want slot/7", got.CredentialSlotRef, got.CredentialRevision)
	}
}