Sign and Verify TPM based JWT using github.com/lestrrat-go/jwx

Using a crypto.Signer from github.com/salrashid123/tpmsigner

Following using a swtpm

rm -rf /tmp/myvtpm && mkdir /tmp/myvtpm
swtpm_setup --tpmstate /tmp/myvtpm --tpm2 --create-ek-cert
swtpm socket --tpmstate dir=/tmp/myvtpm --tpm2 --server type=tcp,port=2321 --ctrl type=tcp,port=2322 --flags not-need-init,startup-clear --log level=2

export TPM2TOOLS_TCTI="swtpm:port=2321"

### create an rsassa key encoded as a PEM key

printf '\x00\x00' > unique.dat
tpm2_createprimary -C o -G ecc  -g sha256  -c primary.ctx -a "fixedtpm|fixedparent|sensitivedataorigin|userwithauth|noda|restricted|decrypt" -u unique.dat

tpm2_createprimary -C o -G rsa2048:aes128cfb -g sha256 -c primary.ctx -a 'restricted|decrypt|fixedtpm|fixedparent|sensitivedataorigin|userwithauth|noda'
tpm2_create -G rsa2048:rsassa:null -g sha256 -u key.pub -r key.priv -C primary.ctx
tpm2_flushcontext -t && tpm2_flushcontext -s && tpm2_flushcontext -l
tpm2_load -C primary.ctx -u key.pub -r key.priv -c key.ctx
tpm2_encodeobject -C primary.ctx -u key.pub -r key.priv -o key.pem

---

package main

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"slices"

	keyfile "github.com/foxboron/go-tpm-keyfiles"
	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpm2/transport"
	"github.com/google/go-tpm/tpmutil"
	"github.com/lestrrat-go/jwx/v3/jwa"
	"github.com/lestrrat-go/jwx/v3/jwt"
	"github.com/salrashid123/tpmsigner"
)

var (
	tpmPath = flag.String("tpm-path", "127.0.0.1:2321", "Path to the TPM device (character device or a Unix socket).")
	pemFile = flag.String("pemFile", "/tmp/key.pem", "KeyFile in PEM format")
)

var TPMDEVICES = []string{"/dev/tpm0", "/dev/tpmrm0"}

func openTPM(path string) (io.ReadWriteCloser, error) {
	if slices.Contains(TPMDEVICES, path) {
		return tpmutil.OpenTPM(path)
	} else {
		return net.Dial("tcp", path)
	}
}

func main() {
	os.Exit(run()) // since defer func() needs to get called first
}

func run() int {

	rwc, err := openTPM(*tpmPath)
	if err != nil {
		fmt.Printf("can't open TPM %q: %v", *tpmPath, err)
		return 1
	}
	defer func() {
		if err := rwc.Close(); err != nil {
			fmt.Printf("can't close TPM %q: %v", *tpmPath, err)
		}
	}()

	rwr := transport.FromReadWriter(rwc)

	c, err := os.ReadFile(*pemFile)
	if err != nil {
		fmt.Printf("can't load keys %q: %v", *tpmPath, err)
		return 1
	}
	key, err := keyfile.Decode(c)
	if err != nil {
		log.Fatalf("can't decode keys %q: %v", *tpmPath, err)
		return 1
	}

	primaryKey, err := tpm2.CreatePrimary{
		PrimaryHandle: key.Parent,
		InPublic:      tpm2.New2B(keyfile.ECCSRK_H2_Template),
	}.Execute(rwr)
	if err != nil {
		fmt.Printf("can't create primary %q: %v", *tpmPath, err)
		return 1
	}

	defer func() {
		flushContextCmd := tpm2.FlushContext{
			FlushHandle: primaryKey.ObjectHandle,
		}
		_, _ = flushContextCmd.Execute(rwr)
	}()

	rsaKey, err := tpm2.Load{
		ParentHandle: tpm2.AuthHandle{
			Handle: primaryKey.ObjectHandle,
			Name:   tpm2.TPM2BName(primaryKey.Name),
			Auth:   tpm2.PasswordAuth([]byte("")),
		},
		InPublic:  key.Pubkey,
		InPrivate: key.Privkey,
	}.Execute(rwr)
	if err != nil {
		fmt.Printf("can't load  hmacKey : %v", err)
		return 1
	}

	defer func() {
		flushContextCmd := tpm2.FlushContext{
			FlushHandle: rsaKey.ObjectHandle,
		}
		_, _ = flushContextCmd.Execute(rwr)
	}()

	r, err := tpmsigner.NewTPMCrypto(&tpmsigner.TPM{
		TpmDevice: rwc,
		Handle:    rsaKey.ObjectHandle,
	})
	if err != nil {
		fmt.Printf("Error signing %v\n", err)
		return 1
	}

	rsaPubKey, ok := r.Public().(*rsa.PublicKey)
	if !ok {
		fmt.Printf("Error reading public %v\n", err)
		return 1
	}
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(rsaPubKey)
	if err != nil {
		fmt.Printf("error when marshalling publickey: %s \n", err)
		return 1
	}
	publicKeyBlock := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicKeyBytes,
	}

	err = pem.Encode(os.Stdout, publicKeyBlock)
	if err != nil {
		fmt.Printf("error when encoding public pem: %s \n", err)
		return 1
	}

	var payload []byte

	token := jwt.New()
	token.Set(`foo`, `bar`)
	payload, err = jwt.Sign(token, jwt.WithKey(jwa.RS256(), r))
	if err != nil {
		fmt.Printf("failed to generate signed payload: %s\n", err)
		return 1
	}

	fmt.Printf("JWT: %s\n", string(payload))

	ptoken, err := jwt.Parse(
		payload,
		jwt.WithValidate(true),
		jwt.WithKey(jwa.RS256(), rsaPubKey),
	)
	if err != nil {
		fmt.Printf("failed to parse JWT token: %s\n", err)
		return 1
	}
	buf, err := json.MarshalIndent(ptoken, "", "  ")
	if err != nil {
		fmt.Printf("failed to generate JSON: %s\n", err)
		return 1
	}
	fmt.Printf("%s\n", buf)

	return 0
}