package worker

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"net/url"
	"testing"

	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/peer"
)

func TestSPIFFEIdentityAuthorizer(t *testing.T) {
	authorizer, err := NewSPIFFEIdentityAuthorizer("proxy.example", "prod")
	if err != nil {
		t.Fatalf("NewSPIFFEIdentityAuthorizer(): %v", err)
	}

	tests := []struct {
		name      string
		workerID  string
		uri       string
		withTLS   bool
		wantError bool
	}{
		{name: "matching worker uri", workerID: "worker-a", uri: "spiffe://proxy.example/prod/worker/worker-a", withTLS: true},
		{name: "different worker", workerID: "worker-a", uri: "spiffe://proxy.example/prod/worker/worker-b", withTLS: true, wantError: true},
		{name: "different environment", workerID: "worker-a", uri: "spiffe://proxy.example/staging/worker/worker-a", withTLS: true, wantError: true},
		{name: "different trust domain", workerID: "worker-a", uri: "spiffe://other.example/prod/worker/worker-a", withTLS: true, wantError: true},
		{name: "missing peer tls", workerID: "worker-a", wantError: true},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			ctx := context.Background()
			if test.withTLS {
				ctx = tlsPeerContext(t, test.uri)
			}
			err := authorizer.Authorize(ctx, test.workerID)
			if (err != nil) != test.wantError {
				t.Fatalf("Authorize() error = %v, wantError %v", err, test.wantError)
			}
		})
	}
}

func TestSPIFFEIdentityAuthorizerRejectsMultipleWorkerURIs(t *testing.T) {
	authorizer, err := NewSPIFFEIdentityAuthorizer("proxy.example", "prod")
	if err != nil {
		t.Fatalf("NewSPIFFEIdentityAuthorizer(): %v", err)
	}
	first, err := url.Parse("spiffe://proxy.example/prod/worker/worker-a")
	if err != nil {
		t.Fatal(err)
	}
	second, err := url.Parse("spiffe://proxy.example/prod/worker/worker-b")
	if err != nil {
		t.Fatal(err)
	}
	ctx := peer.NewContext(context.Background(), &peer.Peer{AuthInfo: credentials.TLSInfo{
		State: tls.ConnectionState{VerifiedChains: [][]*x509.Certificate{{{URIs: []*url.URL{first, second}}}}},
	}})
	if err := authorizer.Authorize(ctx, "worker-a"); err == nil {
		t.Fatal("Authorize() error = nil, want rejection for multiple worker identities")
	}
}

func TestSPIFFEIdentityAuthorizerAuthorizesCheckerOnlyForMatchingResource(t *testing.T) {
	authorizer, err := NewSPIFFEIdentityAuthorizer("proxy.example", "prod")
	if err != nil {
		t.Fatalf("NewSPIFFEIdentityAuthorizer(): %v", err)
	}
	if err := authorizer.AuthorizeChecker(tlsPeerContext(t, "spiffe://proxy.example/prod/checker/checker-a"), "checker-a"); err != nil {
		t.Fatalf("AuthorizeChecker(matching checker) = %v", err)
	}
	if err := authorizer.AuthorizeChecker(tlsPeerContext(t, "spiffe://proxy.example/prod/worker/checker-a"), "checker-a"); err == nil {
		t.Fatal("AuthorizeChecker(worker identity) error = nil")
	}
	if err := authorizer.Authorize(tlsPeerContext(t, "spiffe://proxy.example/prod/checker/checker-a"), "checker-a"); err == nil {
		t.Fatal("Authorize(checker identity) error = nil")
	}
}

func TestAllowLoopbackIdentity(t *testing.T) {
	if err := (AllowLoopbackIdentity{}).Authorize(context.Background(), "worker-a"); err != nil {
		t.Fatalf("Authorize(): %v", err)
	}
	if err := (AllowLoopbackIdentity{}).AuthorizeChecker(context.Background(), "checker-a"); err != nil {
		t.Fatalf("AuthorizeChecker(): %v", err)
	}
}

func tlsPeerContext(t *testing.T, identityURI string) context.Context {
	t.Helper()
	uri, err := url.Parse(identityURI)
	if err != nil {
		t.Fatalf("url.Parse(): %v", err)
	}
	return peer.NewContext(context.Background(), &peer.Peer{AuthInfo: credentials.TLSInfo{
		State: tls.ConnectionState{VerifiedChains: [][]*x509.Certificate{{{URIs: []*url.URL{uri}}}}},
	}})
}