Access resources not managed by Google Cloud IAM

package main

import (
  "context"
  "crypto/rsa"
  "encoding/base64"
  "encoding/json"
  "errors"
  "fmt"
  "io"
  "math/big"
  "net"
  "net/http"
  "strings"

  "github.com/golang-jwt/jwt/v4"
)

const (
  socketPath     = "/run/container_launcher/teeserver.sock"
  expectedIssuer = "https://confidentialcomputing.googleapis.com"
  wellKnownPath  = "/.well-known/openid-configuration"
)

type jwksFile struct {
  Keys []jwk `json:"keys"`
}

type jwk struct {
  N   string `json:"n"`   // "nMMTBwJ7H6Id8zUCZd-L7uoNyz9b7lvoyse9izD9l2rtOhWLWbiG-7pKeYJyHeEpilHP4KdQMfUo8JCwhd-OMW0be_XtEu3jXEFjuq2YnPSPFk326eTfENtUc6qJohyMnfKkcOcY_kTE11jM81-fsqtBKjO_KiSkcmAO4wJJb8pHOjue3JCP09ZANL1uN4TuxbM2ibcyf25ODt3WQn54SRQTV0wn098Y5VDU-dzyeKYBNfL14iP0LiXBRfHd4YtEaGV9SBUuVhXdhx1eF0efztCNNz0GSLS2AEPLQduVuFoUImP4s51YdO9TPeeQ3hI8aGpOdC0syxmZ7LsL0rHE1Q",
  E   string `json:"e"`   // "AQAB" or 65537 as an int
  Kid string `json:"kid"` // "1f12fa916c3a0ef585894b4b420ad17dc9d6cdf5",

  // Unused fields:
  // Alg string `json:"alg"` // "RS256",
  // Kty string `json:"kty"` // "RSA",
  // Use string `json:"use"` // "sig",
}

type wellKnown struct {
  JwksURI string `json:"jwks_uri"` // "https://www.googleapis.com/service_accounts/v1/metadata/jwk/signer@confidentialspace-sign.iam.gserviceaccount.com"

  // Unused fields:
  // Iss                                   string `json:"issuer"`                                // "https://confidentialcomputing.googleapis.com"
  // Subject_types_supported               string `json:"subject_types_supported"`               // [ "public" ]
  // Response_types_supported              string `json:"response_types_supported"`              // [ "id_token" ]
  // Claims_supported                      string `json:"claims_supported"`                      // [ "sub", "aud", "exp", "iat", "iss", "jti", "nbf", "dbgstat", "eat_nonce", "google_service_accounts", "hwmodel", "oemid", "secboot", "submods", "swname", "swversion" ]
  // Id_token_signing_alg_values_supported string `json:"id_token_signing_alg_values_supported"` // [ "RS256" ]
  // Scopes_supported                      string `json:"scopes_supported"`                      // [ "openid" ]
}

func getWellKnownFile() (wellKnown, error) {
  httpClient := http.Client{}
  resp, err := httpClient.Get(expectedIssuer + wellKnownPath)
  if err != nil {
    return wellKnown{}, fmt.Errorf("failed to get raw .well-known response: %w", err)
  }

  wellKnownJSON, err := io.ReadAll(resp.Body)
  if err != nil {
    return wellKnown{}, fmt.Errorf("failed to read .well-known response: %w", err)
  }

  wk := wellKnown{}
  json.Unmarshal(wellKnownJSON, &wk)
  return wk, nil
}

