verify: use context for resource fetching

Author: burgerdev

testing/mocks.go

@@ -15,6 +15,7 @@
 package testing
 
 import (
+	"context"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -246,6 +247,17 @@ func (g *Getter) Get(url string) ([]byte, error) {
 	return body, err
 }
 
+// GetContext checks whether the context expired, returns the context error if that's the case and
+// calls Get otherwise.
+func (g *Getter) GetContext(ctx context.Context, url string) ([]byte, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+		return g.Get(url)
+	}
+}
+
 // Done checks that all configured responses have been consumed, and errors
 // otherwise.
 func (g *Getter) Done(t testing.TB) {

verify/trust/trust.go

@@ -130,6 +130,20 @@ type HTTPSGetter interface {
 	Get(url string) ([]byte, error)
 }
 
+// ContextHTTPSGetter is an HTTPSGetter that accepts a context.Context.
+type ContextHTTPSGetter interface {
+	GetContext(ctx context.Context, url string) ([]byte, error)
+}
+
+// GetWith gets a resource from a URL using an HTTPSGetter.
+// If the HTTPSGetter implements ContextHTTPSGetter, the GetContext method will be used.
+func GetWith(ctx context.Context, getter HTTPSGetter, url string) ([]byte, error) {
+	if contextGetter, ok := getter.(ContextHTTPSGetter); ok {
+		return contextGetter.GetContext(ctx, url)
+	}
+	return getter.Get(url)
+}
+
 // AttestationRecreationErr represents a problem with fetching or interpreting associated
 // certificates for a given attestation report. This is typically due to network unreliability.
 type AttestationRecreationErr struct {
@@ -145,7 +159,16 @@ type SimpleHTTPSGetter struct{}
 
 // Get uses http.Get to return the HTTPS response body as a byte array.
 func (n *SimpleHTTPSGetter) Get(url string) ([]byte, error) {
-	resp, err := http.Get(url)
+	return n.GetContext(context.Background(), url)
+}
+
+// GetContext behaves like get, but forwards the context to the http package.
+func (n *SimpleHTTPSGetter) GetContext(ctx context.Context, url string) ([]byte, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return nil, err
 	} else if resp.StatusCode >= 300 {
@@ -160,9 +183,16 @@ func (n *SimpleHTTPSGetter) Get(url string) ([]byte, error) {
 	return body, nil
 }
 
+var (
+	_ = HTTPSGetter(&SimpleHTTPSGetter{})
+	_ = ContextHTTPSGetter(&SimpleHTTPSGetter{})
+)
+
 // RetryHTTPSGetter is a meta-HTTPS getter that will retry on failure a given number of times.
 type RetryHTTPSGetter struct {
 	// Timeout is how long to retry before failure.
+	// If Timeout is zero, the Get method will retry indefinitely and the GetContext method will
+	// retry until the input context expires.
 	Timeout time.Duration
 	// MaxRetryDelay is the maximum amount of time to wait between retries.
 	MaxRetryDelay time.Duration
@@ -172,11 +202,20 @@ type RetryHTTPSGetter struct {
 
 // Get fetches the body of the URL, retrying a given amount of times on failure.
 func (n *RetryHTTPSGetter) Get(url string) ([]byte, error) {
+	return n.GetContext(context.Background(), url)
+}
+
+// GetContext behaves like get, but forwards the context to the Getter and stops retrying when the
+// context expired.
+func (n *RetryHTTPSGetter) GetContext(ctx context.Context, url string) ([]byte, error) {
 	delay := initialDelay
-	ctx, cancel := context.WithTimeout(context.Background(), n.Timeout)
+	cancel := func() {}
+	if n.Timeout > 0 {
+		ctx, cancel = context.WithTimeout(ctx, n.Timeout)
+	}
 	var returnedError error
 	for {
-		body, err := n.Getter.Get(url)
+		body, err := GetWith(ctx, n.Getter, url)
 		if err == nil {
 			cancel()
 			return body, nil
@@ -189,12 +228,17 @@ func (n *RetryHTTPSGetter) Get(url string) ([]byte, error) {
 		select {
 		case <-ctx.Done():
 			cancel()
-			return nil, multierr.Append(returnedError, fmt.Errorf("timeout")) // context cancelled
+			return nil, multierr.Append(returnedError, ctx.Err())
 		case <-time.After(delay): // wait to retry
 		}
 	}
 }
 
+var (
+	_ = HTTPSGetter(&RetryHTTPSGetter{})
+	_ = ContextHTTPSGetter(&RetryHTTPSGetter{})
+)
+
 // DefaultHTTPSGetter returns the library's default getter implementation. It will
 // retry slowly due to the AMD KDS's rate limiting.
 func DefaultHTTPSGetter() HTTPSGetter {
@@ -311,14 +355,19 @@ func ClearProductCertCache() {
 // GetProductChain returns the ASK and ARK certificates of the given product line, either from getter
 // or from a cache of the results from the last successful call.
 func GetProductChain(productLine string, s abi.ReportSigner, getter HTTPSGetter) (*ProductCerts, error) {
+	return GetProductChainContext(context.Background(), productLine, s, getter)
+}
+
+// GetProductChainContext behaves like GetProductChain but forwards the context to the HTTPSGetter.
+func GetProductChainContext(ctx context.Context, productLine string, s abi.ReportSigner, getter HTTPSGetter) (*ProductCerts, error) {
 	if productLineCertCache == nil {
 		prodCacheMu.Lock()
 		productLineCertCache = make(map[string]*ProductCerts)
 		prodCacheMu.Unlock()
 	}
 	result, ok := productLineCertCache[productLine]
 	if !ok {
-		askark, err := getter.Get(kds.ProductCertChainURL(s, productLine))
+		askark, err := GetWith(ctx, getter, kds.ProductCertChainURL(s, productLine))
 		if err != nil {
 			return nil, &AttestationRecreationErr{
 				Msg: fmt.Sprintf("could not download ASK and ARK certificates: %v", err),

verify/trust/trust_test.go

@@ -16,6 +16,7 @@ package trust_test
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"testing"
 	"time"
@@ -127,10 +128,86 @@ func TestRetryHTTPSGetterAllFail(t *testing.T) {
 
 	body, err := r.Get("https://fetch.me")
 	if !bytes.Equal(body, []byte("")) {
-		t.Errorf("expected '%s' but got '%s'", "content", body)
+		t.Errorf("expected empty body but got %q", body)
 	}
 	if err == nil {
 		t.Errorf("expected error, but got none")
 	}
 	testGetter.Done(t)
 }
+
+func TestRetryHTTPSGetterContext(t *testing.T) {
+	testGetter := &test.Getter{
+		Responses: map[string][]test.GetResponse{
+			"https://fetch.me": {
+				{
+					Occurrences: 1,
+					Body:        []byte("content"),
+					Error:       nil,
+				},
+			},
+		},
+	}
+	r := &trust.RetryHTTPSGetter{
+		MaxRetryDelay: 1 * time.Millisecond,
+		Getter:        testGetter,
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	body, err := r.GetContext(ctx, "https://fetch.me")
+	if !bytes.Equal(body, []byte("")) {
+		t.Errorf("expected empty body but got %q", body)
+	}
+	if !errors.Is(err, context.Canceled) {
+		t.Errorf("expected error %q, but got %q", context.Canceled, err)
+	}
+}
+
+type recordingGetter struct {
+	getCalls int
+}
+
+func (r *recordingGetter) Get(url string) ([]byte, error) {
+	r.getCalls++
+	return []byte{}, nil
+}
+
+type recordingContextGetter struct {
+	recordingGetter
+	getContextCalls int
+}
+
+func (r *recordingContextGetter) GetContext(ctx context.Context, url string) ([]byte, error) {
+	r.getContextCalls++
+	return []byte{}, nil
+}
+
+func TestGetWith(t *testing.T) {
+	url := ""
+	t.Run("HTTPSGetter uses Get", func(t *testing.T) {
+		contextGetter := recordingContextGetter{}
+		if _, err := trust.GetWith(context.Background(), &contextGetter.recordingGetter, url); err != nil {
+			t.Fatalf("trust.GetWith returned an unexpected error: %v", err)
+		}
+		if contextGetter.getContextCalls != 0 {
+			t.Errorf("wrong number of calls to GetContext: got %d, want 0", contextGetter.getContextCalls)
+		}
+		if contextGetter.recordingGetter.getCalls != 1 {
+			t.Errorf("wrong number of calls to Get: got %d, want 1", contextGetter.getCalls)
+		}
+	})
+	t.Run("ContextHTTPSGetter uses GetContext", func(t *testing.T) {
+		contextGetter := recordingContextGetter{}
+		if _, err := trust.GetWith(context.Background(), &contextGetter, url); err != nil {
+			t.Fatalf("trust.GetWith returned an unexpected error: %v", err)
+		}
+		if contextGetter.getContextCalls != 1 {
+			t.Errorf("wrong number of calls to GetContext: got %d, want 1", contextGetter.getContextCalls)
+		}
+		if contextGetter.recordingGetter.getCalls != 0 {
+			t.Errorf("wrong number of calls to Get: got %d, want 0", contextGetter.getCalls)
+		}
+	})
+
+}

verify/verify.go

@@ -16,6 +16,7 @@
 package verify
 
 import (
+	"context"
 	"crypto/ecdsa"
 	"crypto/rsa"
 	"crypto/x509"
@@ -294,6 +295,12 @@ type CRLUnavailableErr struct {
 // GetCrlAndCheckRoot downloads the given cert's CRL from one of the distribution points and
 // verifies that the CRL is valid and doesn't revoke an intermediate key.
 func GetCrlAndCheckRoot(r *trust.AMDRootCerts, opts *Options) (*x509.RevocationList, error) {
+	return GetCrlAndCheckRootContext(context.Background(), r, opts)
+}
+
+// GetCrlAndCheckRootContext behaves like GetCrlAndCheckRoot but forwards the context to the
+// HTTPSGetter.
+func GetCrlAndCheckRootContext(ctx context.Context, r *trust.AMDRootCerts, opts *Options) (*x509.RevocationList, error) {
 	r.Mu.Lock()
 	defer r.Mu.Unlock()
 	getter := opts.Getter
@@ -308,7 +315,7 @@ func GetCrlAndCheckRoot(r *trust.AMDRootCerts, opts *Options) (*x509.RevocationL
 	}
 	var errs error
 	for _, url := range r.ProductCerts.Ask.CRLDistributionPoints {
-		bytes, err := getter.Get(url)
+		bytes, err := trust.GetWith(ctx, getter, url)
 		if err != nil {
 			errs = multierr.Append(errs, err)
 			continue
@@ -354,8 +361,13 @@ func verifyCRL(r *trust.AMDRootCerts) error {
 
 // VcekNotRevoked will consult the online CRL listed in the VCEK certificate for whether this cert
 // has been revoked. Returns nil if not revoked, error on any problem.
-func VcekNotRevoked(r *trust.AMDRootCerts, _ *x509.Certificate, options *Options) error {
-	_, err := GetCrlAndCheckRoot(r, options)
+func VcekNotRevoked(r *trust.AMDRootCerts, cert *x509.Certificate, options *Options) error {
+	return VcekNotRevokedContext(context.Background(), r, cert, options)
+}
+
+// VcekNotRevokedContext behaves like VcekNotRevoked but forwards the context to the HTTPSGetter.
+func VcekNotRevokedContext(ctx context.Context, r *trust.AMDRootCerts, _ *x509.Certificate, options *Options) error {
+	_, err := GetCrlAndCheckRootContext(ctx, r, options)
 	return err
 }
 
@@ -571,7 +583,7 @@ type Options struct {
 	// any missing certificates in an attestation's certificate chain. Uses Getter if false.
 	DisableCertFetching bool
 	// Getter takes a URL and returns the body of its contents. By default uses http.Get and returns
-	// the body.
+	// the body. If Getter implements trust.ContextHTTPSGetter, GetContext will be preferred over Get.
 	Getter trust.HTTPSGetter
 	// Now is the time at which to verify the validity of certificates. If unset, uses time.Now().
 	Now time.Time
@@ -662,6 +674,11 @@ func updateProductExpectation(product **spb.SevProduct, reportProduct *spb.SevPr
 // SnpAttestation verifies the protobuf representation of an attestation report's signature based
 // on the report's SignatureAlgo, provided the certificate chain is valid.
 func SnpAttestation(attestation *spb.Attestation, options *Options) error {
+	return SnpAttestationContext(context.Background(), attestation, options)
+}
+
+// SnpAttestationContext behaves like SnpAttestation but forwards the context to the HTTPSGetter.
+func SnpAttestationContext(ctx context.Context, attestation *spb.Attestation, options *Options) error {
 	if options == nil {
 		return fmt.Errorf("options cannot be nil")
 	}
@@ -670,7 +687,7 @@ func SnpAttestation(attestation *spb.Attestation, options *Options) error {
 	}
 	// Make sure we have the whole certificate chain, or at least the product
 	// info.
-	if err := fillInAttestation(attestation, options); err != nil {
+	if err := fillInAttestation(ctx, attestation, options); err != nil {
 		return err
 	}
 
@@ -786,7 +803,7 @@ func cpuidWorkaround(attestation *spb.Attestation, options *Options) (string, fu
 
 // fillInAttestation uses AMD's KDS to populate any empty certificate field in the attestation's
 // certificate chain.
-func fillInAttestation(attestation *spb.Attestation, options *Options) error {
+func fillInAttestation(ctx context.Context, attestation *spb.Attestation, options *Options) error {
 	if options.DisableCertFetching {
 		return nil
 	}
@@ -810,7 +827,7 @@ func fillInAttestation(attestation *spb.Attestation, options *Options) error {
 		attestation.CertificateChain = chain
 	}
 	if len(chain.GetAskCert()) == 0 || len(chain.GetArkCert()) == 0 {
-		askark, err := trust.GetProductChain(productLine, info.SigningKey, getter)
+		askark, err := trust.GetProductChainContext(ctx, productLine, info.SigningKey, getter)
 		if err != nil {
 			return err
 		}
@@ -826,7 +843,7 @@ func fillInAttestation(attestation *spb.Attestation, options *Options) error {
 	case abi.VcekReportSigner:
 		if len(chain.GetVcekCert()) == 0 {
 			vcekURL := kds.VCEKCertURL(productLine, report.GetChipId(), kds.TCBVersion(report.GetReportedTcb()))
-			vcek, err := getter.Get(vcekURL)
+			vcek, err := trust.GetWith(ctx, getter, vcekURL)
 			if err != nil {
 				return &trust.AttestationRecreationErr{
 					Msg: fmt.Sprintf("could not download VCEK certificate: %v", err),
@@ -854,11 +871,17 @@ func fillInAttestation(attestation *spb.Attestation, options *Options) error {
 // chain for the VCEK that supposedly signed the given report, and returns the Attestation
 // representation of their combination. If getter is nil, uses Golang's http.Get.
 func GetAttestationFromReport(report *spb.Report, options *Options) (*spb.Attestation, error) {
+	return GetAttestationFromReportContext(context.Background(), report, options)
+}
+
+// GetAttestationFromReportContext behaves like GetAttestationFromReport but forwards the context
+// to the HTTPSGetter.
+func GetAttestationFromReportContext(ctx context.Context, report *spb.Report, options *Options) (*spb.Attestation, error) {
 	result := &spb.Attestation{
 		Report:           report,
 		CertificateChain: &spb.CertificateChain{Extras: map[string][]byte{}},
 	}
-	if err := fillInAttestation(result, options); err != nil {
+	if err := fillInAttestation(ctx, result, options); err != nil {
 		return nil, err
 	}
 	// Attempt to fill in the product field of the attestation. Don't error at this
@@ -886,23 +909,33 @@ func GetAttestationFromReport(report *spb.Report, options *Options) (*spb.Attest
 // on the report's SignatureAlgo and uses the AMD Key Distribution Service to download the
 // report's corresponding VCEK certificate.
 func SnpReport(report *spb.Report, options *Options) error {
+	return SnpReportContext(context.Background(), report, options)
+}
+
+// SnpReportContext behaves like SnpReport but forwards the context to the HTTPSGetter.
+func SnpReportContext(ctx context.Context, report *spb.Report, options *Options) error {
 	if options.DisableCertFetching {
 		return errors.New("cannot verify attestation report without fetching certificates")
 	}
-	attestation, err := GetAttestationFromReport(report, options)
+	attestation, err := GetAttestationFromReportContext(ctx, report, options)
 	if err != nil {
 		return fmt.Errorf("could not recreate attestation from report: %w", err)
 	}
-	return SnpAttestation(attestation, options)
+	return SnpAttestationContext(ctx, attestation, options)
 }
 
 // RawSnpReport verifies the raw bytes representation of an attestation report's signature
 // based on the report's SignatureAlgo and uses the AMD Key Distribution Service to download
 // the report's corresponding VCEK certificate.
 func RawSnpReport(rawReport []byte, options *Options) error {
+	return RawSnpReportContext(context.Background(), rawReport, options)
+}
+
+// RawSnpReportContext behaves like RawSnpReport but forwards the context to the HTTPSGetter.
+func RawSnpReportContext(ctx context.Context, rawReport []byte, options *Options) error {
 	report, err := abi.ReportToProto(rawReport)
 	if err != nil {
 		return fmt.Errorf("could not interpret report bytes: %v", err)
 	}
-	return SnpReport(report, options)
+	return SnpReportContext(ctx, report, options)
 }