From 8b58e8ed71b5796931a1124da4985152d269c6e4 Mon Sep 17 00:00:00 2001
From: Paul Meyer <49727155+katexochen@users.noreply.github.com>
Date: Fri, 6 Dec 2024 14:34:41 +0100
Subject: [PATCH] attestation/snp: use context with timeout on THIM request

Signed-off-by: Paul Meyer <49727155+katexochen@users.noreply.github.com>
---
 internal/attestation/snp/issuer/issuer.go | 4 ++--
 internal/attestation/snp/issuer/thim.go   | 7 +++++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/internal/attestation/snp/issuer/issuer.go b/internal/attestation/snp/issuer/issuer.go
index e4f4b4521..532fc445e 100644
--- a/internal/attestation/snp/issuer/issuer.go
+++ b/internal/attestation/snp/issuer/issuer.go
@@ -44,7 +44,7 @@ func (i *Issuer) OID() asn1.ObjectIdentifier {
 }
 
 // Issue the attestation document.
-func (i *Issuer) Issue(_ context.Context, ownPublicKey []byte, nonce []byte) (res []byte, err error) {
+func (i *Issuer) Issue(ctx context.Context, ownPublicKey []byte, nonce []byte) (res []byte, err error) {
 	i.logger.Info("Issue called")
 	defer func() {
 		if err != nil {
@@ -71,7 +71,7 @@ func (i *Issuer) Issue(_ context.Context, ownPublicKey []byte, nonce []byte) (re
 
 	// Get cert chain from THIM
 	var certChain *spb.CertificateChain
-	thimRaw, err := i.thimGetter.GetCertification()
+	thimRaw, err := i.thimGetter.GetCertification(ctx)
 	if err != nil {
 		i.logger.Info("Could not retrieve THIM certification", "error", err)
 	} else {
diff --git a/internal/attestation/snp/issuer/thim.go b/internal/attestation/snp/issuer/thim.go
index 94b4ca8b9..6564f24c9 100644
--- a/internal/attestation/snp/issuer/thim.go
+++ b/internal/attestation/snp/issuer/thim.go
@@ -4,6 +4,7 @@
 package issuer
 
 import (
+	"context"
 	"encoding/json"
 	"encoding/pem"
 	"fmt"
@@ -78,7 +79,7 @@ func NewTHIMGetter(httpClient httpClient) *THIMGetter {
 }
 
 // GetCertification returns the THIM certification.
-func (t *THIMGetter) GetCertification() (THIMSNPCertification, error) {
+func (t *THIMGetter) GetCertification(ctx context.Context) (THIMSNPCertification, error) {
 	// Return cached response if it is still valid.
 	if cached := t.getCached(); cached != nil {
 		var certification THIMSNPCertification
@@ -102,7 +103,9 @@ func (t *THIMGetter) GetCertification() (THIMSNPCertification, error) {
 			"Metadata": {"true"},
 		},
 	}
-	resp, err := t.httpClient.Do(req)
+	reqCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
+	defer cancel()
+	resp, err := t.httpClient.Do(req.WithContext(reqCtx))
 	if err != nil {
 		return THIMSNPCertification{}, fmt.Errorf("getting THIM certification: %w", err)
 	}