[apiserver/tunnelnode] Consolidate Token deps

Author: dilyevsky

pkg/apiserver/controllers/tunnelnode.go

@@ -30,30 +30,28 @@ type TunnelNodeReconciler struct {
 
 	jwksHost              string
 	jwksPort              int
-	jwtPrivateKeyPEM      []byte
-	jwtPublicKeyPEM       []byte
 	tokenRefreshThreshold time.Duration
 	ipamv6, ipamv4        tunnet.IPAM
 
-	validator *token.InMemoryValidator
+	validator token.Validator
 	issuer    token.TokenIssuer
 }
 
 func NewTunnelNodeReconciler(
 	c client.Client,
+	validator token.Validator,
+	issuer token.TokenIssuer,
 	jwksHost string,
 	jwksPort int,
-	jwtPrivateKey []byte,
-	jwtPublicKey []byte,
 	tokenRefreshThreshold time.Duration,
 	ipamv6, ipamv4 tunnet.IPAM,
 ) *TunnelNodeReconciler {
 	return &TunnelNodeReconciler{
 		Client:                c,
+		validator:             validator,
+		issuer:                issuer,
 		jwksHost:              jwksHost,
 		jwksPort:              jwksPort,
-		jwtPrivateKeyPEM:      jwtPrivateKey,
-		jwtPublicKeyPEM:       jwtPublicKey,
 		tokenRefreshThreshold: tokenRefreshThreshold,
 		ipamv6:                ipamv6,
 		ipamv4:                ipamv4,
@@ -222,25 +220,15 @@ func (r *TunnelNodeReconciler) Reconcile(ctx context.Context, req reconcile.Requ
 	return ctrl.Result{}, nil
 }
 
-func (r *TunnelNodeReconciler) SetupWithManager(ctx context.Context, mgr ctrl.Manager) error {
-	var err error
-	r.validator, err = token.NewInMemoryValidator(r.jwtPublicKeyPEM)
-	if err != nil {
-		return fmt.Errorf("failed to create token validator: %w", err)
-	}
-	r.issuer, err = token.NewIssuer(r.jwtPrivateKeyPEM)
-	if err != nil {
-		return fmt.Errorf("failed to create token issuer: %w", err)
-	}
-
+func (r *TunnelNodeReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&corev1alpha.TunnelNode{}).
 		Complete(r)
 }
 
 // ServeJWKS starts an HTTP server to serve JWK sets
 func (r *TunnelNodeReconciler) ServeJWKS(ctx context.Context) error {
-	jwksHandler, err := token.NewJWKSHandler(r.jwtPublicKeyPEM)
+	jwksHandler, err := token.NewJWKSHandler(r.validator.PublicKeyPEM())
 	if err != nil {
 		return fmt.Errorf("failed to create JWKS handler: %w", err)
 	}

pkg/apiserver/controllers/tunnelnode_test.go

@@ -42,27 +42,27 @@ func TestTunnelNodeReconciler(t *testing.T) {
 	privKey, pubKey, err := cryptoutils.GenerateEllipticKeyPair()
 	require.NoError(t, err)
 
+	validator, err := token.NewInMemoryValidator(pubKey)
+	require.NoError(t, err)
+	issuer, err := token.NewIssuer(privKey)
+	require.NoError(t, err)
+
 	systemULA := tunnet.NewULA(context.Background(), net.SystemNetworkID)
 	// Agent prefixes are /96 subnets that can embed IPv4 suffixes.
 	ipamv6, err := systemULA.IPAM(context.Background(), 96)
 	require.NoError(t, err)
 
 	r := NewTunnelNodeReconciler(
 		k8sClient,
+		validator,
+		issuer,
 		"localhost",
 		9444,
-		privKey,
-		pubKey,
 		time.Minute,
 		ipamv6,
 		net.NewIPAMv4(context.Background()),
 	)
 
-	r.validator, err = token.NewInMemoryValidator(r.jwtPublicKeyPEM)
-	require.NoError(t, err)
-	r.issuer, err = token.NewIssuer(r.jwtPrivateKeyPEM)
-	require.NoError(t, err)
-
 	_, err = r.Reconcile(context.TODO(), reconcile.Request{
 		NamespacedName: types.NamespacedName{
 			Name:      "test-tunnelnode",

pkg/apiserver/manager.go

@@ -49,6 +49,7 @@ import (
 	"github.com/apoxy-dev/apoxy/pkg/log"
 	apoxynet "github.com/apoxy-dev/apoxy/pkg/tunnel/net"
 	tunnet "github.com/apoxy-dev/apoxy/pkg/tunnel/net"
+	"github.com/apoxy-dev/apoxy/pkg/tunnel/token"
 	"github.com/apoxy-dev/apoxy/pkg/tunnel/vni"
 
 	corev1alpha "github.com/apoxy-dev/apoxy/api/core/v1alpha"
@@ -512,17 +513,25 @@ func (m *Manager) Start(
 
 	// Legacy v1alpha1 TunnelNode controller
 	log.Infof("Registering TunnelNode controller")
+	tunnelNodeValidator, err := token.NewInMemoryValidator(dOpts.jwtPublicKey)
+	if err != nil {
+		return fmt.Errorf("failed to create tunnel node validator: %v", err)
+	}
+	tunnelNodeIssuer, err := token.NewIssuer(dOpts.jwtPrivateKey)
+	if err != nil {
+		return fmt.Errorf("failed to create tunnel node issuer: %v", err)
+	}
 	tunnelNodeReconciler := controllers.NewTunnelNodeReconciler(
 		m.manager.GetClient(),
+		tunnelNodeValidator,
+		tunnelNodeIssuer,
 		dOpts.jwksHost,
 		dOpts.jwksPort,
-		dOpts.jwtPrivateKey,
-		dOpts.jwtPublicKey,
 		dOpts.jwtRefreshThreshold,
 		dOpts.agentIPAM,
 		apoxynet.NewIPAMv4(context.Background()),
 	)
-	if err := tunnelNodeReconciler.SetupWithManager(ctx, m.manager); err != nil {
+	if err := tunnelNodeReconciler.SetupWithManager(m.manager); err != nil {
 		return fmt.Errorf("failed to set up TunnelNode controller: %v", err)
 	}
 	g.Go(func() error {

pkg/tunnel/token/validator.go

@@ -15,10 +15,21 @@ import (
 	"github.com/apoxy-dev/apoxy/pkg/cryptoutils"
 )
 
+// JWTValidator validates JWT tokens.
 type JWTValidator interface {
+	// Validate validates the token and returns its claims.
 	Validate(tokenStr string) (jwt.Claims, error)
 }
 
+// Validator extends JWTValidator with the ability to retrieve the public key.
+// Use this interface when you need to serve JWKS endpoints locally.
+// For remote validation (e.g., RemoteValidator), use JWTValidator instead.
+type Validator interface {
+	JWTValidator
+	// PublicKeyPEM returns the PEM-encoded public key used for validation.
+	PublicKeyPEM() []byte
+}
+
 func validate(keyFunc jwt.Keyfunc, tokenStr string) (jwt.Claims, error) {
 	token, err := jwt.Parse(
 		tokenStr,
@@ -42,9 +53,11 @@ func validate(keyFunc jwt.Keyfunc, tokenStr string) (jwt.Claims, error) {
 	return tokenClaims, nil
 }
 
-// InMemoryValidator validates JWT tokens signed with an ECDSA public key
+// InMemoryValidator validates JWT tokens signed with an ECDSA public key.
+// It implements the Validator interface.
 type InMemoryValidator struct {
-	publicKey *ecdsa.PublicKey
+	publicKey    *ecdsa.PublicKey
+	publicKeyPEM []byte
 }
 
 // NewInMemoryValidator creates a new Validator with the public key.
@@ -54,7 +67,10 @@ func NewInMemoryValidator(publicKeyPEM []byte) (*InMemoryValidator, error) {
 		return nil, fmt.Errorf("failed to parse ECDSA public key: %w", err)
 	}
 
-	return &InMemoryValidator{publicKey: publicKey}, nil
+	return &InMemoryValidator{
+		publicKey:    publicKey,
+		publicKeyPEM: publicKeyPEM,
+	}, nil
 }
 
 // Validate validates the token is valid.
@@ -64,6 +80,11 @@ func (v *InMemoryValidator) Validate(tokenStr string) (jwt.Claims, error) {
 	}, tokenStr)
 }
 
+// PublicKeyPEM returns the PEM-encoded public key used for validation.
+func (v *InMemoryValidator) PublicKeyPEM() []byte {
+	return v.publicKeyPEM
+}
+
 type RemoteValidator struct {
 	rkf keyfunc.Keyfunc
 }