ecdh-1pu-jwt.go

package main

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"hash"
	"math"

	"github.com/square/go-jose/v3"
	"github.com/square/go-jose/v3/jwt"
)

var (
	// Used as external signing key
	notaryKey         = mustJWK([]byte(`{"kty":"EC","d": "Dy0_L1wKbb3ogotdHFn1J4cA005vVs7sU7Qu3GRIUAx0kdgWrLihmj4eQ6YxJdQL","use": "sig","crv": "P-384","x": "Dv3FuXWopmk0yUqtz9N4fynOAwdOesmYGhyC3pPQP-aYThVlLuLEBF7mOPk_hJGJ","y": "dHevDqS6F2123nDC61NypUJ8t9-nSa2KjpgfdhiPmRGpHHrGH-sd209qX8_e8iMy",	"alg": "ES384"}`))
	// Static Alice Private/Public Key
	aliceStaticJWK    = mustJWK([]byte(`{"kty":"EC", "crv":"P-256", "x":"WKn-ZIGevcwGIyyrzFoZNBdaq9_TsqzGl96oc0CWuis", "y":"y77t-RvAHRKTsSGdIYUfweuOvwrvDD-Q3Hv5J0fSKbE", "d":"Hndv7ZZjs_ke8o9zXYo3iq-Yr8SewI5vrqd0pAvEPqg"}`))
	// Ephemeral Alice Private/Public Key
	aliceEphemeralJWK = mustJWK([]byte(`{"kty":"EC", "crv":"P-256",	"x":"gI0GAILBdu7T53akrFmMyGcsF3n5dO7MmwNBHKW5SV0", "y":"SLW_xSffzlPWrHEVI30DHM_4egVwt3NQqeUD7nMFpps", "d":"0_NxaRPUMQoAJt50Gz8YiTr8gRTwyEaCumd-MToTmIo"}`))
	// Static Bob Public Key
	bobStaticJWK      = mustJWK([]byte(`{"kty":"EC", "crv":"P-256", "x":"weNJy2HscCSM6AEDTDg04biOvhFhyyWvOHQfeF_PxMQ", "y":"e8lnCO-AlStT-NJVX-crhB7QRYhiix03illJOVAOyck", "d":"VEmDZpDXXK8p8N0Cndsxs924q6nS1RXFASRl6BfUqdw"}`))
)

// mustJWK decodes JWK encoded keys and panic if decode fail.
func mustJWK(data []byte) *jose.JSONWebKey {
	var key jose.JSONWebKey
	if err := json.NewDecoder(bytes.NewReader(data)).Decode(&key); err != nil {
		panic(err)
	}

	return &key
}

// uint32ToBytes encodes an uint32 to bigendian byte array
func uint32ToBytes(value uint32) []byte {
	result := make([]byte, 4)
	binary.BigEndian.PutUint32(result, value)

	return result
}

// lengthPrefixedArray retrusn a byte array with a big endian encoded uint32 
// of the value length (TLV without T)
func lengthPrefixedArray(value []byte) []byte {
	result := make([]byte, 4)
	binary.BigEndian.PutUint32(result, uint32(len(value)))

	return append(result, value...)
}

// KDF implements NIST SP 800-56A Concatenation Key Derivation Function.
// Inspired by https://github.com/hackme3/EncryptPad/blob/976da72751b4511a6cad40509d50785cd133e6ee/deps/botan/src/lib/kdf/sp800_56a/sp800_56a.cpp
func kdf(h hash.Hash, sharedSecret, otherInfo []byte, keyDataLen int) []byte {
	keyLenBytes := keyDataLen >> 3
	reps := int(math.Ceil(float64(keyLenBytes) / float64(h.Size())))
	if reps >= 1<<32 {
		panic("SP800-56A KDF requested output too large")
	}

	dk := make([]byte, 0, keyLenBytes)
	for counter := 1; counter <= reps; counter++ {
		h.Reset()
		h.Write(uint32ToBytes(uint32(counter)))
		h.Write(sharedSecret)
		h.Write(otherInfo)
		dk = h.Sum(dk)
	}

	return dk[:keyLenBytes]
}

