Various connect-ip tunnel improvements and get end-to-end test passing

go.mod

@@ -6,6 +6,8 @@ toolchain go1.23.2
 
 require (
 	github.com/ClickHouse/clickhouse-go/v2 v2.23.2
+	github.com/MicahParks/jwkset v0.9.5
+	github.com/MicahParks/keyfunc/v3 v3.3.10
 	github.com/alphadose/haxmap v1.4.1
 	github.com/buraksezer/olric v0.5.6
 	github.com/cncf/xds/go v0.0.0-20240905190251-b4127c9b8d78
@@ -15,7 +17,7 @@ require (
 	github.com/dgraph-io/badger/v4 v4.5.0
 	github.com/docker/docker v27.2.0+incompatible
 	github.com/dpeckett/contextio v0.5.1
-	github.com/dpeckett/network v0.3.3
+	github.com/dpeckett/network v0.3.4
 	github.com/dpeckett/triemap v0.3.1
 	github.com/envoyproxy/gateway v0.5.0-rc.1.0.20240618131507-bdff5d56b59d
 	github.com/envoyproxy/go-control-plane v0.13.0
@@ -26,6 +28,7 @@ require (
 	github.com/getsentry/sentry-go v0.26.0
 	github.com/go-logr/logr v1.4.2
 	github.com/goccy/go-json v0.9.11
+	github.com/golang-jwt/jwt v3.2.2+incompatible
 	github.com/golang-jwt/jwt/v5 v5.2.2
 	github.com/golang-migrate/migrate/v4 v4.17.1
 	github.com/golang/protobuf v1.5.4
@@ -65,7 +68,6 @@ require (
 	golang.org/x/sys v0.29.0
 	golang.org/x/tools v0.26.0
 	golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173
-	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 	google.golang.org/grpc v1.68.0
 	google.golang.org/protobuf v1.35.2
 	gopkg.in/yaml.v2 v2.4.0
@@ -108,8 +110,6 @@ require (
 	github.com/Azure/go-autorest/logger v0.2.1 // indirect
 	github.com/Azure/go-autorest/tracing v0.6.0 // indirect
 	github.com/ClickHouse/ch-go v0.61.5 // indirect
-	github.com/MicahParks/jwkset v0.9.5 // indirect
-	github.com/MicahParks/keyfunc/v3 v3.3.10 // indirect
 	github.com/NYTimes/gziphandler v1.1.1 // indirect
 	github.com/Rican7/retry v0.1.0 // indirect
 	github.com/RoaringBitmap/roaring v1.2.1 // indirect
@@ -178,7 +178,6 @@ require (
 	github.com/gocql/gocql v1.6.0 // indirect
 	github.com/godbus/dbus/v5 v5.1.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
-	github.com/golang-jwt/jwt v3.2.2+incompatible // indirect
 	github.com/golang-jwt/jwt/v4 v4.5.1 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/golang/mock v1.7.0-rc.1 // indirect
@@ -226,7 +225,6 @@ require (
 	github.com/jmoiron/sqlx v1.4.0 // indirect
 	github.com/jonboulle/clockwork v0.2.2 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
-	github.com/josharian/native v1.1.0 // indirect
 	github.com/joyent/triton-go v0.0.0-20180628001255-830d2b111e62 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
@@ -243,9 +241,6 @@ require (
 	github.com/mattn/go-runewidth v0.0.15 // indirect
 	github.com/mattn/go-sqlite3 v1.14.22 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
-	github.com/mdlayher/genetlink v1.3.2 // indirect
-	github.com/mdlayher/netlink v1.7.2 // indirect
-	github.com/mdlayher/socket v0.4.1 // indirect
 	github.com/minio/highwayhash v1.0.2 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/moby/sys/mountinfo v0.7.1 // indirect

go.sum

@@ -897,10 +897,8 @@ github.com/docker/spdystream v0.0.0-20160310174837-449fdfce4d96/go.mod h1:Qh8CwZ
 github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE=
 github.com/dpeckett/contextio v0.5.1 h1:w19s6EThbZuRpa2z/Lu06v6+o3rrZhbBzmkol6en/hA=
 github.com/dpeckett/contextio v0.5.1/go.mod h1:IY/CQ1ee6y4C5j/mU0X0M/D84s2FxNisggbNClTPndc=
-github.com/dpeckett/network v0.3.1 h1:rMDRLc85zc3v4mGcGfbOrNA9Kx69K2Xr8bD/Hc9MERY=
-github.com/dpeckett/network v0.3.1/go.mod h1:83quX+FE+BdOAKFEm5Om+QdI/1ZEQVNUBZSPl7V7erk=
-github.com/dpeckett/network v0.3.3 h1:ghUiAOddpgKSBokp7/CVhjkO+ZjsuczeerwB+i9dBdw=
-github.com/dpeckett/network v0.3.3/go.mod h1:AyOdc+YGAT7zWoDDd/r8Pv4kGNlKrboVcug0FKaoqbA=
+github.com/dpeckett/network v0.3.4 h1:zOa5hIjOKtl13+FGde6Av4VrD+PmD03AadMKYrhkoOg=
+github.com/dpeckett/network v0.3.4/go.mod h1:AyOdc+YGAT7zWoDDd/r8Pv4kGNlKrboVcug0FKaoqbA=
 github.com/dpeckett/triemap v0.3.1 h1:jzxCyKs/ATw9uCdD2bd0xFTPLIP9uZwX0iZUOOOIDoc=
 github.com/dpeckett/triemap v0.3.1/go.mod h1:pBxNH+K6m5I4lVo+W7u6JEanxP13adD4t2XYVMxfmTo=
 github.com/dunglas/httpsfv v1.0.2 h1:iERDp/YAfnojSDJ7PW3dj1AReJz4MrwbECSSE59JWL0=
@@ -1439,8 +1437,6 @@ github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxU
 github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo=
 github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4=
 github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0/go.mod h1:QUyp042oQthUoa9bqDv0ER0wrtXnBruoNd7aNjkbP+k=
-github.com/mdlayher/genetlink v1.3.2 h1:KdrNKe+CTu+IbZnm/GVUMXSqBBLqcGpRDa0xkQy56gw=
-github.com/mdlayher/genetlink v1.3.2/go.mod h1:tcC3pkCrPUGIKKsCsp0B3AdaaKuHtaxoJRz3cc+528o=
 github.com/mdlayher/netlink v1.7.2 h1:/UtM3ofJap7Vl4QWCPDGXY8d3GIY2UGSDbK+QWmY8/g=
 github.com/mdlayher/netlink v1.7.2/go.mod h1:xraEF7uJbxLhc5fpHL4cPe221LI2bdttWlU+ZGLfQSw=
 github.com/mdlayher/socket v0.4.1 h1:eM9y2/jlbs1M615oshPQOHZzj6R6wMT7bX5NPiQvn2U=
@@ -1452,8 +1448,6 @@ github.com/miekg/dns v1.1.26/go.mod h1:bPDLeHnStXmXAq1m/Ch/hvfNHr14JKNPMBo3VZKju
 github.com/miekg/dns v1.1.45/go.mod h1:e3IlAVfNqAllflbibAZEWOXOQ+Ynzk/dDozDxY7XnME=
 github.com/miekg/dns v1.1.63 h1:8M5aAw6OMZfFXTT7K5V0Eu5YiiL8l7nUAkyN6C9YwaY=
 github.com/miekg/dns v1.1.63/go.mod h1:6NGHfjhpmr5lt3XPLuyfDJi5AXbNIPM9PY6H6sF1Nfs=
-github.com/mikioh/ipaddr v0.0.0-20190404000644-d465c8ab6721 h1:RlZweED6sbSArvlE924+mUcZuXKLBHA35U7LN621Bws=
-github.com/mikioh/ipaddr v0.0.0-20190404000644-d465c8ab6721/go.mod h1:Ickgr2WtCLZ2MDGd4Gr0geeCH5HybhRJbonOgQpvSxc=
 github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
 github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE=
 github.com/minio/highwayhash v1.0.2 h1:Aak5U0nElisjDCfPSG79Tgzkn2gl66NxOMspRrKnA/g=
@@ -2330,8 +2324,6 @@ golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11/go.mod h1:tRJNPiyCQ0inRvYxb
 golang.org/x/time v0.0.0-20220922220347-f3bd1da661af/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.1.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
-golang.org/x/time v0.8.0 h1:9i3RxcPv3PZnitoVGMPDKZSq1xW1gK1Xy3ArNOGZfEg=
-golang.org/x/time v0.8.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY=
 golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
@@ -2425,8 +2417,6 @@ golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 h1:B82qJJgjvYKsXS9jeu
 golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2/go.mod h1:deeaetjYA+DHMHg+sMSMI58GrEteJUUzzw7en6TJQcI=
 golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173 h1:/jFs0duh4rdb8uIfPMv78iAJGcPKDeqAFnaLBropIC4=
 golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173/go.mod h1:tkCQ4FQXmpAgYVh++1cq16/dH4QJtmvpRv19DWGAHSA=
-golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6 h1:CawjfCvYQH2OU3/TnxLx97WDSUDRABfT18pCOYwc2GE=
-golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6/go.mod h1:3rxYc4HtVcSG9gVaTs2GEBdehh+sYPOwKtyUWEOTb80=
 gomodules.xyz/jsonpatch/v2 v2.0.1/go.mod h1:IhYNNY4jnS53ZnfE4PAmpKtDpTCj1JFXc+3mwe7XcUU=
 gomodules.xyz/jsonpatch/v2 v2.4.0 h1:Ci3iUJyx9UeRx7CeFN8ARgGbkESwJK+KB9lLcWxY/Zw=
 gomodules.xyz/jsonpatch/v2 v2.4.0/go.mod h1:AH3dM2RI6uoBZxn3LVrfvJ3E0/9dG4cSrbuBJT4moAY=

pkg/apiserver/controllers/tunnelnode.go

@@ -2,13 +2,11 @@ package controllers
 
 import (
 	"context"
-	"encoding/pem"
 	"fmt"
 	"log/slog"
 	"net/http"
 	"time"
 
-	"github.com/MicahParks/jwkset"
 	"github.com/google/uuid"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -26,13 +24,12 @@ type TunnelNodeReconciler struct {
 
 	jwksHost              string
 	jwksPort              int
-	jwtPrivateKey         []byte
-	jwtPublicKey          []byte
+	jwtPrivateKeyPEM      []byte
+	jwtPublicKeyPEM       []byte
 	tokenRefreshThreshold time.Duration
 
 	validator *token.InMemoryValidator
 	issuer    *token.Issuer
-	jwkSet    *jwkset.MemoryJWKSet
 }
 
 func NewTunnelNodeReconciler(
@@ -47,10 +44,9 @@ func NewTunnelNodeReconciler(
 		Client:                c,
 		jwksHost:              jwksHost,
 		jwksPort:              jwksPort,
-		jwtPrivateKey:         jwtPrivateKey,
-		jwtPublicKey:          jwtPublicKey,
+		jwtPrivateKeyPEM:      jwtPrivateKey,
+		jwtPublicKeyPEM:       jwtPublicKey,
 		tokenRefreshThreshold: tokenRefreshThreshold,
-		jwkSet:                jwkset.NewMemoryStorage(),
 	}
 }
 
@@ -123,7 +119,7 @@ func (r *TunnelNodeReconciler) Reconcile(ctx context.Context, req reconcile.Requ
 			return ctrl.Result{}, fmt.Errorf("failed to parse UID as UUID: %w", err)
 		}
 
-		token, claims, err := r.issuer.IssueToken(subj, 2*r.tokenRefreshThreshold)
+		token, claims, err := r.issuer.IssueToken(subj.String(), 2*r.tokenRefreshThreshold)
 		if err != nil {
 			return ctrl.Result{}, fmt.Errorf("failed to issue token: %w", err)
 		}
@@ -146,67 +142,31 @@ func (r *TunnelNodeReconciler) Reconcile(ctx context.Context, req reconcile.Requ
 	return ctrl.Result{}, nil
 }
 
-// JWKSHandler returns an http.HandlerFunc that serves the JWKS at the
-// standard JWKS path.
-func (r *TunnelNodeReconciler) JWKSHandler() http.HandlerFunc {
-	return func(w http.ResponseWriter, req *http.Request) {
-		resp, err := r.jwkSet.JSONPublic(req.Context())
-		if err != nil {
-			slog.Error("Failed to get JWK Set JSON.", slog.Any("error", err))
-			w.WriteHeader(http.StatusInternalServerError)
-			return
-		}
-
-		w.Header().Set("Content-Type", "application/json")
-		_, _ = w.Write(resp)
-	}
-}
-
 func (r *TunnelNodeReconciler) SetupWithManager(ctx context.Context, mgr ctrl.Manager) error {
 	var err error
-	r.validator, err = token.NewInMemoryValidator(r.jwtPublicKey)
+	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.jwtPrivateKey)
+	r.issuer, err = token.NewIssuer(r.jwtPrivateKeyPEM)
 	if err != nil {
 		return fmt.Errorf("failed to create token issuer: %w", err)
 	}
 
-	// TODO(dsky): Implement key rotation.
-	pubKey, _ := pem.Decode(r.jwtPublicKey)
-	if pubKey == nil {
-		return fmt.Errorf("failed to decode private key")
-	}
-	key, err := jwkset.LoadX509KeyInfer(pubKey)
-	if err != nil {
-		return fmt.Errorf("failed to load X509 key: %w", err)
-	}
-	metadata := jwkset.JWKMetadataOptions{
-		KID: r.issuer.KeyID(),
-	}
-	jwk, err := jwkset.NewJWKFromKey(key, jwkset.JWKOptions{
-		Metadata: metadata,
-		Marshal: jwkset.JWKMarshalOptions{
-			Private: false,
-		},
-	})
-	if err != nil {
-		return fmt.Errorf("failed to create JWK: %w", err)
-	}
-	if err := r.jwkSet.KeyWrite(ctx, jwk); err != nil {
-		return fmt.Errorf("failed to write JWK: %w", err)
-	}
-
 	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)
+	if err != nil {
+		return fmt.Errorf("failed to create JWKS handler: %w", err)
+	}
+
 	mux := http.NewServeMux()
-	mux.HandleFunc(token.JWKSURI, r.JWKSHandler())
+	mux.HandleFunc(token.JWKSURI, jwksHandler)
 
 	server := &http.Server{
 		Addr:    fmt.Sprintf("%s:%d", r.jwksHost, r.jwksPort),

pkg/apiserver/controllers/tunnelnode_test.go

@@ -2,11 +2,6 @@ package controllers
 
 import (
 	"context"
-	"crypto/ecdsa"
-	"crypto/elliptic"
-	"crypto/rand"
-	"crypto/x509"
-	"encoding/pem"
 	"testing"
 	"time"
 
@@ -19,17 +14,31 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 
 	corev1alpha "github.com/apoxy-dev/apoxy-cli/api/core/v1alpha"
+	"github.com/apoxy-dev/apoxy-cli/pkg/cryptoutils"
 	"github.com/apoxy-dev/apoxy-cli/pkg/tunnel/token"
 )
 
 func TestTunnelNodeReconciler(t *testing.T) {
 	scheme := runtime.NewScheme()
 	require.NoError(t, corev1alpha.Install(scheme))
 
-	// Create a fake client with the registered scheme.
-	k8sClient := fake.NewClientBuilder().WithScheme(scheme).Build()
+	tunnelNode := &corev1alpha.TunnelNode{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-tunnelnode",
+			Namespace: "default",
+			UID:       types.UID(uuid.New().String()),
+		},
+	}
+
+	// Create a fake client with the registered scheme and the TunnelNode object.
+	k8sClient := fake.NewClientBuilder().WithScheme(scheme).
+		WithObjects(tunnelNode).
+		WithStatusSubresource(tunnelNode).
+		Build()
+
+	privKey, pubKey, err := cryptoutils.GenerateEllipticKeyPair()
+	require.NoError(t, err)
 
-	privKey, pubKey := generateKeyPair(t)
 	r := NewTunnelNodeReconciler(
 		k8sClient,
 		"localhost",
@@ -39,22 +48,9 @@ func TestTunnelNodeReconciler(t *testing.T) {
 		time.Minute,
 	)
 
-	var err error
-	r.validator, err = token.NewInMemoryValidator(r.jwtPublicKey)
+	r.validator, err = token.NewInMemoryValidator(r.jwtPublicKeyPEM)
 	require.NoError(t, err)
-	r.issuer, err = token.NewIssuer(r.jwtPrivateKey)
-	require.NoError(t, err)
-
-	tunnelNode := &corev1alpha.TunnelNode{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      "test-tunnelnode",
-			Namespace: "default",
-			UID:       types.UID(uuid.New().String()),
-		},
-	}
-
-	// Add the TunnelNode to the fake client.
-	err = k8sClient.Create(context.TODO(), tunnelNode)
+	r.issuer, err = token.NewIssuer(r.jwtPrivateKeyPEM)
 	require.NoError(t, err)
 
 	// Call the reconcile method.
@@ -66,26 +62,3 @@ func TestTunnelNodeReconciler(t *testing.T) {
 	})
 	require.NoError(t, err)
 }
-
-func generateKeyPair(t *testing.T) ([]byte, []byte) {
-	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
-	require.NoError(t, err)
-
-	privateKeyBytes, err := x509.MarshalPKCS8PrivateKey(privateKey)
-	require.NoError(t, err)
-
-	privateKeyPem := pem.EncodeToMemory(&pem.Block{
-		Type:  "PRIVATE KEY",
-		Bytes: privateKeyBytes,
-	})
-
-	publicKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
-	require.NoError(t, err)
-
-	pubKeyPem := pem.EncodeToMemory(&pem.Block{
-		Type:  "PUBLIC KEY",
-		Bytes: publicKeyBytes,
-	})
-
-	return privateKeyPem, pubKeyPem
-}