func getJWKFile() (jwksFile, error) {
  wk, err := getWellKnownFile()
  if err != nil {
    return jwksFile{}, fmt.Errorf("failed to get .well-known json: %w", err)
  }

  // Get JWK URI from .wellknown
  uri := wk.JwksURI
  fmt.Printf("jwks URI: %v\n", uri)

  httpClient := http.Client{}
  resp, err := httpClient.Get(uri)
  if err != nil {
    return jwksFile{}, fmt.Errorf("failed to get raw JWK response: %w", err)
  }

  jwkbytes, err := io.ReadAll(resp.Body)
  if err != nil {
    return jwksFile{}, fmt.Errorf("failed to read JWK body: %w", err)
  }

  file := jwksFile{}
  err = json.Unmarshal(jwkbytes, &file)
  if err != nil {
    return jwksFile{}, fmt.Errorf("failed to unmarshall JWK content: %w", err)
  }

  return file, nil
}

// N and E are 'base64urlUInt' encoded: https://www.rfc-editor.org/rfc/rfc7518#section-6.3
func base64urlUIntDecode(s string) (*big.Int, error) {
  b, err := base64.RawURLEncoding.DecodeString(s)
  if err != nil {
    return nil, err
  }
  z := new(big.Int)
  z.SetBytes(b)
  return z, nil
}

func getRSAPublicKeyFromJWKsFile(t *jwt.Token) (any, error) {
  keysfile, err := getJWKFile()
  if err != nil {
    return nil, fmt.Errorf("failed to fetch the JWK file: %w", err)
  }

  // Multiple keys are present in this endpoint to allow for key rotation.
  // This method finds the key that was used for signing to pass to the validator.
  kid := t.Header["kid"]
  for _, key := range keysfile.Keys {
    if key.Kid != kid {
      continue // Select the key used for signing
    }

    n, err := base64urlUIntDecode(key.N)
    if err != nil {
      return nil, fmt.Errorf("failed to decode key.N %w", err)
    }
    e, err := base64urlUIntDecode(key.E)
    if err != nil {
      return nil, fmt.Errorf("failed to decode key.E %w", err)
    }

    // The parser expects an rsa.PublicKey: https://github.com/golang-jwt/jwt/blob/main/rsa.go#L53
    // or an array of keys. We chose to show passing a single key in this example as its possible
    // not all validators accept multiple keys for validation.
    return &rsa.PublicKey{
      N: n,
      E: int(e.Int64()),
    }, nil
  }

  return nil, fmt.Errorf("failed to find key with kid '%v' from well-known endpoint", kid)
}

func decodeAndValidateToken(tokenBytes []byte, keyFunc func(t *jwt.Token) (any, error)) (*jwt.Token, error) {
  var err error
  fmt.Println("Unmarshalling token and checking its validity...")
  token, err := jwt.NewParser().Parse(string(tokenBytes), keyFunc)

  fmt.Printf("Token valid: %v", token.Valid)
  if token.Valid {
    return token, nil
  }
  if ve, ok := err.(*jwt.ValidationError); ok {
    if ve.Errors&jwt.ValidationErrorMalformed != 0 {
      return nil, fmt.Errorf("token format invalid. Please contact the Confidential Space team for assistance")
    }
    if ve.Errors&(jwt.ValidationErrorNotValidYet) != 0 {
      // If device time is not synchronized with the Attestation Service you may need to account for that here.
      return nil, errors.New("token is not active yet")
    }
    if ve.Errors&(jwt.ValidationErrorExpired) != 0 {
      return nil, fmt.Errorf("token is expired")
    }
    return nil, fmt.Errorf("unknown validation error: %v", err)
  }

  return nil, fmt.Errorf("couldn't handle this token or couldn't read a validation error: %v", err)
}

func main() {
  // Get a token from a workload running in Confidential Space
  tokenbytes, err := getTokenBytesFromWorkload()

  // Write a method to return a public key from the well-known endpoint
  keyFunc := getRSAPublicKeyFromJWKsFile

  // Verify properties of the original Confidential Space workload that generated the attestation
  // using the token claims.
  token, err := decodeAndValidateToken(tokenbytes, keyFunc)
  if err != nil {
    panic(err)
  }

  claimsString, err := json.MarshalIndent(token.Claims, "", "  ")
  if err != nil {
    panic(err)
  }
  fmt.Println(string(claimsString))
}

