jwk: RegisterKeyImporter takes KeyImporter, not a typed function (#2152)

* jwk: RegisterKeyImporter takes KeyImporter, not a typed function

* jwk: collapse KeyImportFunc to a single generic name

* jwk: KeyImporter generic interface, T inferable at registration

* jwk: cover struct-shaped KeyImporter[T] dispatch path

Author: lestrrat

MIGRATION.md

@@ -399,7 +399,7 @@ import _ "github.com/jwx-go/asmbase64/v4" // registration only
 ### Recipe 10: Custom Key Importer
 
 ```go
-// Before
+// Before (v3)
 jwk.RegisterKeyImporter(&myKeyType{}, jwk.KeyImportFunc(func(raw any) (jwk.Key, error) {
     src, ok := raw.(*myKeyType)
     if !ok {
@@ -408,10 +408,14 @@ jwk.RegisterKeyImporter(&myKeyType{}, jwk.KeyImportFunc(func(raw any) (jwk.Key,
     // ... convert
 }))
 
-// After
-jwk.RegisterKeyImporter(func(src *myKeyType) (jwk.Key, error) {
-    // ... convert — type parameter inferred, no assertion needed
-})
+// After (v4): RegisterKeyImporter takes a jwk.KeyImporter[T]; use
+// jwk.KeyImportFunc[T] to adapt a typed function. The outer type
+// parameter is inferred from the adapter's typed Import method.
+jwk.RegisterKeyImporter(jwk.KeyImportFunc[*myKeyType](func(src *myKeyType) (jwk.Key, error) {
+    // The interface is typed, so src is *myKeyType in the body —
+    // no manual assertion.
+    // ... convert
+}))
 ```
 
 ### Recipe 11: Iterating Over Sets and Tokens

docs/04-jwk.md

@@ -759,7 +759,7 @@ registration points:
 - `KeyParser` — converts a JSON payload into a `jwk.Key` using the
   probe's hints. Register with [`jwk.RegisterKeyParser`](https://pkg.go.dev/github.com/lestrrat-go/jwx/v4/jwk#RegisterKeyParser).
 - `KeyImporter[T]` — converts a raw Go crypto value into a `jwk.Key`.
-  Register with [`jwk.RegisterKeyImporter[T]`](https://pkg.go.dev/github.com/lestrrat-go/jwx/v4/jwk#RegisterKeyImporter).
+  Register with [`jwk.RegisterKeyImporter`](https://pkg.go.dev/github.com/lestrrat-go/jwx/v4/jwk#RegisterKeyImporter).
 - `KeyExporter` — converts a `jwk.Key` back into a raw Go value.
   Register with [`jwk.RegisterKeyExporter`](https://pkg.go.dev/github.com/lestrrat-go/jwx/v4/jwk#RegisterKeyExporter),
   keyed by `jwk.KeyKind` (case-insensitive — usually
@@ -821,25 +821,46 @@ func (MyKeyParser) ParseKey(probe *jwk.KeyProbe, unmarshaler jwk.KeyUnmarshaler,
 
 ### How import / export dispatch
 
-`jwk.Import[T](raw)` looks up a `KeyImporter[T]` registered for the
+`jwk.Import[T](raw)` looks up the `KeyImporter[T]` registered for the
 exact Go type of `raw`. Pre-registered importers exist for the standard
 crypto types (`*rsa.PrivateKey`, `*ecdsa.PublicKey`, `ed25519.PrivateKey`,
-etc.); add your own:
+etc.); register your own with `jwk.RegisterKeyImporter(KeyImporter[T])`.
+The interface is generic — its `Import(raw T) (Key, error)` method is
+typed — so the registration's type parameter is inferred from whatever
+you pass in. Use `jwk.KeyImportFunc[T]` to adapt a typed function:
 
 ```go
 func init() {
-    if err := jwk.RegisterKeyImporter[*mypkg.SuperSecretKey](
-        func(raw *mypkg.SuperSecretKey) (jwk.Key, error) {
-            // raw is already typed — no `any`-cast or ContinueError
-            // needed. v4's import surface is one importer per Go type.
-            return mypkg.NewSuperSecretJWK(raw), nil
-        },
+    if err := jwk.RegisterKeyImporter(
+        jwk.KeyImportFunc[*mypkg.SuperSecretKey](
+            func(raw *mypkg.SuperSecretKey) (jwk.Key, error) {
+                // The interface is typed, so raw is *mypkg.SuperSecretKey
+                // — no any-cast needed in the body.
+                return mypkg.NewSuperSecretJWK(raw), nil
+            },
+        ),
     ); err != nil {
         panic(err)
     }
 }
 ```
 
+For an importer with non-trivial state, implement the `KeyImporter[T]`
+interface on your own type and pass an instance directly:
+
+```go
+type mySuperSecretImporter struct { /* config, caches, ... */ }
+
+// Typed Import method makes mySuperSecretImporter satisfy
+// jwk.KeyImporter[*mypkg.SuperSecretKey].
+func (i *mySuperSecretImporter) Import(raw *mypkg.SuperSecretKey) (jwk.Key, error) {
+    return i.convert(raw)
+}
+
+// Registration: T inferred from the Import method's parameter type.
+jwk.RegisterKeyImporter(&mySuperSecretImporter{...})
+```
+
 `jwk.Export[T](key)` runs a `KeyExporter` keyed by `jwk.KeyKind`
 (matched case-insensitively against `key.KeyType()`):
 

docs/design/v4-generics.md

@@ -21,16 +21,17 @@ func myFunc(src any) (jwk.Key, error) {
 
 **v4:**
 ```go
-jwk.RegisterKeyImporter(myFunc)
-// myFunc accepts the concrete type — no type switch needed
+jwk.RegisterKeyImporter(jwk.KeyImportFunc[*rsa.PrivateKey](myFunc))
+// myFunc accepts the concrete type. The outer T is inferred from the
+// adapter's typed Import method.
 func myFunc(src *rsa.PrivateKey) (jwk.Key, error) {
     ...
 }
 ```
 
-Implementation: `RegisterKeyImporter[T any](fn func(T) (Key, error))` uses `reflect.TypeFor[T]()` to derive the map key at registration time. The runtime dispatch in `Import()` still uses `reflect.TypeOf(raw)` — this is unavoidable since the input type is only known at call time.
+Implementation: `KeyImporter[T any]` is a generic interface whose `Import(raw T) (Key, error)` method is itself typed; `KeyImportFunc[T any]` is a `func(T) (Key, error)` that satisfies it. `RegisterKeyImporter[T any](ki KeyImporter[T]) error` uses `reflect.TypeFor[T]()` to derive the map key at registration time and stores a closure that re-types `any` → T at dispatch. The runtime lookup in `Import()` still keys by `reflect.TypeOf(raw)`.
 
-Callers holding a full `KeyImporter` value can register it through the same entry point by passing `imp.Import` as the function: `RegisterKeyImporter[T](imp.Import)`.
+The earlier v4.0.0 release shipped `RegisterKeyImporter[T](fn func(T) (Key, error))` — accepting a typed function rather than a `KeyImporter`, with the public `KeyImporter` interface and `KeyImportFunc` adapter exposed but never reachable from registration. The current shape restores the interface as the canonical registration value, makes the interface generic so type-parameter inference covers the registration call, and makes `KeyImportFunc[T]` a generic typed-function adapter that satisfies it.
 
 ### Type-safe generic accessors
 
@@ -66,10 +67,10 @@ jwk.RegisterKeyImporter(&myKey{}, jwk.KeyImportFunc(func(src any) (jwk.Key, erro
     return doImport(k)
 }))
 
-// v4: type-safe, no zero-value discriminator needed
-jwk.RegisterKeyImporter(func(k *myKey) (jwk.Key, error) {
+// v4: pass a KeyImporter[T]; T inferred from the adapter's typed Import.
+jwk.RegisterKeyImporter(jwk.KeyImportFunc[*myKey](func(k *myKey) (jwk.Key, error) {
     return doImport(k)
-})
+}))
 ```
 
 ```go

jwk/convert.go

@@ -22,7 +22,14 @@ import (
 // A converter that converts from a raw key to a `jwk.Key` is called a KeyImporter.
 // A converter that converts from a `jwk.Key` to a raw key is called a KeyExporter.
 
-var keyImporters = make(map[reflect.Type]KeyImporter)
+// keyImporters stores per-Go-type closures that adapt a typed
+// [KeyImporter] (parameterized by the same Go type that keys this map)
+// down to a non-generic any-taking dispatch shape. The closure carries
+// T and performs the raw.(T) type-assertion before delegating to the
+// underlying typed Import method; the assertion is safe because the
+// dispatch in [convertRawKey] keys this entry by reflect.TypeOf(raw)
+// matching reflect.TypeFor[T]().
+var keyImporters = make(map[reflect.Type]func(any) (Key, error))
 
 // builtinImporterTypes is the set of raw key types whose importers
 // are owned by the jwk package and cannot be overridden by callers.
@@ -64,12 +71,14 @@ var keyExporters = make(map[KeyKind][]KeyExporter)
 var muKeyImporters sync.RWMutex
 var muKeyExporters sync.RWMutex
 
-// RegisterKeyImporter registers a typed import function for converting
-// raw keys of type T to jwk.Key. The type is derived from the type
-// parameter, eliminating the need to pass a zero value.
+// RegisterKeyImporter registers a [KeyImporter] for the raw key type T.
+// The type parameter is normally inferred from the importer's typed
+// Import method; pass it explicitly only when inference would fail
+// (e.g. when supplying a [KeyImportFunc] without arguments).
 //
-// When `jwk.Import()` is called, the library looks up the appropriate
-// importer for the given raw key type (via `reflect`) and executes it.
+// When [Import] is called, the library looks up the importer for the
+// raw value's runtime type (via reflect) and invokes its Import method
+// through a closure that handles the any → T re-typing.
 //
 // Importer dispatch is single-valued per Go type T: there is exactly
 // one importer per reflect.Type. Subsequent registrations for the same
@@ -93,7 +102,12 @@ var muKeyExporters sync.RWMutex
 //
 // Extension modules calling this from init() must check the returned
 // error and panic on failure.
-func RegisterKeyImporter[T any](fn func(T) (Key, error)) error {
+//
+// To register from a typed function (the common case for extensions),
+// wrap it with [KeyImportFunc]:
+//
+//	jwk.RegisterKeyImporter(jwk.KeyImportFunc[*mypkg.Key](importMyKey))
+func RegisterKeyImporter[T any](ki KeyImporter[T]) error {
 	muKeyImporters.Lock()
 	defer muKeyImporters.Unlock()
 	t := reflect.TypeFor[T]()
@@ -103,7 +117,13 @@ func RegisterKeyImporter[T any](fn func(T) (Key, error)) error {
 	if _, exists := keyImporters[t]; exists {
 		return fmt.Errorf(`jwk.RegisterKeyImporter: an importer for %s is already registered; call jwk.UnregisterKeyImporter[%s]() first if you need to replace it`, t, t)
 	}
-	keyImporters[t] = &keyImportAdapter[T]{fn: fn}
+	keyImporters[t] = func(raw any) (Key, error) {
+		// Safe: dispatch in convertRawKey keys this entry by
+		// reflect.TypeOf(raw) matching reflect.TypeFor[T](), so the
+		// assertion cannot fail in the lookup path.
+		//nolint:forcetypeassert
+		return ki.Import(raw.(T))
+	}
 	return nil
 }
 
@@ -183,33 +203,36 @@ func RegisterKeyExporter(ident KeyKind, conv KeyExporter) error {
 	return nil
 }
 
-// KeyImporter is used to convert from a raw key to a `jwk.Key`. mneumonic: from the PoV of the `jwk.Key`,
-// we're _importing_ a raw key.
-type KeyImporter interface {
-	// Import takes the raw key to be converted, and returns a `jwk.Key` or an error if the conversion fails.
-	Import(any) (Key, error)
-}
-
-// KeyImportFunc is a convenience type to implement KeyImporter as a function.
-type KeyImportFunc func(any) (Key, error)
+// KeyImporter converts a raw key of Go type T into a [jwk.Key]. The
+// type parameter T is the raw-key type the importer handles. From the
+// point of view of the `jwk.Key`, we're _importing_ a raw key.
+//
+// Implementations are registered with [RegisterKeyImporter]; T flows
+// through to [reflect.TypeFor] at registration time to key the import
+// dispatch table.
+type KeyImporter[T any] interface {
+	// Import takes the raw key to be converted, and returns a
+	// [jwk.Key] or an error if the conversion fails.
+	Import(raw T) (Key, error)
+}
+
+// KeyImportFunc is the typed-function adapter that satisfies
+// [KeyImporter] for type T. The Import method just calls f directly —
+// no internal type assertion is needed because [KeyImporter]'s Import
+// is itself typed.
+//
+// This is the canonical way to register a typed import function:
+//
+//	jwk.RegisterKeyImporter(jwk.KeyImportFunc[*mypkg.Key](importMyKey))
+//
+// For an importer with non-trivial state, implement [KeyImporter] on
+// your own type and pass an instance of it directly.
+type KeyImportFunc[T any] func(T) (Key, error)
 
-func (f KeyImportFunc) Import(raw any) (Key, error) {
+func (f KeyImportFunc[T]) Import(raw T) (Key, error) {
 	return f(raw)
 }
 
-// keyImportAdapter wraps a typed function into a KeyImporter.
-type keyImportAdapter[T any] struct {
-	fn func(T) (Key, error)
-}
-
-func (a *keyImportAdapter[T]) Import(raw any) (Key, error) {
-	v, ok := raw.(T)
-	if !ok {
-		return nil, fmt.Errorf(`cannot convert key type '%T' to %T`, raw, *new(T))
-	}
-	return a.fn(v)
-}
-
 // KeyExporter is used to convert from a `jwk.Key` to a raw key. From the PoV of the `jwk.Key`,
 // we're _exporting_ it to a raw key.
 type KeyExporter interface {
@@ -269,7 +292,13 @@ func init() {
 func registerBuiltinKeyImporter[T any](fn func(T) (Key, error)) {
 	t := reflect.TypeFor[T]()
 	builtinImporterTypes[t] = struct{}{}
-	keyImporters[t] = &keyImportAdapter[T]{fn: fn}
+	keyImporters[t] = func(raw any) (Key, error) {
+		// Safe: dispatch keys this entry by reflect.TypeOf(raw)
+		// matching reflect.TypeFor[T](); the assertion cannot fail
+		// when the closure is invoked from convertRawKey.
+		//nolint:forcetypeassert
+		return fn(raw.(T))
+	}
 }
 
 // panicOnRegistrationError converts a non-nil error returned by a Register*