Fix panic catching in launcher and AKCert

Catch panic using defer recovery in launcher.
Return an empty MachineState pointer instead of nil in validateAKCert.

Signed-off-by: Jiankun Lu <jiankun@google.com>

go.work.sum

@@ -1,4 +1,3 @@
-github.com/cenkalti/backoff/v4 v4.1.3 h1:cFAlzYUlVYDysBEH2T5hyJZMh3+5+WCBvSnK6Q8UtC4=
-github.com/cenkalti/backoff/v4 v4.1.3/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
+github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
 github.com/googleapis/gax-go v2.0.2+incompatible h1:silFMLAnr330+NRuag/VjIGF7TLp/LBrV2CJKFLWEww=
 golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=

launcher/go.mod

@@ -15,16 +15,12 @@ require (
 	github.com/opencontainers/runtime-spec v1.0.3-0.20210326190908-1c3f411f0417
 	golang.org/x/oauth2 v0.0.0-20220622183110-fd043fe589d2
 	google.golang.org/api v0.86.0
-	google.golang.org/genproto v0.0.0-20220624142145-8cd45d7dbd1f
-	google.golang.org/grpc v1.47.0
-	google.golang.org/protobuf v1.28.0
 )
 
 require (
 	cloud.google.com/go v0.102.0 // indirect
 	github.com/Microsoft/go-winio v0.5.2 // indirect
 	github.com/Microsoft/hcsshim v0.9.3 // indirect
-	github.com/cenkalti/backoff v2.2.1+incompatible // indirect
 	github.com/containerd/cgroups v1.0.3 // indirect
 	github.com/containerd/continuity v0.3.0 // indirect
 	github.com/containerd/fifo v1.0.0 // indirect
@@ -57,6 +53,9 @@ require (
 	golang.org/x/sys v0.0.0-20220624220833-87e55d714810 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/genproto v0.0.0-20220624142145-8cd45d7dbd1f // indirect
+	google.golang.org/grpc v1.47.0 // indirect
+	google.golang.org/protobuf v1.28.0 // indirect
 	gopkg.in/yaml.v3 v3.0.0 // indirect
 )
 

launcher/go.sum

@@ -209,7 +209,6 @@ github.com/caarlos0/ctrlc v1.0.0/go.mod h1:CdXpj4rmq0q/1Eb44M9zi2nKB0QraNKuRGYGr
 github.com/campoy/unique v0.0.0-20180121183637-88950e537e7e/go.mod h1:9IOqJGCPMSc6E5ydlp5NIonxObaeu/Iub/X03EKPVYo=
 github.com/casbin/casbin/v2 v2.1.2/go.mod h1:YcPU1XXisHhLzuxH9coDNf2FbKpjGlbCg3n9yuLkIJQ=
 github.com/cavaliercoder/go-cpio v0.0.0-20180626203310-925f9528c45e/go.mod h1:oDpT4efm8tSYHXV5tHSdRvBet/b/QzxZ+XyyPehvm3A=
-github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
 github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/cenkalti/backoff/v4 v4.1.1/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
 github.com/cenkalti/backoff/v4 v4.1.3 h1:cFAlzYUlVYDysBEH2T5hyJZMh3+5+WCBvSnK6Q8UtC4=

launcher/launcher/main.go

@@ -35,13 +35,23 @@ var launchSpec spec.LaunchSpec
 
 func main() {
 	var exitCode int
-	defer func() {
-		os.Exit(exitCode)
-	}()
 
 	logger = log.Default()
+	// log.Default() outputs to stderr; change to stdout.
+	log.SetOutput(os.Stdout)
 	logger.Println("TEE container launcher initiating")
 
+	defer func() {
+		// catch panic, will only output to stdout, because cloud logging closed
+		// This should rarely happen (almost impossible), the only place can panic
+		// recover here is in the deferred logClient.Close().
+		if r := recover(); r != nil {
+			logger.Println("Panic:", r)
+			exitCode = 2
+		}
+		os.Exit(exitCode)
+	}()
+
 	mdsClient = metadata.NewClient(nil)
 	projectID, err := mdsClient.ProjectID()
 	if err != nil {
@@ -71,6 +81,14 @@ func main() {
 		return
 	}
 
+	defer func() {
+		// catch panic, will also output to cloud logging if possible
+		if r := recover(); r != nil {
+			logger.Println("Panic:", r)
+			exitCode = 2
+		}
+		logger.Println("TEE container launcher exiting with exit code:", exitCode)
+	}()
 	if err = startLauncher(); err != nil {
 		logger.Println(err)
 	}

server/verify.go

@@ -100,7 +100,7 @@ func VerifyAttestation(attestation *pb.Attestation, opts VerifyOpts) (*pb.Machin
 	}
 
 	var akPubKey crypto.PublicKey
-	machineState := &pb.MachineState{}
+	var machineState *pb.MachineState
 	if len(attestation.GetAkCert()) == 0 {
 		// If the AK Cert is not in the attestation, use the AK Public Area.
 		akPubArea, err := tpm2.DecodePublic(attestation.GetAkPub())
@@ -111,7 +111,8 @@ func VerifyAttestation(attestation *pb.Attestation, opts VerifyOpts) (*pb.Machin
 		if err != nil {
 			return nil, fmt.Errorf("failed to get AK public key: %w", err)
 		}
-		if err := validateAKPub(akPubKey, opts); err != nil {
+		machineState, err = validateAKPub(akPubKey, opts)
+		if err != nil {
 			return nil, fmt.Errorf("failed to validate AK public key: %w", err)
 		}
 	} else {
@@ -234,18 +235,18 @@ func validateOpts(opts VerifyOpts) error {
 	return nil
 }
 
-func validateAKPub(ak crypto.PublicKey, opts VerifyOpts) error {
+func validateAKPub(ak crypto.PublicKey, opts VerifyOpts) (*pb.MachineState, error) {
 	for _, trusted := range opts.TrustedAKs {
 		if internal.PubKeysEqual(ak, trusted) {
-			return nil
+			return &pb.MachineState{}, nil
 		}
 	}
-	return fmt.Errorf("key not trusted")
+	return nil, fmt.Errorf("key not trusted")
 }
 
 func validateAKCert(akCert *x509.Certificate, opts VerifyOpts) (*pb.MachineState, error) {
 	if len(opts.TrustedRootCerts) == 0 {
-		return nil, validateAKPub(akCert.PublicKey.(crypto.PublicKey), opts)
+		return validateAKPub(akCert.PublicKey.(crypto.PublicKey), opts)
 	}
 
 	// We manually handle the SAN extension because x509 marks it unhandled if

server/verify_test.go

@@ -244,6 +244,32 @@ func TestVerifyBasicAttestation(t *testing.T) {
 	}
 }
 
+func TestVerifyWithTrustedAK(t *testing.T) {
+	rwc := test.GetTPM(t)
+	defer client.CheckedClose(t, rwc)
+
+	ak, err := client.AttestationKeyRSA(rwc)
+	if err != nil {
+		t.Fatalf("failed to generate AK: %v", err)
+	}
+	defer ak.Close()
+
+	nonce := []byte("super secret nonce")
+	attestation, err := ak.Attest(client.AttestOpts{Nonce: nonce})
+	if err != nil {
+		t.Fatalf("failed to attest: %v", err)
+	}
+
+	opts := VerifyOpts{
+		Nonce:      nonce,
+		TrustedAKs: []crypto.PublicKey{ak.PublicKey()},
+	}
+	_, err = VerifyAttestation(attestation, opts)
+	if err != nil {
+		t.Errorf("failed to verify: %v", err)
+	}
+}
+
 func TestVerifySHA1Attestation(t *testing.T) {
 	rwc := test.GetTPM(t)
 	defer client.CheckedClose(t, rwc)