// Used for computing ECDH-1PU
// https://tools.ietf.org/html/draft-madden-jose-ecdh-1pu-03
func main() {
	var (
		alice                 = []byte("Alice")
		aliceStaticPrivateKey = aliceStaticJWK.Key.(*ecdsa.PrivateKey)
		aliceEphPrivateKey    = aliceEphemeralJWK.Key.(*ecdsa.PrivateKey)

		bob                = []byte("Bob")
		bobStaticPublicKey = bobStaticJWK.Public().Key.(*ecdsa.PublicKey)
	)

	// Compute Ze : ECDH Agreement - Alice Ephemeral private key / Bob Static public key
	Ze, _ := aliceEphPrivateKey.Curve.ScalarMult(bobStaticPublicKey.X, bobStaticPublicKey.Y, aliceEphPrivateKey.D.Bytes())
	fmt.Printf("Ze: %x\n", Ze)

	// Compute Zs : ECDH Agreement - Alice Static private key / Bob Static public key
	Zs, _ := aliceStaticPrivateKey.Curve.ScalarMult(bobStaticPublicKey.X, bobStaticPublicKey.Y, aliceStaticPrivateKey.D.Bytes())
	fmt.Printf("Zs: %x\n", Zs)

	// Compute Z (shared secret)
	Z := Ze.Bytes()
	Z = append(Z, Zs.Bytes()...)
	fmt.Printf("Z (sharedsecret): %x\n", Z)

	// 256 bit for AES256GCM
	keyDataLen := 256

	// Prepare KDF info components
	algorithmID := lengthPrefixedArray([]byte("A256GCM"))
	partyUInfo := lengthPrefixedArray(alice)
	partyVInfo := lengthPrefixedArray(bob)
	suppPubInfo := uint32ToBytes(uint32(keyDataLen))

	// Assembling (AlgorithmID || PartyUInfo || PartyVInfo || SubPubInfo)
	info := append(algorithmID, partyUInfo...)
	info = append(info, partyVInfo...)
	info = append(info, suppPubInfo...)
	fmt.Printf("info: %x\n", info)

	outKey := kdf(sha256.New(), Z, info, keyDataLen)
	fmt.Printf("Key: %x\n", outKey)
	fmt.Printf("Key (base64url): %s\n", base64.RawURLEncoding.EncodeToString(outKey))

	// Create sign-crypted token
	signer, err := jose.NewSigner(jose.SigningKey{
		Algorithm: jose.ES256,
		Key:       aliceStaticJWK,
	}, &jose.SignerOptions{
		ExtraHeaders: map[jose.HeaderKey]interface{}{
			jose.HeaderContentType: "JWT",
		},
	})
	if err != nil {
		panic(err)
	}

	encrypter, err := jose.NewEncrypter(jose.A256GCM, jose.Recipient{
		Algorithm: jose.DIRECT,
		Key:       outKey,
	}, &jose.EncrypterOptions{
		ExtraHeaders: map[jose.HeaderKey]interface{}{
			jose.HeaderContentType: jose.ContentType("JWT"),
			"alg":                  "ECDH-1PU",
			"apu":                  base64.RawURLEncoding.EncodeToString(alice),
			"apv":                  base64.RawURLEncoding.EncodeToString(bob),
			"epk":                  aliceEphemeralJWK.Public(),
		},
	})
	if err != nil {
		panic(err)
	}

	token, err := jwt.SignedAndEncrypted(signer, encrypter).Claims(&jwt.Claims{
		ID:      "12345",
		Subject: fmt.Sprintf("%s", alice),
	}).CompactSerialize()
	if err != nil {
		panic(err)
	}

	fmt.Printf("SignedEncrypted Token: %s\n", token)

	// Check JWT

	_, err = jwt.ParseSignedAndEncrypted(token)
	if err != nil {
		panic(err)
	}
}

Awesome!