NOISSUE - Refactor attestation handling: rename AttestationResult to AzureAttestationToken (#504)

* Refactor attestation handling: rename AttestationResult to AzureAttestationToken

- Updated the protobuf definition to change azureAttestationResponse to azureAttestationToken.
- Refactored the Service interface and its implementation to replace AttestationResult with AzureAttestationToken.
- Modified mock functions and tests to reflect the new naming and functionality.
- Adjusted CLI commands to use the new AzureAttestationToken method.
- Removed the AzureToken constant from the attestation package as it is no longer needed.

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Remove redundant data checks and logging in SendData and sendData methods

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Update agent/api/grpc/server_test.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update agent/api/grpc/endpoint_test.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Refactor attestation handling: rename AttestationToken to AzureAttestationToken in server and test files

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Refactor attestation command output messages for clarity and consistency

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Rename AttestationToken to AzureAttestationToken in TestAttestationToken for consistency

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Refactor TestChangeAttestationConfiguration to use vtpm.ConvertPolicyToJSON for JSON conversion

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Fix: reset temporary file pointer after zipping directory

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Sammy Kerata Oina
2025-08-26 15:42:33 +03:00
committed by GitHub
parent 654e22bba5
commit 4b27b98edb
26 changed files with 398 additions and 411 deletions
+10 -11
View File
@@ -148,25 +148,24 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
attType = attestation.SNPvTPM
case AzureToken:
cmd.Println("Fetching Azure token")
attType = attestation.AzureToken
case TDX:
cmd.Println("Fetching TDX attestation report")
attType = attestation.TDX
}
if (attType == attestation.VTPM || attType == attestation.SNPvTPM) && len(nonce) == 0 {
if (attestationType == VTPM || attestationType == SNPvTPM) && len(nonce) == 0 {
msg := color.New(color.FgRed).Sprint("vTPM nonce must be defined for vTPM attestation ❌ ")
cmd.Println(msg)
return
}
if (attType == attestation.SNP || attType == attestation.SNPvTPM) && len(teeNonce) == 0 {
if (attestationType == SNP || attestationType == SNPvTPM) && len(teeNonce) == 0 {
msg := color.New(color.FgRed).Sprint("TEE nonce must be defined for SEV-SNP attestation ❌ ")
cmd.Println(msg)
return
}
if (attType == attestation.AzureToken) && len(tokenNonce) == 0 {
if (attestationType == AzureToken) && len(tokenNonce) == 0 {
msg := color.New(color.FgRed).Sprint("Token nonce must be defined for Azure attestation ❌ ")
cmd.Println(msg)
return
@@ -184,13 +183,13 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
}
var fixedVtpmNonceByte [vtpm.Nonce]byte
if attType != attestation.SNP {
if attType != attestation.SNP || attestationType == AzureToken {
if (len(nonce) > vtpm.Nonce) || (len(tokenNonce) > vtpm.Nonce) {
msg := color.New(color.FgRed).Sprintf("vTPM nonce must be a hex encoded string of length lesser or equal %d bytes ❌ ", vtpm.Nonce)
cmd.Println(msg)
return
}
if attType == attestation.AzureToken {
if attestationType == AzureToken {
copy(fixedVtpmNonceByte[:], tokenNonce)
} else {
copy(fixedVtpmNonceByte[:], nonce)
@@ -199,7 +198,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
filename := attestationFilePath
if attType == attestation.AzureToken {
if attestationType == AzureToken {
filename = azureAttestResultFilePath
}
@@ -217,10 +216,10 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
var returnJsonAzureToken bool
if attType == attestation.AzureToken {
err := cli.agentSDK.AttestationResult(cmd.Context(), fixedVtpmNonceByte, int(attType), attestationFile)
if attestationType == AzureToken {
err := cli.agentSDK.AttestationToken(cmd.Context(), fixedVtpmNonceByte, int(attType), attestationFile)
if err != nil {
printError(cmd, "Failed to get attestation result due to error: %v ❌", err)
printError(cmd, "Failed to get attestation token due to error: %v ❌", err)
return
}
returnJsonAzureToken = !getAzureTokenJWT
@@ -279,7 +278,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
}
}
cmd.Println("Attestation result retrieved and saved successfully!")
cmd.Println("Attestation retrieved and saved successfully!")
},
}
+1 -2
View File
@@ -5,7 +5,6 @@ package cli
import (
"bytes"
"encoding/base64"
"encoding/json"
"os"
"testing"
@@ -23,7 +22,7 @@ func TestChangeAttestationConfiguration(t *testing.T) {
initialConfig := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
initialJSON, err := json.Marshal(initialConfig)
initialJSON, err := vtpm.ConvertPolicyToJSON(&initialConfig)
require.NoError(t, err)
err = os.WriteFile(tmpfile.Name(), initialJSON, 0o644)
require.NoError(t, err)
+7 -7
View File
@@ -67,21 +67,21 @@ func TestNewGetAttestationCmd(t *testing.T) {
args: []string{"snp", "--tee", teeNonce},
mockResponse: []byte("mock attestation"),
mockError: nil,
expectedOut: "Attestation result retrieved and saved successfully!",
expectedOut: "Attestation retrieved and saved successfully!",
},
{
name: "successful vTPM attestation retrieval",
args: []string{"vtpm", "--vtpm", vtpmNonce},
mockResponse: []byte("mock attestation"),
mockError: nil,
expectedOut: "Attestation result retrieved and saved successfully!",
expectedOut: "Attestation retrieved and saved successfully!",
},
{
name: "successful SNP-vTPM attestation retrieval",
args: []string{"snp-vtpm", "--tee", teeNonce, "--vtpm", vtpmNonce},
mockResponse: []byte("mock attestation"),
mockError: nil,
expectedOut: "Attestation result retrieved and saved successfully!",
expectedOut: "Attestation retrieved and saved successfully!",
},
{
name: "missing vTPM nonce",
@@ -137,7 +137,7 @@ func TestNewGetAttestationCmd(t *testing.T) {
args: []string{"snp", "--tee", teeNonce, "--reporttextproto"},
mockResponse: validattestation,
mockError: nil,
expectedOut: "Attestation result retrieved and saved successfully!",
expectedOut: "Attestation retrieved and saved successfully!",
},
{
name: "connection error",
@@ -151,14 +151,14 @@ func TestNewGetAttestationCmd(t *testing.T) {
args: []string{"azure-token", "--token", tokenNonce},
mockResponse: []byte("eyJhbGciOiAiUlMyNTYifQ.eyJzdWIiOiAidGVzdC11c2VyIn0.signature"),
mockError: nil,
expectedOut: "Fetching Azure token\nAttestation result retrieved and saved successfully!\n",
expectedOut: "Fetching Azure token\nAttestation retrieved and saved successfully!\n",
},
{
name: "failed to retrieve Azure token",
args: []string{"azure-token", "--token", tokenNonce},
mockResponse: nil,
mockError: errors.New("error"),
expectedErr: "Fetching Azure token\nFailed to get attestation result due to error: error ❌\n",
expectedErr: "Fetching Azure token\nFailed to get attestation token due to error: error ❌\n",
},
{
name: "invalid token nonce size",
@@ -189,7 +189,7 @@ func TestNewGetAttestationCmd(t *testing.T) {
require.NoError(t, err)
})
mockSDK.On("AttestationResult", mock.Anything, [vtpm.Nonce]byte(bytes.Repeat([]byte{0x00}, vtpm.Nonce)), mock.Anything, mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
mockSDK.On("AttestationToken", mock.Anything, [vtpm.Nonce]byte(bytes.Repeat([]byte{0x00}, vtpm.Nonce)), mock.Anything, mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
_, err := args.Get(3).(*os.File).Write(tc.mockResponse)
require.NoError(t, err)
})
+1
View File
@@ -44,6 +44,7 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
var dataset *os.File
if f.IsDir() {
cmd.Println("Detected directory, zipping dataset...")
dataset, err = internal.ZipDirectoryToTempFile(datasetPath)
if err != nil {
printError(cmd, "Error zipping dataset directory: %v ❌ ", err)