Add url and kid to jws. (#4340)

* Add url and kid to jws.

This will be required in order to implement the latest ACME spec, which uses
these protected header fields.

* Add comments and fix lint.

* Enforce mutual exclusivity of jwk and kid.

2 files changed, 35 insertions, 11 deletions

diff --git a/acme/acme/jws.py b/acme/acme/jws.py
index 54bc26d94..79e96edcb 100644
--- a/acme/acme/jws.py
+++ b/acme/acme/jws.py
@@ -1,14 +1,18 @@
-"""ACME JOSE JWS."""
+"""ACME-specific JWS.
+
+The JWS implementation in acme.jose only implements the base JOSE standard. In
+order to support the new header fields defined in ACME, this module defines some
+ACME-specific classes that layer on top of acme.jose.
+"""
 from acme import jose
 
 
 class Header(jose.Header):
-    """ACME JOSE Header.
-
-    .. todo:: Implement ``acmePath``.
-
+    """ACME-specific JOSE Header. Implements nonce, kid, and url.
     """
     nonce = jose.Field('nonce', omitempty=True, encoder=jose.encode_b64jose)
+    kid = jose.Field('kid', omitempty=True)
+    url = jose.Field('url', omitempty=True)
 
     @nonce.decoder
     def nonce(value):  # pylint: disable=missing-docstring,no-self-argument
@@ -20,7 +24,7 @@ class Header(jose.Header):
 
 
 class Signature(jose.Signature):
-    """ACME Signature."""
+    """ACME-specific Signature. Uses ACME-specific Header for customer fields."""
     __slots__ = jose.Signature._orig_slots  # pylint: disable=no-member
 
     # TODO: decoder/encoder should accept cls? Otherwise, subclassing
@@ -34,11 +38,17 @@ class Signature(jose.Signature):
 
 
 class JWS(jose.JWS):
-    """ACME JWS."""
+    """ACME-specific JWS. Includes none, url, and kid in protected header."""
     signature_cls = Signature
     __slots__ = jose.JWS._orig_slots  # pylint: disable=no-member
 
     @classmethod
-    def sign(cls, payload, key, alg, nonce):  # pylint: disable=arguments-differ
+    # pylint: disable=arguments-differ,too-many-arguments
+    def sign(cls, payload, key, alg, nonce, url=None, kid=None):
+        # Per ACME spec, jwk and kid are mutually exclusive, so only include a
+        # jwk field if kid is not provided.
+        include_jwk = kid is None
         return super(JWS, cls).sign(payload, key=key, alg=alg,
-                                    protect=frozenset(['nonce']), nonce=nonce)
+                                    protect=frozenset(['nonce', 'url', 'kid']),
+                                    nonce=nonce, url=url, kid=kid,
+                                    include_jwk=include_jwk)
diff --git a/acme/acme/jws_test.py b/acme/acme/jws_test.py
index e8f8e871a..62409227e 100644
--- a/acme/acme/jws_test.py
+++ b/acme/acme/jws_test.py
@@ -37,16 +37,30 @@ class JWSTest(unittest.TestCase):
         self.privkey = KEY
         self.pubkey = self.privkey.public_key()
         self.nonce = jose.b64encode(b'Nonce')
+        self.url = 'hi'
+        self.kid = 'baaaaa'
 
-    def test_it(self):
+    def test_kid_serialize(self):
         from acme.jws import JWS
         jws = JWS.sign(payload=b'foo', key=self.privkey,
-                       alg=jose.RS256, nonce=self.nonce)
+                       alg=jose.RS256, nonce=self.nonce,
+                       url=self.url, kid=self.kid)
         self.assertEqual(jws.signature.combined.nonce, self.nonce)
+        self.assertEqual(jws.signature.combined.url, self.url)
+        self.assertEqual(jws.signature.combined.kid, self.kid)
+        self.assertEqual(jws.signature.combined.jwk, None)
         # TODO: check that nonce is in protected header
 
         self.assertEqual(jws, JWS.from_json(jws.to_json()))
 
+    def test_jwk_serialize(self):
+        from acme.jws import JWS
+        jws = JWS.sign(payload=b'foo', key=self.privkey,
+                       alg=jose.RS256, nonce=self.nonce,
+                       url=self.url)
+        self.assertEqual(jws.signature.combined.kid, None)
+        self.assertEqual(jws.signature.combined.jwk, self.pubkey)
+
 
 if __name__ == '__main__':
     unittest.main()  # pragma: no cover