Files
cocos/pkg/sdk/agent_test.go
T
Sammy Kerata Oina 207bfd99af COCOS-525-487 - Refactor attestation and atls (#562)
* Refactor attestation handling to remove quoteprovider dependency

- Removed references to quoteprovider in various files, replacing them with vtpm where necessary.
- Updated function signatures and implementations to use SEVNonce instead of quoteprovider.Nonce.
- Introduced new vtpm package to handle SEV-related attestation logic, including fetching and verifying attestation reports.
- Adjusted tests to reflect changes in the attestation logic and ensure compatibility with the new structure.
- Deleted the now redundant quoteprovider/sev_test.go file.

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

* fix: Add veraison/go-cose dependency to go.mod

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

* feat: Introduce TLS package for enhanced security configuration and refactor client code to utilize new TLS utilities

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

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.com>
2026-02-18 11:53:04 +01:00

663 lines
15 KiB
Go

// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package sdk_test
import (
"context"
"crypto/ecdsa"
"crypto/ed25519"
"crypto/elliptic"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"os"
"testing"
"github.com/absmach/supermq/pkg/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
"github.com/ultravioletrs/cocos/pkg/sdk"
"golang.org/x/crypto/sha3"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/status"
)
var (
algoPath = "../../test/manual/algo/lin_reg.py"
dataPath = "../../test/manual/data/iris.csv"
errInappropriateIoctl = errors.New("failed to get terminal width: inappropriate ioctl for device")
)
func TestAlgo(t *testing.T) {
conn, err := grpc.NewClient("passthrough://bufnet", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(bufDialer))
if err != nil {
t.Fatalf("Failed to dial bufnet: %v", err)
}
defer conn.Close()
client := agent.NewAgentServiceClient(conn)
sdk := sdk.NewAgentSDK(client)
algo, err := os.ReadFile(algoPath)
require.NoError(t, err)
algoHash := sha3.Sum256(algo)
algorithmProviderKey, algorithmProviderPubKey := generateKeys(t, "ed25519")
algoProvider1Key, algoProvider1PubKey := generateKeys(t, "ed25519")
algorithm := agent.Algorithm{
Algorithm: algo,
Hash: algoHash,
UserKey: algorithmProviderPubKey,
}
cases := []struct {
name string
err error
algo agent.Algorithm
userKey any
}{
{
name: "Test Algo successfully",
algo: algorithm,
userKey: algorithmProviderKey,
err: nil,
},
{
name: "hash mismatch",
algo: agent.Algorithm{
Algorithm: algo,
Hash: sha3.Sum256([]byte("algo")),
UserKey: algoProvider1PubKey,
},
userKey: algoProvider1Key,
err: errInappropriateIoctl,
},
{
name: "no manifest expected",
algo: agent.Algorithm{
Algorithm: algo,
UserKey: algorithmProviderPubKey,
Hash: algoHash,
},
userKey: algorithmProviderKey,
err: errInappropriateIoctl,
},
{
name: "state not ready",
algo: agent.Algorithm{
Algorithm: algo,
UserKey: algorithmProviderPubKey,
Hash: algoHash,
},
userKey: algorithmProviderKey,
err: errInappropriateIoctl,
},
{
name: "gRPC client error",
algo: agent.Algorithm{
Algorithm: algo,
Hash: algoHash,
UserKey: algoProvider1PubKey,
},
userKey: algoProvider1Key,
err: errors.New("gRPC client error"),
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
svcCall := svc.On("Algo", mock.Anything, mock.Anything).Return(tc.err)
algo, err := os.CreateTemp("", "algo")
require.NoError(t, err)
defer os.Remove(algo.Name())
_, err = algo.Write(algorithm.Algorithm)
require.NoError(t, err)
err = algo.Close()
require.NoError(t, err)
algo, err = os.Open(algo.Name())
require.NoError(t, err)
err = sdk.Algo(context.Background(), algo, nil, tc.userKey)
st, _ := status.FromError(err)
if tc.err != nil {
if st.Message() != tc.err.Error() {
t.Errorf("%s : Expected error message %q, but got %q", tc.name, tc.err.Error(), st.Message())
}
}
svcCall.Unset()
})
}
}
func TestData(t *testing.T) {
conn, err := grpc.NewClient("passthrough://bufnet", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(bufDialer))
if err != nil {
t.Fatalf("Failed to dial bufnet: %v", err)
}
defer conn.Close()
client := agent.NewAgentServiceClient(conn)
sdk := sdk.NewAgentSDK(client)
data, err := os.ReadFile(dataPath)
require.NoError(t, err)
dataHash := sha3.Sum256(data)
dataProviderKey, dataProviderPubKey := generateKeys(t, "ecdsa")
dataProvider1Key, dataProvider1PubKey := generateKeys(t, "ed25519")
dataset := agent.Dataset{
Hash: dataHash,
Dataset: data,
UserKey: dataProviderPubKey,
}
cases := []struct {
name string
data agent.Dataset
userKey any
svcErr error
}{
{
name: "Test data successfully",
data: dataset,
userKey: dataProviderKey,
},
{
name: "undeclared dataset",
data: agent.Dataset{
Dataset: data,
UserKey: dataProvider1PubKey,
Hash: dataHash,
},
userKey: dataProvider1Key,
svcErr: errInappropriateIoctl,
},
{
name: "hash mismatch",
data: agent.Dataset{
Dataset: data,
UserKey: dataProvider1PubKey,
Hash: dataHash,
},
userKey: dataProvider1Key,
svcErr: errInappropriateIoctl,
},
{
name: "all manifest items received",
data: agent.Dataset{
Dataset: data,
UserKey: dataProvider1PubKey,
Hash: dataHash,
},
userKey: dataProvider1Key,
svcErr: errInappropriateIoctl,
},
{
name: "missing dataset file",
data: agent.Dataset{
UserKey: dataProvider1PubKey,
Hash: dataHash,
},
userKey: dataProvider1Key,
svcErr: errors.New("dataset is required"),
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
dataCall := svc.On("Data", mock.Anything, mock.Anything).Return(tc.svcErr)
data, err := os.CreateTemp("", "data")
require.NoError(t, err)
_, err = data.Write(dataset.Dataset)
require.NoError(t, err)
err = data.Close()
require.NoError(t, err)
data, err = os.Open(data.Name())
require.NoError(t, err)
err = sdk.Data(context.Background(), data, tc.data.Filename, tc.userKey)
st, _ := status.FromError(err)
if tc.svcErr != nil {
if st.Message() != tc.svcErr.Error() {
t.Errorf("%s: Expected error message %q, but got %q", tc.name, tc.svcErr.Error(), st.Message())
}
}
dataCall.Unset()
})
}
}
func TestResult(t *testing.T) {
conn, err := grpc.NewClient("passthrough://bufnet", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(bufDialer))
if err != nil {
t.Fatalf("Failed to dial bufnet: %v", err)
}
defer conn.Close()
client := agent.NewAgentServiceClient(conn)
sdk := sdk.NewAgentSDK(client)
resultConsumerKey, _ := generateKeys(t, "ecdsa")
resultConsumer1Key, _ := generateKeys(t, "ed25519")
response := &agent.ResultResponse{
File: []byte{
0x01, 0x02, 0x03, 0x04,
0x05, 0x06, 0x07, 0x08,
},
}
cases := []struct {
name string
userKey any
response *agent.ResultResponse
svcRes []byte
err error
}{
{
name: "Test result successfully",
userKey: resultConsumerKey,
response: response,
svcRes: response.File,
err: nil,
},
{
name: "Test result successfully with ed25519 key type",
userKey: resultConsumer1Key,
response: response,
svcRes: response.File,
err: nil,
},
{
name: "Results not ready",
userKey: resultConsumer1Key,
response: &agent.ResultResponse{
File: []byte{},
},
svcRes: nil,
err: agent.ErrResultsNotReady,
},
{
name: "All manifest items received",
userKey: resultConsumer1Key,
response: &agent.ResultResponse{
File: []byte{},
},
svcRes: nil,
err: agent.ErrAllManifestItemsReceived,
},
{
name: "Undeclared consumer",
userKey: resultConsumer1Key,
response: &agent.ResultResponse{
File: []byte{},
},
svcRes: nil,
err: agent.ErrUndeclaredConsumer,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
svcCall := svc.On("Result", mock.Anything, mock.Anything).Return(tc.svcRes, tc.err)
resultFile, err := os.CreateTemp("", "result")
require.NoError(t, err)
t.Cleanup(func() {
os.Remove(resultFile.Name())
})
err = sdk.Result(context.Background(), tc.userKey, resultFile)
require.NoError(t, resultFile.Close())
st, ok := status.FromError(err)
if !ok {
t.Fatalf("Expected gRPC status error, but got: %v", err)
}
if tc.err != nil {
if st.Message() != tc.err.Error() {
t.Errorf("%s: Expected error message %q, but got %q", tc.name, tc.err.Error(), st.Message())
}
}
res, err := os.ReadFile(resultFile.Name())
require.NoError(t, err)
assert.Equal(t, tc.response.File, res, tc.name)
svcCall.Unset()
})
}
}
func TestAttestation(t *testing.T) {
resultConsumerKey, _ := generateKeys(t, "rsa")
resultConsumer1Key, _ := generateKeys(t, "ed25519")
reportData := make([]byte, 64)
nonce := make([]byte, 64)
report := []byte{
0x01, 0x02, 0x03, 0x04,
0x05, 0x06, 0x07, 0x08,
}
conn, err := grpc.NewClient("passthrough://bufnet", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(bufDialer))
if err != nil {
t.Fatalf("Failed to dial bufnet: %v", err)
}
defer conn.Close()
client := agent.NewAgentServiceClient(conn)
sdk := sdk.NewAgentSDK(client)
_, err = rand.Read(reportData)
require.NoError(t, err)
cases := []struct {
name string
userKey any
reportData [vtpm.SEVNonce]byte
nonce [vtpm.Nonce]byte
response *agent.AttestationResponse
svcRes []byte
err error
}{
{
name: "fetch attestation report successfully",
userKey: resultConsumerKey,
reportData: [vtpm.SEVNonce]byte(reportData),
nonce: [vtpm.Nonce]byte(nonce),
response: &agent.AttestationResponse{
File: report,
},
svcRes: report,
err: nil,
},
{
name: "fetch attestation report with different key type",
userKey: resultConsumer1Key,
reportData: [vtpm.SEVNonce]byte(reportData),
nonce: [vtpm.Nonce]byte(nonce),
response: &agent.AttestationResponse{
File: report,
},
svcRes: report,
err: nil,
},
{
name: "failed to fetch attestation report",
userKey: resultConsumerKey,
reportData: [vtpm.SEVNonce]byte(reportData),
nonce: [vtpm.Nonce]byte(nonce),
response: &agent.AttestationResponse{
File: []byte{},
},
err: nil,
},
{
name: "invalid report data",
userKey: resultConsumerKey,
reportData: [vtpm.SEVNonce]byte{},
nonce: [vtpm.Nonce]byte(nonce),
response: &agent.AttestationResponse{
File: []byte{},
},
svcRes: nil,
err: errors.New("invalid report data"),
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
svcCall := svc.On("Attestation", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(tc.svcRes, tc.err)
file, err := os.CreateTemp("", "attestation")
require.NoError(t, err)
t.Cleanup(func() {
os.Remove(file.Name())
})
err = sdk.Attestation(context.Background(), tc.reportData, tc.nonce, 0, file)
require.NoError(t, file.Close())
st, ok := status.FromError(err)
if !ok {
t.Fatalf("Expected gRPC status error, but got: %v", err)
}
if tc.err != nil {
if st.Message() != tc.err.Error() {
t.Errorf("%s: Expected error message %q, but got %q", tc.name, tc.err.Error(), st.Message())
}
}
res, err := os.ReadFile(file.Name())
require.NoError(t, err)
assert.Equal(t, tc.response.File, res, tc.name)
svcCall.Unset()
})
}
}
func TestAttestationToken(t *testing.T) {
reportData := make([]byte, 64)
nonce := make([]byte, 64)
report := []byte{
0x01, 0x02, 0x03, 0x04,
0x05, 0x06, 0x07, 0x08,
}
conn, err := grpc.NewClient("passthrough://bufnet", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(bufDialer))
if err != nil {
t.Fatalf("Failed to dial bufnet: %v", err)
}
defer conn.Close()
client := agent.NewAgentServiceClient(conn)
sdk := sdk.NewAgentSDK(client)
_, err = rand.Read(reportData)
require.NoError(t, err)
cases := []struct {
name string
nonce [vtpm.Nonce]byte
response *agent.AttestationTokenResponse
svcRes []byte
err error
}{
{
name: "fetch attestation token successfully",
nonce: [vtpm.Nonce]byte(nonce),
response: &agent.AttestationTokenResponse{
File: report,
},
svcRes: report,
err: nil,
},
{
name: "failed to fetch attestation token",
nonce: [vtpm.Nonce]byte(nonce),
response: &agent.AttestationTokenResponse{
File: []byte{},
},
err: nil,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
svcCall := svc.On("AzureAttestationToken", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(tc.svcRes, tc.err)
file, err := os.CreateTemp("", "attestation")
require.NoError(t, err)
t.Cleanup(func() {
os.Remove(file.Name())
})
err = sdk.AttestationToken(context.Background(), tc.nonce, 0, file)
require.NoError(t, file.Close())
st, ok := status.FromError(err)
if !ok {
t.Fatalf("Expected gRPC status error, but got: %v", err)
}
if tc.err != nil {
if st.Message() != tc.err.Error() {
t.Errorf("%s: Expected error message %q, but got %q", tc.name, tc.err.Error(), st.Message())
}
}
res, err := os.ReadFile(file.Name())
require.NoError(t, err)
assert.Equal(t, tc.response.File, res, tc.name)
svcCall.Unset()
})
}
}
func TestIMAMeasurements(t *testing.T) {
conn, err := grpc.NewClient("passthrough://bufnet", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(bufDialer))
if err != nil {
t.Fatalf("Failed to dial bufnet: %v", err)
}
defer conn.Close()
client := agent.NewAgentServiceClient(conn)
sdk := sdk.NewAgentSDK(client)
response := &agent.IMAMeasurementsResponse{
File: []byte{
0x01, 0x02, 0x03, 0x04,
0x05, 0x06, 0x07, 0x08,
0x01, 0x02, 0x03, 0x04,
0x05, 0x06, 0x07, 0x08,
0x01, 0x02, 0x03, 0x04,
0x05, 0x06, 0x07, 0x08,
},
}
cases := []struct {
name string
response *agent.IMAMeasurementsResponse
svcRes []byte
err error
}{
{
name: "fetch IMA measurements successfully",
response: response,
svcRes: response.File,
err: nil,
},
{
name: "failed to fetch IMA measurements",
response: &agent.IMAMeasurementsResponse{File: []byte{}},
svcRes: nil,
err: errors.New("failed to fetch IMA measurements"),
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
svcCall := svc.On("IMAMeasurements", mock.Anything).Return(tc.svcRes, tc.svcRes, tc.err)
file, err := os.CreateTemp("", "ima_measurements")
require.NoError(t, err)
t.Cleanup(func() {
os.Remove(file.Name())
})
_, err = sdk.IMAMeasurements(context.Background(), file)
require.NoError(t, file.Close())
st, ok := status.FromError(err)
if !ok {
t.Fatalf("Expected gRPC status error, but got: %v", err)
}
if tc.err != nil {
if st.Message() != tc.err.Error() {
t.Errorf("%s: Expected error message %q, but got %q", tc.name, tc.err.Error(), st.Message())
}
}
res, err := os.ReadFile(file.Name())
require.NoError(t, err)
assert.Equal(t, tc.response.File, res, tc.name)
svcCall.Unset()
})
}
}
func generateKeys(t *testing.T, keyType string) (priv any, pub []byte) {
switch keyType {
case "ecdsa":
privEcdsaKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privEcdsaKey.PublicKey)
require.NoError(t, err)
return privEcdsaKey, pubKeyBytes
case "ed25519":
pubEd25519Key, privEd25519Key, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
pubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
require.NoError(t, err)
return privEd25519Key, pubKey
default:
privKey, err := rsa.GenerateKey(rand.Reader, 2048)
require.NoError(t, err)
pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
require.NoError(t, err)
return privKey, pubKeyBytes
}
}