// This code is an example of how to validate a PKI token. This library is not an official library,
// nor is it endorsed by Google.

// ValidatePKIToken validates the PKI token returned from the attestation service is valid.
// Returns a valid jwt.Token or returns an error if invalid.
func ValidatePKIToken(storedRootCertificate x509.Certificate, attestationToken string) (jwt.Token, error) {
  // IMPORTANT: The attestation token should be considered untrusted until the certificate chain and
  // the signature is verified.

  jwtHeaders, err := ExtractJWTHeaders(attestationToken)
  if err != nil {
    return jwt.Token{}, fmt.Errorf("ExtractJWTHeaders(token) returned error: %v", err)
  }

  if jwtHeaders["alg"] != "RS256" {
    return jwt.Token{}, fmt.Errorf("ValidatePKIToken(string, *attestpb.Attestation, *v1mainpb.VerifyAttestationRequest) - got Alg: %v, want: %v", jwtHeaders["alg"], "RS256")
  }

  // Additional Check: Validate the ALG in the header matches the certificate SPKI.
  // https://datatracker.ietf.org/doc/html/rfc5280#section-4.1.2.7
  // This is included in golangs jwt.Parse function

  x5cHeaders := jwtHeaders["x5c"].([]any)
  certificates, err := ExtractCertificatesFromX5CHeader(x5cHeaders)
  if err != nil {
    return jwt.Token{}, fmt.Errorf("ExtractCertificatesFromX5CHeader(x5cHeaders) returned error: %v", err)
  }

  // Verify the leaf certificate signature algorithm is an RSA key
  if certificates.LeafCert.SignatureAlgorithm != x509.SHA256WithRSA {
    return jwt.Token{}, fmt.Errorf("leaf certificate signature algorithm is not SHA256WithRSA")
  }

  // Verify the leaf certificate public key algorithm is RSA
  if certificates.LeafCert.PublicKeyAlgorithm != x509.RSA {
    return jwt.Token{}, fmt.Errorf("leaf certificate public key algorithm is not RSA")
  }

  // Verify the storedRootCertificate is the same as the root certificate returned in the token.
  // storedRootCertificate is downloaded from the confidential computing well known endpoint
  // https://confidentialcomputing.googleapis.com/.well-known/attestation-pki-root
  err = CompareCertificates(storedRootCertificate, *certificates.RootCert)
  if err != nil {
    return jwt.Token{}, fmt.Errorf("failed to verify certificate chain: %v", err)
  }

  err = VerifyCertificateChain(certificates)
  if err != nil {
    return jwt.Token{}, fmt.Errorf("VerifyCertificateChain(string, *attestpb.Attestation, *v1mainpb.VerifyAttestationRequest) - error verifying x5c chain: %v", err)
  }

  keyFunc := func(token *jwt.Token) (any, error) {
    return certificates.LeafCert.PublicKey, nil
  }

  verifiedJWT, err := jwt.Parse(attestationToken, keyFunc)
  return *verifiedJWT, err
}

// ExtractJWTHeaders parses the JWT and returns the headers.
func ExtractJWTHeaders(token string) (map[string]any, error) {
  parser := &jwt.Parser{}

  // The claims returned from the token are unverified at this point
  // Do not use the claims until the algorithm, certificate chain verification and root certificate
  // comparison is successful
  unverifiedClaims := &jwt.MapClaims{}
  parsedToken, _, err := parser.ParseUnverified(token, unverifiedClaims)
  if err != nil {
    return nil, fmt.Errorf("Failed to parse claims token: %v", err)
  }

  return parsedToken.Header, nil
}

// PKICertificates contains the certificates extracted from the x5c header.
type PKICertificates struct {
  LeafCert         *x509.Certificate
  IntermediateCert *x509.Certificate
  RootCert         *x509.Certificate
}

