Headers missing in custom key provider when decrypting with jwe.WithKeyProvider

[ilya-korotya]

Describe the bug
I updated the jwx library from v3.0.10 to v3.0.12. I started using the new opt jwe.WithLegacyHeaderMerging(false) to prevent legacy header merging that was discussed in the PR. With this new options, I call the encryption function in this way:

	jweOpts := append([]jwe.EncryptOption{
		jwe.WithJSON(),
		jwe.WithContentEncryption(cea),
		jwe.WithProtectedHeaders(headers),
		jwe.WithLegacyHeaderMerging(false),
	}, withKeys...)
    ret, err := jwe.Encrypt(payload, jweOpts...)

As a result, I get a flattened token (because len(withKeys) is 1). Token:

{
  "ciphertext": "b3t1n6vUbELEkuZJ9_4faebSpy421MfucLLfHAoxGpPBptNn9sK2QidSbSX1KAQTDGFdWUM35m_B4wXTRP42JSrfFRAkLKfDtnFXxwaTN4eety0KUR9nxbXNXWPjrXD5uODkARszCofznKuYDJ67hjJQsPY16cj3c2_2a-L5eRmr2D5nagphrS4Jfd6kqzr7eQEJ7x01UX7z5jer_5TAjLhLrlZx2oHjH-OHGBjCplFlVayZJ1v1WzcGWL7ErILvLUcIuhVTdIJ0OCV5U_Qmm6Fx9069WWptBh23ELDMeSpU3QYqDnAjlhz-shrhoFIAJ7uREw9SAoCVV7rgK1_0T2m41q0UunAvXoH5q_IknAyaUmJAWTmuoxxaBA-GjYDLL0lEqK-yrKpCVX6p26Afh3RhJXCDrJIziQIcAILRoMqT8V12V4ni",
  "encrypted_key": "Vpv7bxGfT8fBhNQYvB2Cxpm2O4fEEqj7x8ZMDbIv5IrWC7Ad9a37Uw",
  "iv": "zwftaPgcs5FBGHKe",
  "protected": "eyJhbGciOiJFQ0RILUVTK0EyNTZLVyIsImVuYyI6IkEyNTZHQ00iLCJlcGsiOnsiY3J2IjoiUC0zODQiLCJrdHkiOiJFQyIsIngiOiJhT0ZoRV9rZk9lQW9xc2ZlNlZRelY0SnREbGhzMzZEN2dud1RZN1p4NmVmUXVkMmp6YzgzNlR0aF9lYi1xWm0yIiwieSI6Im9idVdNS1A2UEFEcENCRjJYMXZmdnZDMXRNUXBHRjl2ams4RFdOZ044NjF4LXRiNEpfUHhuWlJPRmxFM21MdFQifSwia2lkIjoiMTIzNDU2Nzg5IiwidHlwIjoiYXBwbGljYXRpb24vaWRlbjNjb21tLWVuY3J5cHRlZC1qc29uIn0",
  "tag": "-k7eGDDgoeRBFtwp02F3gg"
}

If I decode the protected field from base64, I get the following payload:

{
  "alg": "ECDH-ES+A256KW",
  "enc": "A256GCM",
  "epk": {
    "crv": "P-384",
    "kty": "EC",
    "x": "aOFhE_kfOeAoqsfe6VQzV4JtDlhs36D7gnwTY7Zx6efQud2jzc836Tth_eb-qZm2",
    "y": "obuWMKP6PADpCBF2X1vfvvC1tMQpGF9vjk8DWNgN861x-tb4J_PxnZROFlE3mLtT"
  },
  "kid": "123456789",
  "typ": "application/iden3comm-encrypted-json"
}

As you can see, the alg field was presented during encryption. But when I try to decrypt the token with the custom provider:

func Decrypt(envelope []byte, fn KeyResolutionFunc) ([]byte, error) {
	customKeyProvider := func(ctx context.Context, sink jwe.KeySink, r jwe.Recipient, msg *jwe.Message) error {
		alg, ok := r.Headers().Algorithm()
		if !ok || alg.String() == "" {
			return fmt.Errorf("recipient has no algorithm")
		}
		if !IsSupportedKeyEncryptionAlgorithm(alg.String()) {
			return fmt.Errorf("unsupported key encryption algorithm: %s", alg.String())
		}

		enc, ok := msg.ProtectedHeaders().ContentEncryption()
		if !ok || enc.String() == "" {
			return fmt.Errorf("message has no content encryption algorithm")
		}
		if !IsSupportedContentEncryptionAlgorithm(enc.String()) {
			return fmt.Errorf("unsupported content encryption algorithm: %s", enc.String())
		}

		kid, ok := r.Headers().KeyID()
		if !ok || kid == "" {
			return fmt.Errorf("recipient has no key ID")
		}
		decryptionKey, err := fn(kid)
		if err != nil {
			return err
		}
		sink.Key(alg, decryptionKey)
		return nil
	}

	payload, err := jwe.Decrypt(envelope, jwe.WithKeyProvider(jwe.KeyProviderFunc(customKeyProvider)))
	if err != nil {
		return nil, fmt.Errorf("%s: %w", err, ErrDecryptionKeyNotFound)
	}
	return payload, nil
}

