Often when writing programs and functions, one starts off with concrete types that solve the problem at hand. At some later time, due to emerging requirements or observed patterns, or just to improve code readability and reusability, we refactor to make our code more polymorphic. The importance of not breaking your API typically ranges from nice to have (e.g. minimises rework but not essential) to paramount (e.g. in a popular, foundational library). This post is a case study of a refactoring in the jose library demonstrating how type classes aid API stability, with a side helping of classy optics.
In the jose library, the verifyJWS function lets you verify a JSON Web Signature (JWS) object:
verifyJWS
:: ( HasAlgorithms a, HasValidationPolicy a
, AsError e, MonadError e m
, HasJWSHeader h, HasParams h
)
=> a
-> JWK
-> JWS h
-> m ()The function is applied to a configuration value, a JSON Web Key (JWK) to use for validation, and the JWS, and returns () on success otherwise throws an error.
A JWS can have multiple signatures, each by a different key. If an application requires all signatures to be valid, it is difficult to perform the correct validation using the existing verifyJWS function.
Unsurprisingly, someone raised an issue to address this. Specifically, it asks for a way to use a JWK Set instead of a single JWK, where the JWK Set contains all the keys that can be used for validation. JWK Set is defined by RFC 7515 as an array of JWKs. Its definition in jose is below (it is a newtype because it needs particular aeson instances).
newtype JWKSet = JWKSet [JWK]How can we support verification when the caller has a JWKSet? Do we create a new verification function that gets applied to JWKSet instead of JWK? I did not want multiple functions in the API for essentially the same thing. We could make the caller construct a singleton JWKSet, but changing the signature of verifyJWS would break the API, so we won't do that either.
At first, we could only validate with a single JWK at a time. Now, we want to be able to handle a JWKSet too. Are there other structures we might want to handle?
Of course there are. One might want to use a [JWK], NonEmpty JWK or some other kind of container. These containers (indeed, most containers) have instances for Foldable t, so we should try to use that abstraction. This is rather boring, but it's an additional valid use case, and we'll keep it in mind as we start the refactoring.
An immediate question is: can I redefine ``verifyJWS`` in terms of ``Foldable``? It is a nice idea, but the answer is no. First, Foldable t => t has kind (* -> *), but JWKSet has kind *. Turning JWKSet into a generic container doesn't make sense either, because, well it just isn't, according to the RFCs. Second, we still want our function to work with a plain old JWK, which also has kind *.
It's clear at this point that if we want to make verifyJWS polymorphic but avoid breaking the API, we have to define our own type class:
class JWKStore a where
keys :: ???What should the type of keys be? Conceptually, it would be applied to a JWKStore a => a and yield the JWK values contained within. For validation use case there is no need to be able to add or update keys (it is read-only), so Fold a JWK is a good fit for the requirement. A Fold is a read-only optic that that can retrieve multiple values. The beauty of optics is that they can be composed together and Fold is no different. The Fold type is provided by the lens library, along with a bunch of useful helper functions.
So our type class and instances are:
class JWKStore a where
keys :: Fold a JWK
instance JWKStore JWK where
keys = id
instance JWKStore JWKSet where
keys = folding (\(JWKSet xs) -> xs)
data JWKFoldable t = JWKFoldable { unJWKFoldable :: t JWK }
instance Foldable t => JWKStore (JWKFoldable t) where
keys = folding unJWKFoldable(We newtype the instance for generic containers to avoid the possibility of overlapping instances. Users of the library wishing to use a Foldable t => t JWK would have to wrap it with the JWKFoldable constructor. In fact, we do not include JWKFoldable or the associated instance in jose; it is just included here as an example.)
Next we update the type signature for verifyJWS:
verifyJWS
:: ( HasAlgorithms a, HasValidationPolicy a
, AsError e, MonadError e m
, HasJWSHeader h, HasParams h
, JWKStore k
)
=> a
-> k
-> JWS h
-> m ()Existing code applying verifyJWS to a JWK works without changes. The only internal change that was needed was to apply anyOf keys to the existing test. (anyOf is a function provided in lens that returns True if any target of a Fold satisfies a predicate.) The line:
validate s = verifySig p s k == Right Truebecame:
validate s = anyOf keys ((== Right True) . verifySig p s) kThe technique of using type classes to select optics is commonly called classy optics. The terms classy lenses and classy prisms are also used when referring to those kinds of optics.
Our goal was to support validation using a JWKSet without breaking the API or adding more functions. At this point we have accomplished that goal (elegantly). But we are not even close to done. I mentioned in passing that JWKSet and Foldable t => t JWK were the boring use cases. Let's talk about some interesting ones!
JWS signatures each have a header which, at minimum, indicates the algorithm used (the "alg" member). It can optionally contain other data including a Key ID ("kid"), thumbprint of an X.509 certificate containing the key used make the signature ("x5t@S256"), a JWK for the signing key ("jwk"), and so on. It is not a stretch to imagine that if your use case involves many signing keys, you might want to use data available in the signature header to speed up key lookup. Such schemes are common in PKI, e.g., X.509 certificates contain an Authority Key Identifier field, and certificate databases usually provide efficient lookup by key identifier.
So in addition to being able to enumerate keys, we want JWKStore instances to potentially be able to look up keys based on data in the JWS header. We can add another function to the type class to accomplish this, along with a sensible default implementation:
class JWKStore a where
keys :: Fold a JWK
keysFor :: (HasJWSHeader h) => h -> Fold a JWK
keysFor _ = keysNow, for example, we shall instantiate JWKStore at a data type based on HashMap. The keysFor function will efficiently search for a JWK based on the "kid" (Key ID) header parameter if it appears in the JWS header, otherwise it will fold over all the keys.
newtype KidMap = KidMap { getMap :: HashMap String JWK }
instance JWKStore KidMap where
keys = folding getMap
keyFor h = case preview (kid . _Just . param) h of
Just k -> folding (lookup k . getMap)
Nothing -> keysRecall the type of keysFor:
keysFor :: (HasJWSHeader h) => h -> Fold a JWKh has an explicit HasJWSHeader type class constraint, which allows the implementation to use any information available via that type class to decide what to do. For JWS we're basically done, but we have forgotten about JSON Web Encryption (JWE). The concept of looking up keys in a key database applies to JWE as well as JWS, but the HasJWSHeader constraint is not suitable when you have a JWE header.
Lucky for us, JWE headers and JWS headers consist of mostly the same fields, with the same types and semantics. So instead of having a type class constraint mentioning the kind of header, we can define a type class for every shared header parameter and constrain keysFor to all of them. We will use classy optics again (lenses this time). There are eleven header fields shared by JWS and JWE headers, but for brevity we'll pretend there are only three.
class HasAlg a where
alg :: Lens' a (HeaderParam JWA.JWS.Alg)
class HasKid a where
kid :: Lens' a (Maybe (HeaderParam String))
class HasX5tS256 a where
x5tS256 :: Lens' a (Maybe (HeaderParam Types.Base64SHA256))The class definitions are mundane, as are the instances for JWSHeader:
instance HasAlg JWSHeader where
alg f h@(JWSHeader { _jwsHeaderAlg = a }) =
fmap (\a' -> h { _jwsHeaderAlg = a' }) (f a)
instance HasKid JWSHeader where
kid f h@(JWSHeader { _jwsHeaderKid = a }) =
fmap (\a' -> h { _jwsHeaderKid = a' }) (f a)
instance HasX5tS256 JWSHeader where
x5tS256 f h@(JWSHeader { _jwsHeaderX5tS256 = a }) =
fmap (\a' -> h { _jwsHeaderX5tS256 = a' }) (f a)And finally, the updated keysFor type signature:
class JWKStore a where
keys :: Fold a JWK
keysFor
:: (HasAlg h, HasKid h, HasX5tS256)
=> h
-> Fold a JWK
keysFor _ = keysI would lie if I told you that the refactoring ended here. But it is enough for one blog post! Let's recap what we have covered in this post.
First, we discussed a requirement to generalise the jose library's JWS validation code, which because of its concrete type could only work with a single JWK, to encompass additional valid use cases regarding key storage. We introduced the JWKStore type class, which provides access to JWKs inside arbitrary data types. The refactor was a generalisation of the existing verifyJWS function, so existing code using our library continues to work without changes. After this, we added another function to JWKStore to allow implementors to support efficient key lookup. Finally, we observed that key databases are needed for JWE as well as JWS, and further generalised KeyStore to account for this.
Classy optics are an important part of the implementation resulting from this refactoring effort; they are employed in two different ways. One was to provide a read-only, composable interface to keys in the JWKStore. The other was to parameterise the key lookup function over the fields that are common to both JWS and JWE headers, providing the generality we need whilst improving readbility by making explicit what data can (and cannot) be used during key lookup.
One important principle the reader can draw from this case study, which applies to all programming no matter what language, is that when you have to refactor or augment some part of a program or library to encompass a new (or YAGNI'd) use case, it pays to have a good hard think about what other use cases exist, and how you can improve the generality (and therefore readability and reuse) of the code. (And you don't necessarily have to do the thinking all on your own: talking it over with colleagues or users can reveal things you might never have noticed.)
I mentioned above that there was more to this refactor, but actually there was not much more, so I will conclude with a brief discussion. The final change was to add a KeyUse parameter to keysFor, so that calling code can indicate what operation it wants to use the key(s) for, and JWKStore implementations can search or filter keys accordingly. I am not 100% sold on this approach, and I can forsee that changes or additions to JWKStore may be needed in the future, but it suffices for now. I hope you enjoyed this refactoring case study.