// ExtractCertificatesFromX5CHeader extracts the certificates from the given x5c header.
func ExtractCertificatesFromX5CHeader(x5cHeaders []any) (PKICertificates, error) {
  if x5cHeaders == nil {
    return PKICertificates{}, fmt.Errorf("VerifyAttestation(string, *attestpb.Attestation, *v1mainpb.VerifyAttestationRequest) - x5c header not set")
  }

  x5c := []string{}
  for _, header := range x5cHeaders {
    x5c = append(x5c, header.(string))
  }

  // The PKI token x5c header should have 3 certificates - leaf, intermediate and root
  if len(x5c) != 3 {
    return PKICertificates{}, fmt.Errorf("incorrect number of certificates in x5c header, expected 3 certificates, but got %v", len(x5c))
  }

  leafCert, err := DecodeAndParseCertificate(x5c[0])
  if err != nil {
    return PKICertificates{}, fmt.Errorf("cannot parse leaf certificate: %v", err)
  }

  intermediateCert, err := DecodeAndParseCertificate(x5c[1])
  if err != nil {
    return PKICertificates{}, fmt.Errorf("cannot parse intermediate certificate: %v", err)
  }

  rootCert, err := DecodeAndParseCertificate(x5c[2])
  if err != nil {
    return PKICertificates{}, fmt.Errorf("cannot parse root certificate: %v", err)
  }

  certificates := PKICertificates{
    LeafCert:         leafCert,
    IntermediateCert: intermediateCert,
    RootCert:         rootCert,
  }
  return certificates, nil
}

// DecodeAndParseCertificate decodes the given PEM certificate string and parses it into an x509 certificate.
func DecodeAndParseCertificate(certificate string) (*x509.Certificate, error) {
  block, _ := pem.Decode([]byte(certificate))
  if block == nil {
    return nil, fmt.Errorf("cannot decode certificate")
  }

  cert, err := x509.ParseCertificate(block.Bytes)
  if err != nil {
    return nil, fmt.Errorf("cannot parse certificate: %v", err)
  }

  return cert, nil
}

// VerifyCertificateChain verifies the certificate chain from leaf to root.
// It also checks that all certificate lifetimes are valid.
func VerifyCertificateChain(certificates PKICertificates) error {
  if isCertificateLifetimeValid(certificates.LeafCert) {
    return fmt.Errorf("leaf certificate is not valid")
  }

  if isCertificateLifetimeValid(certificates.IntermediateCert) {
    return fmt.Errorf("intermediate certificate is not valid")
  }
  interPool := x509.NewCertPool()
  interPool.AddCert(certificates.IntermediateCert)

  if isCertificateLifetimeValid(certificates.RootCert) {
    return fmt.Errorf("root certificate is not valid")
  }
  rootPool := x509.NewCertPool()
  rootPool.AddCert(certificates.RootCert)

  _, err := certificates.LeafCert.Verify(x509.VerifyOptions{
    Intermediates: interPool,
    Roots:         rootPool,
    KeyUsages:     []x509.ExtKeyUsage{x509.ExtKeyUsageAny},
  })

  if err != nil {
    return fmt.Errorf("failed to verify certificate chain: %v", err)
  }

  return nil
}

func isCertificateLifetimeValid(certificate *x509.Certificate) bool {
  currentTime := time.Now()
  // check the current time is after the certificate NotBefore time
  if !currentTime.After(certificate.NotBefore) {
    return false
  }

  // check the current time is before the certificate NotAfter time
  if currentTime.Before(certificate.NotAfter) {
    return false
  }

  return true
}

// CompareCertificates compares two certificate fingerprints.
func CompareCertificates(cert1 x509.Certificate, cert2 x509.Certificate) error {
  fingerprint1 := sha256.Sum256(cert1.Raw)
  fingerprint2 := sha256.Sum256(cert2.Raw)
  if fingerprint1 != fingerprint2 {
    return fmt.Errorf("certificate fingerprint mismatch")
  }
  return nil
}