I get an error saying the alg header was not provided:

jwe.Decrypt: failed to decrypt any of the recipients: key provider 0 failed: recipient has no algorithm: decryption key not found
        	            	failed to decrypt message

Here is an unit test to reproduce the issue:

func TestAlgInHeaders(t *testing.T) {
	token := `{
  "ciphertext": "b3t1n6vUbELEkuZJ9_4faebSpy421MfucLLfHAoxGpPBptNn9sK2QidSbSX1KAQTDGFdWUM35m_B4wXTRP42JSrfFRAkLKfDtnFXxwaTN4eety0KUR9nxbXNXWPjrXD5uODkARszCofznKuYDJ67hjJQsPY16cj3c2_2a-L5eRmr2D5nagphrS4Jfd6kqzr7eQEJ7x01UX7z5jer_5TAjLhLrlZx2oHjH-OHGBjCplFlVayZJ1v1WzcGWL7ErILvLUcIuhVTdIJ0OCV5U_Qmm6Fx9069WWptBh23ELDMeSpU3QYqDnAjlhz-shrhoFIAJ7uREw9SAoCVV7rgK1_0T2m41q0UunAvXoH5q_IknAyaUmJAWTmuoxxaBA-GjYDLL0lEqK-yrKpCVX6p26Afh3RhJXCDrJIziQIcAILRoMqT8V12V4ni",
  "encrypted_key": "Vpv7bxGfT8fBhNQYvB2Cxpm2O4fEEqj7x8ZMDbIv5IrWC7Ad9a37Uw",
  "iv": "zwftaPgcs5FBGHKe",
  "protected": "eyJhbGciOiJFQ0RILUVTK0EyNTZLVyIsImVuYyI6IkEyNTZHQ00iLCJlcGsiOnsiY3J2IjoiUC0zODQiLCJrdHkiOiJFQyIsIngiOiJhT0ZoRV9rZk9lQW9xc2ZlNlZRelY0SnREbGhzMzZEN2dud1RZN1p4NmVmUXVkMmp6YzgzNlR0aF9lYi1xWm0yIiwieSI6Im9idVdNS1A2UEFEcENCRjJYMXZmdnZDMXRNUXBHRjl2ams4RFdOZ044NjF4LXRiNEpfUHhuWlJPRmxFM21MdFQifSwia2lkIjoiMTIzNDU2Nzg5IiwidHlwIjoiYXBwbGljYXRpb24vaWRlbjNjb21tLWVuY3J5cHRlZC1qc29uIn0",
  "tag": "-k7eGDDgoeRBFtwp02F3gg"
}`

	var mapToken map[string]interface{}
	err := json.Unmarshal([]byte(token), &mapToken)
	require.NoError(t, err)

	// decode protected header
	protectedB64, ok := mapToken["protected"].(string)
	require.True(t, ok)
	protectedBytes, err := base64.RawURLEncoding.DecodeString(protectedB64)
	require.NoError(t, err)

	var protectedHeader map[string]interface{}
	err = json.Unmarshal(protectedBytes, &protectedHeader)
	require.NoError(t, err)
	require.NotEmpty(t, protectedHeader["alg"])

	checkIfAlgHeaderPresented := func(ctx context.Context, sink jwe.KeySink, r jwe.Recipient, msg *jwe.Message) error {
		recipientAlgHeader, ok := r.Headers().Algorithm()
		require.True(t, ok)
		require.NotEmpty(t, recipientAlgHeader.String())
		// code to resolve key would go here
		// headers for jwe.Recipient weren't merged with protected headers
		return nil
	}

	jwe.Decrypt([]byte(token),
		jwe.WithKeyProvider(
			jwe.KeyProviderFunc(
				checkIfAlgHeaderPresented)))
}

[lestrrat]

Sorry, it's bed time in this part of the world, so I'm just going to leave you with a random comment that I haven't verified...

The fixes that went in with 3.0.12 should have changed the code such that it saves everything in the protected headers now. Since r.Headers() returns the public headers, I'd expect it to be empty. So I'm kind of thinking alg could be in the protected headers, i.e. msg.ProtectedHeaders()? If not I'll take a look tomorrow

[lestrrat]

Yeah, I changed

  recipientAlgHeader, ok := r.Headers().Algorithm()

to

  recipientAlgHeader, ok := msg.ProtectedHeaders().Algorithm()

And it works. Yes, it's awkward, but I'd have to say this is on the spec itself.

The spec introduces different serialization formats, which doesn't specifically say where the "alg" parameter must reside. Furthermore, in compact mode the equivalent of "header" does not exist. In flattened mode it exists, so "alg" could reside in either "header" or "protected" headers. When serializing, this library chooses to use it as part of the AAD (i.e. protected headers) so that it's part of the integrity check. In the wild, who knows which the sender may choose to store this in, so you would have to check both. In full JSON mode, you can do pretty much any combination, and again, you would have to check both.

So I hope you can see it's hard to standardize on the library's side. The end user needs to adapt to the actual payload.

---

BTW, this is a nit, but I see you doing json.Unmarshal(payload, map) all the time, but you could just do jwe.Parse(payload) to get a full jwe.Message object -- then you won't have to fiddle with maps.

[ilya-korotya]

Makes sense to me. Closing the issue.