Some test code I wrote to play around with signing JWS objects

test_pktoken.go

package pktoken_test

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/sha256"
	_ "embed"
	"encoding/asn1"
	"encoding/json"
	"fmt"
	"math/big"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/openpubkey/openpubkey/pktoken"
	"github.com/openpubkey/openpubkey/pktoken/clientinstance"
	"github.com/openpubkey/openpubkey/pktoken/mocks"

	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jws"

	"github.com/openpubkey/openpubkey/util"
)

func TestPkToken(t *testing.T) {
	alg := jwa.ES256

	signingKey, err := util.GenKeyPair(alg)
	if err != nil {
		t.Fatal(err)
	}

	pkt, err := mocks.GenerateMockPKToken(signingKey, alg)
	if err != nil {
		t.Fatal(err)
	}

	testPkTokenMessageSigning(t, pkt, signingKey)
	testPkTokenSerialization(t, pkt)
}

func testPkTokenMessageSigning(t *testing.T, pkt *pktoken.PKToken, signingKey crypto.Signer) {
	// Create new OpenPubKey Signed Message (OSM)
	msg := "test message!"
	osm, err := pkt.NewSignedMessage([]byte(msg), signingKey)
	if err != nil {
		t.Fatal(err)
	}

	// Verify our OSM is valid
	payload, err := pkt.VerifySignedMessage(osm)
	if err != nil {
		t.Fatal(err)
	}

	if string(payload) != msg {
		t.Fatal("OSM payload did not match what we initially wrapped")
	}
}

func testPkTokenSerialization(t *testing.T, pkt *pktoken.PKToken) {
	// Test json serialization/deserialization
	pktJson, err := json.Marshal(pkt)
	if err != nil {
		t.Fatal(err)
	}
	fmt.Println(string(pktJson))

	var newPkt *pktoken.PKToken
	if err := json.Unmarshal(pktJson, &newPkt); err != nil {
		t.Fatal(err)
	}

	newPktJson, err := json.Marshal(newPkt)
	if err != nil {
		t.Fatal(err)
	}

	require.JSONEq(t, string(pktJson), string(newPktJson))
}

func TestPkTokenJwsUnchanged(t *testing.T) {

	payload := `{
		"aud": [
		  "also_me"
		],
		"email": "[email protected]",
		"exp": 1708641372,
		"iat": 1708554972,
		"iss": "me",
		"nonce": "iOqVQfpJsbt4gpcGUX0lJLT82bTm8PU1fwNghiGau0M",
		"sub": "1234567890"
	  }`

	// {"payload":"ewoJCSJhdWQiOiBbCgkJICAiYWxzb19tZSIKCQldLAoJCSJlbWFpbCI6ICJhcnRodXIuYWFyZHZhcmtAZXhhbXBsZS5jb20iLAoJCSJleHAiOiAxNzA4NjQxMzcyLAoJCSJpYXQiOiAxNzA4NTU0OTcyLAoJCSJpc3MiOiAibWUiLAoJCSJub25jZSI6ICJpT3FWUWZwSnNidDRncGNHVVgwbEpMVDgyYlRtOFBVMWZ3TmdoaUdhdTBNIiwKCQkic3ViIjogIjEyMzQ1Njc4OTAiCgkgIH0","signatures":[{"protected":"eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9","signature":"QqJFsnGTemdcy5MnsqhfTS4kpBUCW7UXeahGwLtT0C9XqKiLKZNMhU6GYQRoaqTPzrEEGVzutoeqFflKYh1n-FOmN92KwY0PgAwTFOGzQxJl_SkLyWSZQS-aqPWFnuVUjX-1fjtgxWEfkOFJ7NPE7vaIaAdsbjVIJNvZB9NJm5ECJ2PemkdF1jwIiQ3b4Mbw0lchrPydeznWQIdC5MPxLGUE6AiB6ljO4fOdE1Rgwv2SQk2vbueTXQNDYrsf14qe7oYAx-ZMrgW8WrMW7BGUh2pL-kkf19L708P9sJlvvwcPnYilCV9BLLbJFfHJqXslHGK2hTBy7QOiqoTrU9QHRw"},{"protected":"eyJhbGciOiJFUzI1NiIsInJ6IjoiODcyYzYzOTlmNDQwZDgwYThjMjg5MzVkOGRkODRkYTEzZWNkZmM4ZTk5YjNkZmJmOTJiZGYxYTMxMzNhMGI1ZSIsInR5cCI6IkNJQyIsInVwayI6eyJhbGciOiJFUzI1NiIsImNydiI6IlAtMjU2Iiwia3R5IjoiRUMiLCJ4IjoiMVV4Q3REQ2p5YjBiU3o5UDgxNXNNVHFHalNkRjJ1LXNZazBlZ3k0eWlncyIsInkiOiIwcVFuSGtPTE15UVk1V3ducGphRk8yVHpHQ3RxX25GZzEwZkkxNkxjZXhFIn19","signature":"V5PC7-HSGMnX6Q0qZufpWzQsNlYcNNfTjSUNrrlQzy2LBF3RJvQmHAwMCMDh5rBHGLOgjagW_wOGf0Ssf2Fu0g"}]}

	// Alphabetical order A to Z
	pheaderOpAz := `{"alg":"RS256","typ":"JWT"}`
	pheaderCicAz := `{"alg":"ES256","rz":"872c6399f440d80a8c28935d8dd84da13ecdfc8e99b3dfbf92bdf1a3133a0b5e","typ":"CIC","upk":{"alg":"ES256","crv":"P-256","kty":"EC","x":"1UxCtDCjyb0bSz9P815sMTqGjSdF2u-sYk0egy4yigs","y":"0qQnHkOLMyQY5WwnpjaFO2TzGCtq_nFg10fI16LcexE"}}`
	testCreatePKTokenFromJson(t, payload, pheaderOpAz, pheaderCicAz)

	// Reverse Alphabetical order Z to A
	pheaderOpZa := `{"typ":"JWT","alg":"RS256"}`
	pheaderCicZa := `{"upk":{"alg":"ES256","crv":"P-256","kty":"EC","x":"1UxCtDCjyb0bSz9P815sMTqGjSdF2u-sYk0egy4yigs","y":"0qQnHkOLMyQY5WwnpjaFO2TzGCtq_nFg10fI16LcexE"}, "typ":"CIC", "rz":"872c6399f440d80a8c28935d8dd84da13ecdfc8e99b3dfbf92bdf1a3133a0b5e", "alg":"ES256"}`
	testCreatePKTokenFromJson(t, payload, pheaderOpZa, pheaderCicZa)

	// Whitespace
	pheaderOpWs := `{  "alg":"RS256",     "typ":"JWT" }`
	pheaderCicWs := `{ "alg":    "ES256",  "rz": "872c6399f440d80a8c28935d8dd84da13ecdfc8e99b3dfbf92bdf1a3133a0b5e" , "typ":"CIC","upk":{"alg":"ES256","crv":"P-256","kty":"EC","x":"1UxCtDCjyb0bSz9P815sMTqGjSdF2u-sYk0egy4yigs","y":"0qQnHkOLMyQY5WwnpjaFO2TzGCtq_nFg10fI16LcexE"}}`

	testCreatePKTokenFromJson(t, payload, pheaderOpWs, pheaderCicWs)
}

