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 True
became:
validate s = anyOf keys ((== Right True) . verifySig p s) k
The 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 _ = keys
Now, 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 -> keys
Recall the type of keysFor
:
keysFor :: (HasJWSHeader h) => h -> Fold a JWK
h
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 _ = keys
I 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.