func testCreatePKTokenFromJson(t *testing.T, payload string, opPheader string, cicPheader string) []byte {
	algOP := jwa.RS256
	signerOP, err := util.GenKeyPair(algOP)
	require.NoError(t, err)

	// Build ID Token
	buf := bytes.Buffer{}
	buf.WriteString(string(util.Base64EncodeForJWT([]byte(opPheader))))
	buf.WriteByte('.')
	encoded := util.Base64EncodeForJWT([]byte(payload))
	buf.WriteString(string(encoded))

	sigHash := sha256.Sum256(buf.Bytes())
	signature, err := signerOP.Sign(rand.Reader, sigHash[:], crypto.SHA256)
	require.NoError(t, err)

	buf.WriteByte('.')
	buf.WriteString(string(util.Base64EncodeForJWT(signature)))
	idtIn := buf.Bytes()
	require.NotNil(t, idtIn)

	_, err = jws.Verify(idtIn, jws.WithKey(algOP, signerOP.Public()))
	require.NoError(t, err)

	pkt := &pktoken.PKToken{}
	err = pkt.AddSignature(idtIn, pktoken.OIDC)
	require.NoError(t, err)

	idtOut, err := pkt.Compact(pkt.Op)

	require.EqualValues(t, string(idtIn), string(idtOut), "danger, signed values in ID Token being changed by PK Token")
	_, err = jws.Verify(idtOut, jws.WithKey(algOP, signerOP.Public()))
	require.NoError(t, err, "signature does not verify")

	// CIC CIC CIC CIC CIC
	// --- --- --- --- ---
	// CIC CIC CIC CIC CIC

	alg := jwa.ES256
	signerCic, err := jwk.ParseKey([]byte(`{"crv":"P-256","d":"U39LdEPqc7eXE_KS3F6mVjNFy7OF_r7SFED0fSpATF0","kty":"EC","x":"1UxCtDCjyb0bSz9P815sMTqGjSdF2u-sYk0egy4yigs","y":"0qQnHkOLMyQY5WwnpjaFO2TzGCtq_nFg10fI16LcexE"}`))
	require.NoError(t, err)

	jwkKey, err := jwk.PublicKeyOf(signerCic)
	require.NoError(t, err)
	err = jwkKey.Set(jwk.AlgorithmKey, alg)
	require.NoError(t, err)

	cic, err := clientinstance.NewClaims(jwkKey, map[string]any{})
	require.NoError(t, err)

	var rawkey interface{} // This is the raw key, like *rsa.PrivateKey or *ecdsa.PrivateKey
	err = signerCic.Raw(&rawkey)
	require.NoError(t, err)

	// Build CIC Token
	buf2 := bytes.Buffer{}
	buf2.WriteString(string(util.Base64EncodeForJWT([]byte(cicPheader))))
	buf2.WriteByte('.')
	encoded2 := util.Base64EncodeForJWT([]byte(payload))
	buf2.WriteString(string(encoded2))

	sigHash2 := sha256.Sum256(buf2.Bytes())
	sig2Der, err := rawkey.(*ecdsa.PrivateKey).Sign(rand.Reader, sigHash2[:], crypto.SHA256)
	require.NoError(t, err)
	signature2, err := RemoveDer(sig2Der)
	require.NoError(t, err)
	buf2.WriteByte('.')
	buf2.WriteString(string(util.Base64EncodeForJWT(signature2)))
	cicIn := buf2.Bytes()
	require.NotNil(t, cicIn)

	// // TODO: Sign cicPh after making it base64
	cicToken, err := cic.Sign(rawkey.(*ecdsa.PrivateKey), jwkKey.Algorithm(), idtOut)
	require.NoError(t, err)
	fmt.Println(string(cicToken))
	fmt.Println(string(cicIn))

	pkt.AddSignature(cicIn, pktoken.CIC)
	// zzz, err := pkt.Compact(pkt.Cic)
	require.NoError(t, err)
	pubkey, _ := signerCic.PublicKey()
	_, err = jws.Verify(cicIn, jws.WithKey(jwa.ES256, pubkey))

	err = pkt.VerifyCicSig()
	require.NoError(t, err)

	pkt2Json, err := pkt.MarshalJSON()
	require.NoError(t, err)
	require.NotNil(t, pkt2Json)
	fmt.Println(string(pkt2Json))

	var pkt2 pktoken.PKToken
	err = json.Unmarshal(pkt2Json, &pkt2)
	require.NoError(t, err)

	err = pkt2.VerifyCicSig()
	require.NoError(t, err)

	return idtOut
}

//go:embed test_jwk.json
var test_jwk []byte

// based on https://www.rfc-editor.org/rfc/rfc7638.html
func TestThumprintCalculation(t *testing.T) {
	fromRfc := "NzbLsXh8uDCcd-6MNwXF4W_7noWXFZAfHkxZsRGC9Xs"
	pub, err := jwk.ParseKey(test_jwk)
	if err != nil {
		t.Fatal(err)
	}
	thumb, err := pub.Thumbprint(crypto.SHA256)
	if err != nil {
		t.Fatal(err)
	}
	thumbEnc := util.Base64EncodeForJWT(thumb)
	if string(thumbEnc) != fromRfc {
		t.Fatalf("thumbprint %s did not match expected value %s", thumbEnc, fromRfc)
	}
}

func TestWeirdSigErr(t *testing.T) {
	payload := `{
		"aud": [
		  "also_me"
		],
		"email": "[email protected]",
		"exp": 1708641372,
		"iat": 1708554972,
		"iss": "me",
		"nonce": "iOqVQfpJsbt4gpcGUX0lJLT82bTm8PU1fwNghiGau0M",
		"sub": "1234567890"
	  }`

	// Set alg to ES256 in the pheader, but sign and verify with RSA256
	pheaderCicAz := `{"alg":"ES256","typ":"JWT"}`
	// pheaderCicAz := `{"alg":"RS256","typ":"JWT"}`

	// algOP := jwa.RS256
	algOP := jwa.ES256
	signerOP, err := util.GenKeyPair(algOP)
	require.NoError(t, err)

	cicIn := []byte(string(util.Base64EncodeForJWT([]byte(pheaderCicAz))) + "." + string(util.Base64EncodeForJWT([]byte(payload))))

	sigHash := sha256.Sum256(cicIn)
	signature, err := signerOP.Sign(rand.Reader, sigHash[:], crypto.SHA256)
	require.NoError(t, err)
	var sigNoDER noDER
	asn1.Unmarshal([]byte(signature), &sigNoDER)
	require.NoError(t, err)
	fmt.Println(sigNoDER)
	require.NoError(t, err)
	sssig := append(sigNoDER.R.Bytes(), sigNoDER.S.Bytes()...)
	cicIn = []byte(string(cicIn) + "." + string(util.Base64EncodeForJWT(sssig)))
	_, err = jws.Verify([]byte(cicIn), jws.WithKey(algOP, signerOP))
	require.NoError(t, err)
}

type noDER struct {
	R *big.Int
	S *big.Int
}

func RemoveDer(sigDer []byte) ([]byte, error) {
	var sigNoDER noDER
	_, err := asn1.Unmarshal([]byte(sigDer), &sigNoDER)
	if err != nil {
		return nil, err
	}
	return append(sigNoDER.R.Bytes(), sigNoDER.S.Bytes()...), nil
}