Compare commits

..

1 Commits

Author SHA1 Message Date
Drasko Draskovic 36fefd554e Fix usage
Signed-off-by: Drasko Draskovic <drasko.draskovic@gmail.com>
2024-04-16 00:41:51 +02:00
243 changed files with 5058 additions and 31531 deletions
+33
View File
@@ -0,0 +1,33 @@
## Jaeger
COCOS_JAEGER_PORT=6831
COCOS_JAEGER_FRONTEND=16686
COCOS_JAEGER_COLLECTOR=14268
COCOS_JAEGER_CONFIGS=5778
COCOS_JAEGER_URL=http://jaeger:14268/api/traces
COCOS_JAEGER_TRACE_RATIO=1.0
## Core Services
### Manager
MANAGER_HTTP_HOST="cocos-manager"
MANAGER_HTTP_PORT=9021
MANAGER_HTTP_SERVER_CERT=""
MANAGER_HTTP_SERVER_KEY=""
MANAGER_GRPC_HOST="cocos-manager"
MANAGER_GRPC_PORT=7003
MANAGER_GRPC_SERVER_CERT=""
MANAGER_GRPC_SERVER_KEY=""
AGENT_GRPC_URL="192.168.100.4:7002"
AGENT_GRPC_TIMEOUT=""
AGENT_GRPC_CA_CERTS=""
AGENT_GRPC_CLIENT_TLS=""
MANAGER_INSTANCE_ID=""
MANAGER_LOG_LEVEL=debug
MANAGER_QEMU_USE_SUDO=false
MANAGER_QEMU_ENABLE_SEV=false
MANAGER_QEMU_SEV_CBITPOS=51
MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/OVMF/OVMF_CODE.fd
MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/OVMF/OVMF_VARS.fd
# Docker image tag
COCOS_RELEASE_TAG=latest
-18
View File
@@ -1,18 +0,0 @@
version: 2
updates:
- package-ecosystem: "cargo"
directory: "/scripts/attestation_policy"
schedule:
interval: "weekly"
day: "monday"
groups:
rs-dependencies:
patterns:
- "*"
- package-ecosystem: "gomod"
directories:
- "/"
schedule:
interval: "weekly"
day: "monday"
+6 -7
View File
@@ -9,7 +9,6 @@ on:
- "pkg/manager/*.pb.go"
- "agent/agent.proto"
- "agent/*.pb.go"
pull_request:
branches:
- main
@@ -30,13 +29,13 @@ jobs:
- name: Install Go
uses: actions/setup-go@v5
with:
go-version: 1.23.x
go-version: 1.21.x
- name: Set up protoc
run: |
PROTOC_VERSION=29.0
PROTOC_GEN_VERSION=v1.36.5
PROTOC_GRPC_VERSION=v1.5.1
PROTOC_VERSION=25.3
PROTOC_GEN_VERSION=v1.31.0
PROTOC_GRPC_VERSION=v1.3.0
# Download and install protoc
PROTOC_ZIP=protoc-$PROTOC_VERSION-linux-x86_64.zip
@@ -56,7 +55,7 @@ jobs:
- name: Set up Cocos-AI
run: |
# Rename .pb.go files to .pb.go.tmp to prevent conflicts
for p in $(ls manager/*.pb.go); do
for p in $(ls pkg/manager/*.pb.go); do
mv $p $p.tmp
done
@@ -68,7 +67,7 @@ jobs:
make protoc
# Compare generated Go files with the original ones
for p in $(ls manager/*.pb.go); do
for p in $(ls pkg/manager/*.pb.go); do
if ! cmp -s $p $p.tmp; then
echo "Proto file and generated Go file $p are out of sync!"
exit 1
+58 -65
View File
@@ -1,80 +1,73 @@
name: Build and Release Hal
name: Build and Release
on:
push:
tags:
- "*"
- '*'
jobs:
build:
runs-on: ubuntu-latest
timeout-minutes: 120
steps:
- name: Free Disk Space
run: |
sudo apt-get clean
sudo rm -rf /usr/share/dotnet
sudo rm -rf /opt/ghc
sudo rm -rf "/usr/local/share/boost"
sudo rm -rf "$AGENT_TOOLSDIRECTORY"
df -h
- name: Update Ubuntu
run: |
sudo apt-get update
sudo apt-get upgrade -y
- name: Install Go
uses: actions/setup-go@v5
with:
go-version: 1.21.x
cache-dependency-path: "go.sum"
- name: Update Ubuntu
run: |
sudo apt-get update
sudo apt-get upgrade -y
sudo apt-get clean
sudo apt-get autoremove -y
df -h
- name: Checkout cocos
uses: actions/checkout@v4
with:
repository: 'ultravioletrs/cocos'
path: cocos
- name: Maximize build space
uses: easimon/maximize-build-space@master
with:
root-reserve-mb: 35000
swap-size-mb: 1024
remove-dotnet: "true"
remove-android: "true"
- name: Check free space
run: |
echo "Free space:"
df -h
- name: Checkout buildroot
uses: actions/checkout@v4
with:
repository: 'buildroot/buildroot'
path: buildroot
- name: Install Go
uses: actions/setup-go@v5
with:
go-version: 1.23.x
cache-dependency-path: "go.sum"
- name: Build
run: |
cd buildroot
make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
make
- name: Checkout cocos
uses: actions/checkout@v4
with:
repository: "ultravioletrs/cocos"
path: cocos
- name: Create Release
id: create_release
uses: actions/create-release@latest
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
tag_name: ${{ github.ref }}
release_name: Release ${{ github.ref }}
draft: false
prerelease: false
- name: Checkout buildroot
uses: actions/checkout@v4
with:
repository: "buildroot/buildroot"
path: buildroot
ref: 2025.05-rc1
- name: Upload Release Asset
id: upload-release-kernel
uses: actions/upload-release-asset@latest
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
upload_url: ${{ steps.create_release.outputs.upload_url }}
asset_path: ./buildroot/output/images/bzImage
asset_name: bzImage
asset_content_type: application/octet-stream
- name: Build hal
run: |
cd buildroot
make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
make
- name: Build cocos
run: |
cd cocos
make
- name: Release
uses: softprops/action-gh-release@v2
with:
files: |
buildroot/output/images/bzImage
buildroot/output/images/rootfs.cpio.gz
cocos/build/cocos-agent
cocos/build/cocos-cli
cocos/build/cocos-manager
- name: Upload Release Asset
id: upload-release-rootfs
uses: actions/upload-release-asset@latest
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
upload_url: ${{ steps.create_release.outputs.upload_url }}
asset_path: ./buildroot/output/images/rootfs.cpio.gz
asset_name: rootfs.cpio.gz
asset_content_type: application/gzip
+5 -33
View File
@@ -19,44 +19,16 @@ jobs:
- name: Install Go
uses: actions/setup-go@v5
with:
go-version: 1.23.x
go-version: 1.21.x
- name: golangci-lint
uses: golangci/golangci-lint-action@v7
uses: golangci/golangci-lint-action@v3
with:
version: v2.0.2
version: latest
- name: Build
run: |
make
- name: Create coverage directory
run: |
mkdir coverage
- name: Run Agent tests
run: go test -v --race -covermode=atomic -coverprofile coverage/agent.out ./agent/...
- name: Run cli tests
run: go test -v --race -covermode=atomic -coverprofile coverage/cli.out ./cli/...
- name: Run cmd tests
run: go test -v --race -covermode=atomic -coverprofile coverage/cmd.out ./cmd/...
- name: Run internal tests
run: go test -v --race -covermode=atomic -coverprofile coverage/internal.out ./internal/...
- name: Run pkg tests
run: go test -v --race -covermode=atomic -coverprofile coverage/pkg.out ./pkg/...
- name: Run manager tests
run: sudo go test -v --race -covermode=atomic -coverprofile coverage/manager.out ./manager/...
- name: Upload results to Codecov
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
files: ./coverage/*.out
codecov_yml_path: codecov.yml
verbose: true
- name: Run tests
run: go test -v --race -covermode=atomic -coverprofile cover.out ./...
-41
View File
@@ -1,41 +0,0 @@
name: Rust CI Pipeline
on:
push:
branches:
- main
paths:
- "scripts/attestation_policy/**"
- ".github/workflows/rust.yaml"
pull_request:
branches:
- main
paths:
- "scripts/attestation_policy/**"
- ".github/workflows/rust.yaml"
env:
CARGO_TERM_COLOR: always
jobs:
rust-check:
runs-on: ubuntu-latest
defaults:
run:
working-directory: ./scripts/attestation_policy
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Check cargo
run: cargo check --release --all-targets
- name: Check formatting
run: cargo fmt --all -- --check
- name: Run linter
run: cargo clippy -- -D warnings
- name: Build for all features
run: cargo build --release --all-features
+3 -22
View File
@@ -1,27 +1,8 @@
build
build
cmd/manager/img
cmd/manager/iso
cmd/manager/tmp
.cov
*.pem
dist/
*.zip
*.spec
*.tar
# Generated by Cargo
# will have compiled files and executables
debug/
target/
# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
Cargo.lock
# These are backup files generated by rustfmt
**/*.rs.bk
# MSVC Windows builds of rustc generate these, which store debugging information
*.pdb
+69 -83
View File
@@ -1,100 +1,86 @@
version: "2"
run:
timeout: 3m
issues:
max-issues-per-linter: 10
max-same-issues: 10
linters-settings:
importas:
no-unaliased: true
no-extra-aliases: false
alias:
- pkg: github.com/absmach/magistrala/logger
alias: mglog
gocritic:
enabled-checks:
- captLocal
- dupImport
- singleCaseSwitch
- switchTrue
- importShadow
- httpNoBody
- paramTypeCombine
- emptyStringTest
- builtinShadow
- exposedSyncMutex
disabled-checks:
- appendAssign
enabled-tags:
- diagnostic
disabled-tags:
- performance
- style
- experimental
- opinionated
misspell:
ignore-words:
- "mosquitto"
stylecheck:
checks: ["-ST1000", "-ST1003", "-ST1020", "-ST1021", "-ST1022"]
goheader:
template: |-
Copyright (c) Ultraviolet
SPDX-License-Identifier: Apache-2.0
linters:
default: none
disable-all: true
enable:
- importas
- gocritic
- gosimple
- errcheck
- govet
- unused
- goconst
- godot
- godox
- ineffassign
- misspell
- stylecheck
- whitespace
- gci
- gofmt
- goimports
- loggercheck
- goheader
- asasalint
- asciicheck
- bidichk
- contextcheck
- copyloopvar
- decorder
- dogsled
- dupword
- errcheck
- errchkjson
- errname
- execinquery
- exportloopref
- ginkgolinter
- gocheckcompilerdirectives
- goconst
- gocritic
- godot
- godox
- goheader
- gofumpt
- goprintffuncname
- govet
- importas
- ineffassign
- loggercheck
- makezero
- mirror
- misspell
- nakedret
- staticcheck
- unused
- whitespace
settings:
gocritic:
enabled-checks:
- dupImport
- importShadow
- httpNoBody
- paramTypeCombine
- emptyStringTest
- builtinShadow
- exposedSyncMutex
disabled-checks:
- appendAssign
enabled-tags:
- diagnostic
disabled-tags:
- performance
- style
- experimental
- opinionated
goheader:
template: |-
Copyright (c) Ultraviolet
SPDX-License-Identifier: Apache-2.0
importas:
alias:
- pkg: github.com/absmach/magistrala/logger
alias: mglog
no-unaliased: true
no-extra-aliases: false
staticcheck:
checks:
- -ST1000
- -ST1003
- -ST1020
- -ST1021
- -ST1022
exclusions:
generated: lax
presets:
- comments
- common-false-positives
- legacy
- std-error-handling
rules:
- linters:
- makezero
text: with non-zero initialized length
paths:
- third_party$
- builtin$
- examples$
issues:
max-issues-per-linter: 10
max-same-issues: 10
formatters:
enable:
- gci
- gofmt
- gofumpt
- goimports
exclusions:
generated: lax
paths:
- third_party$
- builtin$
- examples$
- dupword
+7 -47
View File
@@ -1,18 +1,12 @@
BUILD_DIR = build
SERVICES = manager agent cli
ATTESTATION_POLICY = attestation_policy
CGO_ENABLED ?= 1
CGO_ENABLED ?= 0
GOARCH ?= amd64
VERSION ?= $(shell git describe --abbrev=0 --tags --always)
COMMIT ?= $(shell git rev-parse HEAD)
TIME ?= $(shell date +%F_%T)
EMBED_ENABLED ?= 0
INSTALL_DIR ?= /usr/local/bin
CONFIG_DIR ?= /etc/cocos
SERVICE_NAME ?= cocos-manager
SERVICE_DIR ?= /etc/systemd/system
SERVICE_FILE = init/systemd/$(SERVICE_NAME).service
IGVM_BUILD_SCRIPT := ./scripts/igvmmeasure/igvm.sh
empty:=
space:= $(empty) $(empty)
define compile_service
CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) \
@@ -20,50 +14,16 @@ define compile_service
-X 'github.com/absmach/magistrala.BuildTime=$(TIME)' \
-X 'github.com/absmach/magistrala.Version=$(VERSION)' \
-X 'github.com/absmach/magistrala.Commit=$(COMMIT)'" \
$(if $(filter 1,$(EMBED_ENABLED)),-tags "embed",) \
-o ${BUILD_DIR}/cocos-$(1) cmd/$(1)/main.go
endef
.PHONY: all $(SERVICES) $(ATTESTATION_POLICY) install clean
.PHONY: all $(SERVICES)
all: $(SERVICES)
$(SERVICES):
$(call compile_service,$@)
@if [ "$@" = "cli" ] || [ "$@" = "manager" ]; then $(MAKE) build-igvm; fi
$(ATTESTATION_POLICY):
$(MAKE) -C ./scripts/attestation_policy
$(SERVICES):
$(call compile_service,$(@))
protoc:
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/agent.proto
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative manager/manager.proto
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/events/events.proto
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/cvms/cvms.proto
mocks:
mockery --config ./mockery.yml
install: $(SERVICES)
install -d $(INSTALL_DIR)
install $(BUILD_DIR)/cocos-cli $(INSTALL_DIR)/cocos-cli
install $(BUILD_DIR)/cocos-manager $(INSTALL_DIR)/cocos-manager
install -d $(CONFIG_DIR)
install cocos-manager.env $(CONFIG_DIR)/cocos-manager.env
clean:
rm -rf $(BUILD_DIR)
run: install_service
sudo systemctl start $(SERVICE_NAME).service
stop:
sudo systemctl stop $(SERVICE_NAME).service
install_service:
sudo install -m 644 $(SERVICE_FILE) $(SERVICE_DIR)/$(SERVICE_NAME).service
sudo systemctl daemon-reload
build-igvm:
@echo "Running build script for igvmmeasure..."
@$(IGVM_BUILD_SCRIPT)
protoc -I. --go_out=./pkg --go_opt=paths=source_relative --go-grpc_out=./pkg --go-grpc_opt=paths=source_relative manager/manager.proto
+38 -57
View File
@@ -1,80 +1,61 @@
<div align="center">
# Cocos AI 🥥
**Confidential Computing System for AI**
**Made with ❤️ by [Ultraviolet](https://ultraviolet.rs/)**
[![codecov](https://codecov.io/gh/ultravioletrs/cocos/graph/badge.svg?token=HX01LR01K9)](https://codecov.io/gh/ultravioletrs/cocos)
[![Go report card](https://goreportcard.com/badge/github.com/ultravioletrs/cocos)](https://goreportcard.com/report/github.com/ultravioletrs/cocos)
[![License](https://img.shields.io/badge/license-Apache--2.0-blue)](LICENSE)
### [Guide](https://docs.cocos.ultraviolet.rs) | [Contributing](CONTRIBUTING.md) | [Website](https://cocos.ai/)
</div>
## Introduction 🚀
Cocos AI is a **cutting-edge platform** designed to enable secure multiparty computation (SMPC) using **Confidential Computing** and **Trusted Execution Environments (TEEs)**.
It empowers organizations to collaboratively process sensitive data for AI/ML workloads while ensuring:
- 🔒 **Data Privacy**: Your data stays encrypted and secure throughout the computation.
- 🛡️ **Trust and Integrity**: Protected by hardware enclaves with robust remote attestation protocols.
- 🤝 **Seamless Collaboration**: Multiple organizations can work together without exposing sensitive information.
# Cocos AI
[Cocos AI (Confdential Computing System for AI/ML)][cocos] is a platform for secure multiparty computation (SMPC)
based on the [Confidential Computing][cc] and [Trusted Execution Environments (TEEs)][tee].
<p align="center">
<img src="https://cocos.ai/images/Collaborative%20AI.drawio.svg" alt="Cocos AI Illustration" width="400" height="400">
<img src="https://cocos.ai/images/Collaborative%20AI.drawio.svg" width="500" height="500">
</p>
## Features 🛠️
With Cocos AI it becomes possible to run AI/ML workloads on combined datasets from multiple organizations
while guaranteeing the privacy and security of the data and the algorithms.
Data is always encrypted, protected by hardware secure enclaves (Trusted Execution Environments),
attested via secure remote attestation protocols, and invisible to cloud processors or any other
3rd party to which computation is offloaded.
Cocos AI provides essential features for secure and efficient collaborative AI/ML:
## Features
- 🖥️ **TEE Enablement and Monitoring**: Secure VM management for deploying and monitoring workloads.
- 🛡️ **Hardware Abstraction Layer (HAL)**: Built on a hardened Linux kernel, secure bootloader, and minimal root filesystem (minimal TCB).
- 🕵️ **In-Enclave Agent and Networking Controller**: Essential system software for managing secure workloads.
- 🔒 **Encrypted Data Transfer**: Asynchronous data transfer and secure result delivery.
- 🛠️ **API for Platform Manipulation**: Programmatic control for managing workloads.
-**Attestation and Verification Tools**: Hardware- and software-supported attestation for integrity assurance.
- 🖱️ **Command-Line Interface (CLI)**: A user-friendly CLI for system interaction.
Cocos AI is implementing the following features:
## 🚀 Quick Start
- TEE enablement, deployment and monitoring (secure VM manager)
- HAL for TEEs based on hardened Linux kernel, secure bootloader and custom-tailored embedded rootfs for minimal TCB
- In-enclave agent, netowrking controller and other system software
- Encrypted asynchronous data transfer and result delivery
- API for programmable platform manipulation
- HW and SW supported attestation with verification tools
- CLI for system interaction
## Usage
Clone the repo and create binaries:
### Clone the Repository and Build Binaries
```bash
git clone git@github.com:ultravioletrs/cocos.git
make
```
This will generate three binaries:
This will create 3 binaries:
```bash
ls build/
# cocos-agent cocos-cli cocos-manager
```
### Deployment Overview:
- **Manager**: Deploy on the AMD SEV-SNP host to orchestrate workloads.
- **Agent**: Build into the [EOS](https://github.com/ultravioletrs/eos)-based HAL for secure enclave management.
- **CLI**: Interact with remote agents to control operations.
- Manager can be deployed on the AMD SEV-SNP host
- Agent can be built into [EOS][eos]-based HAL
- CLI can be used to communicate to remote Agent.
## 📚 Documentation
## Documentation
Comprehensive documentation is available at the [official documentation page](https://docs.cocos.ultraviolet.rs).
For CLI usage details, visit the [CLI Documentation](https://docs.cocos.ultraviolet.rs/cli).
Project documentation is hosted at [Cocos AI official docs page][docs].
Documentation is automatically generated from the [docs repository](https://github.com/ultravioletrs/docs). Contributions to documentation are welcome!
Documentation is generated from the [docs repository](https://github.com/ultravioletrs/docs).
## 🛡️ License
## License
Cocos AI is published under permissive open-source [Apache-2.0](LICENSE) license.
Cocos AI is published under the permissive open-source [Apache-2.0](LICENSE) license. Contributions are encouraged and appreciated!
## 🌐 Links and Resources
- [Cocos AI Website](https://cocos.ai/)
- [Official Releases](https://github.com/ultravioletrs/cocos/releases)
- [Confidential Computing Overview](https://confidentialcomputing.io/white-papers-reports/)
- [Trusted Execution Environments (TEEs)](https://en.wikipedia.org/wiki/Trusted_execution_environment)
>This work has been partially supported by the ELASTIC project, which received funding from the Smart Networks and Services Joint Undertaking (SNS JU) under the European Unions Horizon Europe research and innovation programme under Grant Agreement No 101139067. Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union. Neither the European Union nor the granting authority can be held responsible for them.
[cc]: https://confidentialcomputing.io/white-papers-reports/
[cocos]: https://cocos.ai/
[rel]: https://github.com/ultravioletrs/cocos/releases
[tee]: https://en.wikipedia.org/wiki/Trusted_execution_environment
[docs]: https://docs.cocos.ultraviolet.rs
[cli]: https://docs.cocos.ultraviolet.rs/cli
[eos]: https://github.com/ultravioletrs/eos
+11 -10
View File
@@ -6,16 +6,17 @@ Agent service provides a barebones HTTP and gRPC API and Service interface imple
The service is configured using the environment variables from the following table. Note that any unset variables will be replaced with their default values.
| Variable | Description | Default |
| ------------------------------ | ------------------------------------------------------------------------------------------------------------- | ------------------------------ |
| AGENT_LOG_LEVEL | Log level for agent service (debug, info, warn, error) | debug |
| AGENT_CVM_GRPC_HOST | Agent service gRPC host | "" |
| AGENT_CVM_GRPC_PORT | Agent service gRPC port | 7001 |
| AGENT_CVM_GRPC_SERVER_CERT | Path to gRPC server certificate in pem format | "" |
| AGENT_CVM_GRPC_SERVER_KEY | Path to gRPC server key in pem format | "" |
| AGENT_CVM_GRPC_SERVER_CA_CERTS | Path to gRPC server CA certificate | "" |
| AGENT_CVM_GRPC_CLIENT_CA_CERTS | Path to gRPC client CA certificate | "" |
| AGENT_CVM_CA_URL | URL for CA service, if provided it will be used for certificate generation, used only with aTLS at the moment | "" |
| Variable | Description | Default |
| ----------------------------- | ------------------------------------------------------ | ------------------------------ |
| AGENT_LOG_LEVEL | Log level for agent service (debug, info, warn, error) | info |
| AGENT_GRPC_HOST | Agent service gRPC host | "" |
| AGENT_GRPC_PORT | Agent service gRPC port | 7002 |
| AGENT_GRPC_SERVER_CERT | Path to gRPC server certificate in pem format | "" |
| AGENT_GRPC_SERVER_KEY | Path to gRPC server key in pem format | "" |
| AGENT_GRPC_SERVER_CA_CERTS | Path to gRPC server CA certificate | "" |
| AGENT_GRPC_CLIENT_CA_CERTS | Path to gRPC client CA certificate | "" |
| COCOS_NOTIFICATION_SERVER_URL | Server to receive notification events from agent. | http:/localhost:9000 |
## Deployment
+278 -367
View File
@@ -3,8 +3,8 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.5
// protoc v5.29.0
// protoc-gen-go v1.31.0
// protoc v4.25.3
// source: agent/agent.proto
package agent
@@ -14,7 +14,6 @@ import (
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
reflect "reflect"
sync "sync"
unsafe "unsafe"
)
const (
@@ -25,18 +24,22 @@ const (
)
type AlgoRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
Requirements []byte `protobuf:"bytes,2,opt,name=requirements,proto3" json:"requirements,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
Provider string `protobuf:"bytes,2,opt,name=provider,proto3" json:"provider,omitempty"`
Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"`
}
func (x *AlgoRequest) Reset() {
*x = AlgoRequest{}
mi := &file_agent_agent_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_agent_agent_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *AlgoRequest) String() string {
@@ -47,7 +50,7 @@ func (*AlgoRequest) ProtoMessage() {}
func (x *AlgoRequest) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[0]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -69,24 +72,33 @@ func (x *AlgoRequest) GetAlgorithm() []byte {
return nil
}
func (x *AlgoRequest) GetRequirements() []byte {
func (x *AlgoRequest) GetProvider() string {
if x != nil {
return x.Requirements
return x.Provider
}
return nil
return ""
}
func (x *AlgoRequest) GetId() string {
if x != nil {
return x.Id
}
return ""
}
type AlgoResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
}
func (x *AlgoResponse) Reset() {
*x = AlgoResponse{}
mi := &file_agent_agent_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_agent_agent_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *AlgoResponse) String() string {
@@ -97,7 +109,7 @@ func (*AlgoResponse) ProtoMessage() {}
func (x *AlgoResponse) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[1]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -113,18 +125,22 @@ func (*AlgoResponse) Descriptor() ([]byte, []int) {
}
type DataRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
Dataset []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
Filename string `protobuf:"bytes,2,opt,name=filename,proto3" json:"filename,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Dataset []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
Provider string `protobuf:"bytes,2,opt,name=provider,proto3" json:"provider,omitempty"`
Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"`
}
func (x *DataRequest) Reset() {
*x = DataRequest{}
mi := &file_agent_agent_proto_msgTypes[2]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_agent_agent_proto_msgTypes[2]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *DataRequest) String() string {
@@ -135,7 +151,7 @@ func (*DataRequest) ProtoMessage() {}
func (x *DataRequest) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[2]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -157,24 +173,33 @@ func (x *DataRequest) GetDataset() []byte {
return nil
}
func (x *DataRequest) GetFilename() string {
func (x *DataRequest) GetProvider() string {
if x != nil {
return x.Filename
return x.Provider
}
return ""
}
func (x *DataRequest) GetId() string {
if x != nil {
return x.Id
}
return ""
}
type DataResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
}
func (x *DataResponse) Reset() {
*x = DataResponse{}
mi := &file_agent_agent_proto_msgTypes[3]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_agent_agent_proto_msgTypes[3]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *DataResponse) String() string {
@@ -185,7 +210,7 @@ func (*DataResponse) ProtoMessage() {}
func (x *DataResponse) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[3]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -201,16 +226,20 @@ func (*DataResponse) Descriptor() ([]byte, []int) {
}
type ResultRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Consumer string `protobuf:"bytes,1,opt,name=consumer,proto3" json:"consumer,omitempty"`
}
func (x *ResultRequest) Reset() {
*x = ResultRequest{}
mi := &file_agent_agent_proto_msgTypes[4]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_agent_agent_proto_msgTypes[4]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *ResultRequest) String() string {
@@ -221,7 +250,7 @@ func (*ResultRequest) ProtoMessage() {}
func (x *ResultRequest) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[4]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -236,18 +265,28 @@ func (*ResultRequest) Descriptor() ([]byte, []int) {
return file_agent_agent_proto_rawDescGZIP(), []int{4}
}
func (x *ResultRequest) GetConsumer() string {
if x != nil {
return x.Consumer
}
return ""
}
type ResultResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
}
func (x *ResultResponse) Reset() {
*x = ResultResponse{}
mi := &file_agent_agent_proto_msgTypes[5]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_agent_agent_proto_msgTypes[5]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *ResultResponse) String() string {
@@ -258,7 +297,7 @@ func (*ResultResponse) ProtoMessage() {}
func (x *ResultResponse) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[5]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -281,19 +320,20 @@ func (x *ResultResponse) GetFile() []byte {
}
type AttestationRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
TeeNonce []byte `protobuf:"bytes,1,opt,name=teeNonce,proto3" json:"teeNonce,omitempty"` // Should be less or equal 64 bytes.
VtpmNonce []byte `protobuf:"bytes,2,opt,name=vtpmNonce,proto3" json:"vtpmNonce,omitempty"` // Should be less or equal 32 bytes.
Type int32 `protobuf:"varint,3,opt,name=type,proto3" json:"type,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
ReportData []byte `protobuf:"bytes,1,opt,name=report_data,json=reportData,proto3" json:"report_data,omitempty"` // Should be of length 64.
}
func (x *AttestationRequest) Reset() {
*x = AttestationRequest{}
mi := &file_agent_agent_proto_msgTypes[6]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_agent_agent_proto_msgTypes[6]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *AttestationRequest) String() string {
@@ -304,7 +344,7 @@ func (*AttestationRequest) ProtoMessage() {}
func (x *AttestationRequest) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[6]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -319,39 +359,28 @@ func (*AttestationRequest) Descriptor() ([]byte, []int) {
return file_agent_agent_proto_rawDescGZIP(), []int{6}
}
func (x *AttestationRequest) GetTeeNonce() []byte {
func (x *AttestationRequest) GetReportData() []byte {
if x != nil {
return x.TeeNonce
return x.ReportData
}
return nil
}
func (x *AttestationRequest) GetVtpmNonce() []byte {
if x != nil {
return x.VtpmNonce
}
return nil
}
func (x *AttestationRequest) GetType() int32 {
if x != nil {
return x.Type
}
return 0
}
type AttestationResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
}
func (x *AttestationResponse) Reset() {
*x = AttestationResponse{}
mi := &file_agent_agent_proto_msgTypes[7]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_agent_agent_proto_msgTypes[7]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *AttestationResponse) String() string {
@@ -362,7 +391,7 @@ func (*AttestationResponse) ProtoMessage() {}
func (x *AttestationResponse) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[7]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -384,308 +413,91 @@ func (x *AttestationResponse) GetFile() []byte {
return nil
}
type IMAMeasurementsRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *IMAMeasurementsRequest) Reset() {
*x = IMAMeasurementsRequest{}
mi := &file_agent_agent_proto_msgTypes[8]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *IMAMeasurementsRequest) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*IMAMeasurementsRequest) ProtoMessage() {}
func (x *IMAMeasurementsRequest) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[8]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use IMAMeasurementsRequest.ProtoReflect.Descriptor instead.
func (*IMAMeasurementsRequest) Descriptor() ([]byte, []int) {
return file_agent_agent_proto_rawDescGZIP(), []int{8}
}
type IMAMeasurementsResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
Pcr10 []byte `protobuf:"bytes,2,opt,name=pcr10,proto3" json:"pcr10,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *IMAMeasurementsResponse) Reset() {
*x = IMAMeasurementsResponse{}
mi := &file_agent_agent_proto_msgTypes[9]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *IMAMeasurementsResponse) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*IMAMeasurementsResponse) ProtoMessage() {}
func (x *IMAMeasurementsResponse) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[9]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use IMAMeasurementsResponse.ProtoReflect.Descriptor instead.
func (*IMAMeasurementsResponse) Descriptor() ([]byte, []int) {
return file_agent_agent_proto_rawDescGZIP(), []int{9}
}
func (x *IMAMeasurementsResponse) GetFile() []byte {
if x != nil {
return x.File
}
return nil
}
func (x *IMAMeasurementsResponse) GetPcr10() []byte {
if x != nil {
return x.Pcr10
}
return nil
}
type AttestationResultRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
TokenNonce []byte `protobuf:"bytes,1,opt,name=tokenNonce,proto3" json:"tokenNonce,omitempty"` // Should be less or equal 32 bytes
Type int32 `protobuf:"varint,3,opt,name=type,proto3" json:"type,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *AttestationResultRequest) Reset() {
*x = AttestationResultRequest{}
mi := &file_agent_agent_proto_msgTypes[10]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *AttestationResultRequest) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*AttestationResultRequest) ProtoMessage() {}
func (x *AttestationResultRequest) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[10]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use AttestationResultRequest.ProtoReflect.Descriptor instead.
func (*AttestationResultRequest) Descriptor() ([]byte, []int) {
return file_agent_agent_proto_rawDescGZIP(), []int{10}
}
func (x *AttestationResultRequest) GetTokenNonce() []byte {
if x != nil {
return x.TokenNonce
}
return nil
}
func (x *AttestationResultRequest) GetType() int32 {
if x != nil {
return x.Type
}
return 0
}
type AttestationResultResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *AttestationResultResponse) Reset() {
*x = AttestationResultResponse{}
mi := &file_agent_agent_proto_msgTypes[11]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *AttestationResultResponse) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*AttestationResultResponse) ProtoMessage() {}
func (x *AttestationResultResponse) ProtoReflect() protoreflect.Message {
mi := &file_agent_agent_proto_msgTypes[11]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use AttestationResultResponse.ProtoReflect.Descriptor instead.
func (*AttestationResultResponse) Descriptor() ([]byte, []int) {
return file_agent_agent_proto_rawDescGZIP(), []int{11}
}
func (x *AttestationResultResponse) GetFile() []byte {
if x != nil {
return x.File
}
return nil
}
var File_agent_agent_proto protoreflect.FileDescriptor
var file_agent_agent_proto_rawDesc = string([]byte{
var file_agent_agent_proto_rawDesc = []byte{
0x0a, 0x11, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x70, 0x72,
0x6f, 0x74, 0x6f, 0x12, 0x05, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x22, 0x4f, 0x0a, 0x0b, 0x41, 0x6c,
0x6f, 0x74, 0x6f, 0x12, 0x05, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x22, 0x57, 0x0a, 0x0b, 0x41, 0x6c,
0x67, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1c, 0x0a, 0x09, 0x61, 0x6c, 0x67,
0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x09, 0x61, 0x6c,
0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x12, 0x22, 0x0a, 0x0c, 0x72, 0x65, 0x71, 0x75, 0x69,
0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0c, 0x72,
0x65, 0x71, 0x75, 0x69, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x22, 0x0e, 0x0a, 0x0c, 0x41,
0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x43, 0x0a, 0x0b, 0x44,
0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x61,
0x74, 0x61, 0x73, 0x65, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x61, 0x74,
0x61, 0x73, 0x65, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65,
0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65,
0x22, 0x0e, 0x0a, 0x0c, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
0x22, 0x0f, 0x0a, 0x0d, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
0x74, 0x22, 0x24, 0x0a, 0x0e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f,
0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28,
0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x22, 0x62, 0x0a, 0x12, 0x41, 0x74, 0x74, 0x65, 0x73,
0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1a, 0x0a,
0x08, 0x74, 0x65, 0x65, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52,
0x08, 0x74, 0x65, 0x65, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x12, 0x1c, 0x0a, 0x09, 0x76, 0x74, 0x70,
0x6d, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x09, 0x76, 0x74,
0x70, 0x6d, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18,
0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x22, 0x29, 0x0a, 0x13, 0x41,
0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c,
0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x22, 0x18, 0x0a, 0x16, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61,
0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74,
0x22, 0x43, 0x0a, 0x17, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65,
0x6e, 0x74, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66,
0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x12,
0x14, 0x0a, 0x05, 0x70, 0x63, 0x72, 0x31, 0x30, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x05,
0x70, 0x63, 0x72, 0x31, 0x30, 0x22, 0x4e, 0x0a, 0x18, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61,
0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
0x74, 0x12, 0x1e, 0x0a, 0x0a, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18,
0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0a, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x4e, 0x6f, 0x6e, 0x63,
0x65, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52,
0x04, 0x74, 0x79, 0x70, 0x65, 0x22, 0x2f, 0x0a, 0x19, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61,
0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c,
0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x32, 0xad, 0x03, 0x0a, 0x0c, 0x41, 0x67, 0x65, 0x6e, 0x74,
0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x33, 0x0a, 0x04, 0x41, 0x6c, 0x67, 0x6f, 0x12,
0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52, 0x65, 0x71, 0x75,
0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f,
0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x28, 0x01, 0x12, 0x33, 0x0a, 0x04,
0x44, 0x61, 0x74, 0x61, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x44, 0x61, 0x74,
0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74,
0x2e, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x28,
0x01, 0x12, 0x39, 0x0a, 0x06, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x12, 0x14, 0x2e, 0x61, 0x67,
0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
0x74, 0x1a, 0x15, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74,
0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x48, 0x0a, 0x0b,
0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x19, 0x2e, 0x61, 0x67,
0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52,
0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41,
0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x54, 0x0a, 0x0f, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61,
0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x12, 0x1d, 0x2e, 0x61, 0x67, 0x65, 0x6e,
0x74, 0x2e, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74,
0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74,
0x2e, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73,
0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x58, 0x0a, 0x11,
0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c,
0x74, 0x12, 0x1f, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74,
0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65,
0x73, 0x74, 0x1a, 0x20, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73,
0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70,
0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61, 0x67, 0x65, 0x6e,
0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
})
0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x12, 0x1a, 0x0a, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69,
0x64, 0x65, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69,
0x64, 0x65, 0x72, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52,
0x02, 0x69, 0x64, 0x22, 0x0e, 0x0a, 0x0c, 0x41, 0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f,
0x6e, 0x73, 0x65, 0x22, 0x53, 0x0a, 0x0b, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65,
0x73, 0x74, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x18, 0x01, 0x20,
0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x12, 0x1a, 0x0a, 0x08,
0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08,
0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x03,
0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x22, 0x0e, 0x0a, 0x0c, 0x44, 0x61, 0x74, 0x61,
0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x2b, 0x0a, 0x0d, 0x52, 0x65, 0x73, 0x75,
0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x63, 0x6f, 0x6e,
0x73, 0x75, 0x6d, 0x65, 0x72, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x63, 0x6f, 0x6e,
0x73, 0x75, 0x6d, 0x65, 0x72, 0x22, 0x24, 0x0a, 0x0e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52,
0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18,
0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x22, 0x35, 0x0a, 0x12, 0x41,
0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
0x74, 0x12, 0x1f, 0x0a, 0x0b, 0x72, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x5f, 0x64, 0x61, 0x74, 0x61,
0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0a, 0x72, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x44, 0x61,
0x74, 0x61, 0x22, 0x29, 0x0a, 0x13, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f,
0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c,
0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x32, 0xf5, 0x01,
0x0a, 0x0c, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x31,
0x0a, 0x04, 0x41, 0x6c, 0x67, 0x6f, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41,
0x6c, 0x67, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65,
0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22,
0x00, 0x12, 0x31, 0x0a, 0x04, 0x44, 0x61, 0x74, 0x61, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e,
0x74, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e,
0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
0x73, 0x65, 0x22, 0x00, 0x12, 0x37, 0x0a, 0x06, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x12, 0x14,
0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71,
0x75, 0x65, 0x73, 0x74, 0x1a, 0x15, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73,
0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x46, 0x0a,
0x0b, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x19, 0x2e, 0x61,
0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e,
0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e,
0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f,
0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61, 0x67, 0x65, 0x6e, 0x74,
0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
}
var (
file_agent_agent_proto_rawDescOnce sync.Once
file_agent_agent_proto_rawDescData []byte
file_agent_agent_proto_rawDescData = file_agent_agent_proto_rawDesc
)
func file_agent_agent_proto_rawDescGZIP() []byte {
file_agent_agent_proto_rawDescOnce.Do(func() {
file_agent_agent_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_agent_proto_rawDesc), len(file_agent_agent_proto_rawDesc)))
file_agent_agent_proto_rawDescData = protoimpl.X.CompressGZIP(file_agent_agent_proto_rawDescData)
})
return file_agent_agent_proto_rawDescData
}
var file_agent_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 12)
var file_agent_agent_proto_goTypes = []any{
(*AlgoRequest)(nil), // 0: agent.AlgoRequest
(*AlgoResponse)(nil), // 1: agent.AlgoResponse
(*DataRequest)(nil), // 2: agent.DataRequest
(*DataResponse)(nil), // 3: agent.DataResponse
(*ResultRequest)(nil), // 4: agent.ResultRequest
(*ResultResponse)(nil), // 5: agent.ResultResponse
(*AttestationRequest)(nil), // 6: agent.AttestationRequest
(*AttestationResponse)(nil), // 7: agent.AttestationResponse
(*IMAMeasurementsRequest)(nil), // 8: agent.IMAMeasurementsRequest
(*IMAMeasurementsResponse)(nil), // 9: agent.IMAMeasurementsResponse
(*AttestationResultRequest)(nil), // 10: agent.AttestationResultRequest
(*AttestationResultResponse)(nil), // 11: agent.AttestationResultResponse
var file_agent_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 8)
var file_agent_agent_proto_goTypes = []interface{}{
(*AlgoRequest)(nil), // 0: agent.AlgoRequest
(*AlgoResponse)(nil), // 1: agent.AlgoResponse
(*DataRequest)(nil), // 2: agent.DataRequest
(*DataResponse)(nil), // 3: agent.DataResponse
(*ResultRequest)(nil), // 4: agent.ResultRequest
(*ResultResponse)(nil), // 5: agent.ResultResponse
(*AttestationRequest)(nil), // 6: agent.AttestationRequest
(*AttestationResponse)(nil), // 7: agent.AttestationResponse
}
var file_agent_agent_proto_depIdxs = []int32{
0, // 0: agent.AgentService.Algo:input_type -> agent.AlgoRequest
2, // 1: agent.AgentService.Data:input_type -> agent.DataRequest
4, // 2: agent.AgentService.Result:input_type -> agent.ResultRequest
6, // 3: agent.AgentService.Attestation:input_type -> agent.AttestationRequest
8, // 4: agent.AgentService.IMAMeasurements:input_type -> agent.IMAMeasurementsRequest
10, // 5: agent.AgentService.AttestationResult:input_type -> agent.AttestationResultRequest
1, // 6: agent.AgentService.Algo:output_type -> agent.AlgoResponse
3, // 7: agent.AgentService.Data:output_type -> agent.DataResponse
5, // 8: agent.AgentService.Result:output_type -> agent.ResultResponse
7, // 9: agent.AgentService.Attestation:output_type -> agent.AttestationResponse
9, // 10: agent.AgentService.IMAMeasurements:output_type -> agent.IMAMeasurementsResponse
11, // 11: agent.AgentService.AttestationResult:output_type -> agent.AttestationResultResponse
6, // [6:12] is the sub-list for method output_type
0, // [0:6] is the sub-list for method input_type
0, // [0:0] is the sub-list for extension type_name
0, // [0:0] is the sub-list for extension extendee
0, // [0:0] is the sub-list for field type_name
0, // 0: agent.AgentService.Algo:input_type -> agent.AlgoRequest
2, // 1: agent.AgentService.Data:input_type -> agent.DataRequest
4, // 2: agent.AgentService.Result:input_type -> agent.ResultRequest
6, // 3: agent.AgentService.Attestation:input_type -> agent.AttestationRequest
1, // 4: agent.AgentService.Algo:output_type -> agent.AlgoResponse
3, // 5: agent.AgentService.Data:output_type -> agent.DataResponse
5, // 6: agent.AgentService.Result:output_type -> agent.ResultResponse
7, // 7: agent.AgentService.Attestation:output_type -> agent.AttestationResponse
4, // [4:8] is the sub-list for method output_type
0, // [0:4] is the sub-list for method input_type
0, // [0:0] is the sub-list for extension type_name
0, // [0:0] is the sub-list for extension extendee
0, // [0:0] is the sub-list for field type_name
}
func init() { file_agent_agent_proto_init() }
@@ -693,13 +505,111 @@ func file_agent_agent_proto_init() {
if File_agent_agent_proto != nil {
return
}
if !protoimpl.UnsafeEnabled {
file_agent_agent_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*AlgoRequest); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_agent_agent_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*AlgoResponse); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_agent_agent_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*DataRequest); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_agent_agent_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*DataResponse); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_agent_agent_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*ResultRequest); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_agent_agent_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*ResultResponse); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_agent_agent_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*AttestationRequest); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_agent_agent_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*AttestationResponse); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
}
type x struct{}
out := protoimpl.TypeBuilder{
File: protoimpl.DescBuilder{
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_agent_proto_rawDesc), len(file_agent_agent_proto_rawDesc)),
RawDescriptor: file_agent_agent_proto_rawDesc,
NumEnums: 0,
NumMessages: 12,
NumMessages: 8,
NumExtensions: 0,
NumServices: 1,
},
@@ -708,6 +618,7 @@ func file_agent_agent_proto_init() {
MessageInfos: file_agent_agent_proto_msgTypes,
}.Build()
File_agent_agent_proto = out.File
file_agent_agent_proto_rawDesc = nil
file_agent_agent_proto_goTypes = nil
file_agent_agent_proto_depIdxs = nil
}
+10 -27
View File
@@ -8,29 +8,30 @@ package agent;
option go_package = "./agent";
service AgentService {
rpc Algo(stream AlgoRequest) returns (AlgoResponse) {}
rpc Data(stream DataRequest) returns (DataResponse) {}
rpc Result(ResultRequest) returns (stream ResultResponse) {}
rpc Attestation(AttestationRequest) returns (stream AttestationResponse) {}
rpc IMAMeasurements(IMAMeasurementsRequest) returns (stream IMAMeasurementsResponse) {}
rpc AttestationResult(AttestationResultRequest) returns (AttestationResultResponse) {}
rpc Algo(AlgoRequest) returns (AlgoResponse) {}
rpc Data(DataRequest) returns (DataResponse) {}
rpc Result(ResultRequest) returns (ResultResponse) {}
rpc Attestation(AttestationRequest) returns (AttestationResponse) {}
}
message AlgoRequest {
bytes algorithm = 1;
bytes requirements = 2;
string provider = 2;
string id = 3;
}
message AlgoResponse {}
message DataRequest {
bytes dataset = 1;
string filename = 2;
string provider = 2;
string id = 3;
}
message DataResponse {}
message ResultRequest {
string consumer = 1;
}
message ResultResponse {
@@ -38,27 +39,9 @@ message ResultResponse {
}
message AttestationRequest {
bytes teeNonce = 1; // Should be less or equal 64 bytes.
bytes vtpmNonce = 2; // Should be less or equal 32 bytes.
int32 type = 3;
bytes report_data = 1; // Should be of length 64.
}
message AttestationResponse {
bytes file = 1;
}
message IMAMeasurementsRequest {
}
message IMAMeasurementsResponse {
bytes file = 1;
bytes pcr10 = 2;
}
message AttestationResultRequest{
bytes tokenNonce = 1; // Should be less or equal 32 bytes
int32 type = 3;
}
message AttestationResultResponse{
bytes file = 1;
}
+117 -204
View File
@@ -3,8 +3,8 @@
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
// versions:
// - protoc-gen-go-grpc v1.5.1
// - protoc v5.29.0
// - protoc-gen-go-grpc v1.3.0
// - protoc v4.25.3
// source: agent/agent.proto
package agent
@@ -18,28 +18,24 @@ import (
// This is a compile-time assertion to ensure that this generated file
// is compatible with the grpc package it is being compiled against.
// Requires gRPC-Go v1.64.0 or later.
const _ = grpc.SupportPackageIsVersion9
// Requires gRPC-Go v1.32.0 or later.
const _ = grpc.SupportPackageIsVersion7
const (
AgentService_Algo_FullMethodName = "/agent.AgentService/Algo"
AgentService_Data_FullMethodName = "/agent.AgentService/Data"
AgentService_Result_FullMethodName = "/agent.AgentService/Result"
AgentService_Attestation_FullMethodName = "/agent.AgentService/Attestation"
AgentService_IMAMeasurements_FullMethodName = "/agent.AgentService/IMAMeasurements"
AgentService_AttestationResult_FullMethodName = "/agent.AgentService/AttestationResult"
AgentService_Algo_FullMethodName = "/agent.AgentService/Algo"
AgentService_Data_FullMethodName = "/agent.AgentService/Data"
AgentService_Result_FullMethodName = "/agent.AgentService/Result"
AgentService_Attestation_FullMethodName = "/agent.AgentService/Attestation"
)
// AgentServiceClient is the client API for AgentService service.
//
// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
type AgentServiceClient interface {
Algo(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[AlgoRequest, AlgoResponse], error)
Data(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[DataRequest, DataResponse], error)
Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ResultResponse], error)
Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AttestationResponse], error)
IMAMeasurements(ctx context.Context, in *IMAMeasurementsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[IMAMeasurementsResponse], error)
AttestationResult(ctx context.Context, in *AttestationResultRequest, opts ...grpc.CallOption) (*AttestationResultResponse, error)
Algo(ctx context.Context, in *AlgoRequest, opts ...grpc.CallOption) (*AlgoResponse, error)
Data(ctx context.Context, in *DataRequest, opts ...grpc.CallOption) (*DataResponse, error)
Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (*ResultResponse, error)
Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error)
}
type agentServiceClient struct {
@@ -50,93 +46,36 @@ func NewAgentServiceClient(cc grpc.ClientConnInterface) AgentServiceClient {
return &agentServiceClient{cc}
}
func (c *agentServiceClient) Algo(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[AlgoRequest, AlgoResponse], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[0], AgentService_Algo_FullMethodName, cOpts...)
func (c *agentServiceClient) Algo(ctx context.Context, in *AlgoRequest, opts ...grpc.CallOption) (*AlgoResponse, error) {
out := new(AlgoResponse)
err := c.cc.Invoke(ctx, AgentService_Algo_FullMethodName, in, out, opts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[AlgoRequest, AlgoResponse]{ClientStream: stream}
return x, nil
return out, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_AlgoClient = grpc.ClientStreamingClient[AlgoRequest, AlgoResponse]
func (c *agentServiceClient) Data(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[DataRequest, DataResponse], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[1], AgentService_Data_FullMethodName, cOpts...)
func (c *agentServiceClient) Data(ctx context.Context, in *DataRequest, opts ...grpc.CallOption) (*DataResponse, error) {
out := new(DataResponse)
err := c.cc.Invoke(ctx, AgentService_Data_FullMethodName, in, out, opts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[DataRequest, DataResponse]{ClientStream: stream}
return x, nil
return out, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_DataClient = grpc.ClientStreamingClient[DataRequest, DataResponse]
func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ResultResponse], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[2], AgentService_Result_FullMethodName, cOpts...)
func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (*ResultResponse, error) {
out := new(ResultResponse)
err := c.cc.Invoke(ctx, AgentService_Result_FullMethodName, in, out, opts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[ResultRequest, ResultResponse]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
return out, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_ResultClient = grpc.ServerStreamingClient[ResultResponse]
func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AttestationResponse], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[3], AgentService_Attestation_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[AttestationRequest, AttestationResponse]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_AttestationClient = grpc.ServerStreamingClient[AttestationResponse]
func (c *agentServiceClient) IMAMeasurements(ctx context.Context, in *IMAMeasurementsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[IMAMeasurementsResponse], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[4], AgentService_IMAMeasurements_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[IMAMeasurementsRequest, IMAMeasurementsResponse]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_IMAMeasurementsClient = grpc.ServerStreamingClient[IMAMeasurementsResponse]
func (c *agentServiceClient) AttestationResult(ctx context.Context, in *AttestationResultRequest, opts ...grpc.CallOption) (*AttestationResultResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(AttestationResultResponse)
err := c.cc.Invoke(ctx, AgentService_AttestationResult_FullMethodName, in, out, cOpts...)
func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error) {
out := new(AttestationResponse)
err := c.cc.Invoke(ctx, AgentService_Attestation_FullMethodName, in, out, opts...)
if err != nil {
return nil, err
}
@@ -145,44 +84,32 @@ func (c *agentServiceClient) AttestationResult(ctx context.Context, in *Attestat
// AgentServiceServer is the server API for AgentService service.
// All implementations must embed UnimplementedAgentServiceServer
// for forward compatibility.
// for forward compatibility
type AgentServiceServer interface {
Algo(grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]) error
Data(grpc.ClientStreamingServer[DataRequest, DataResponse]) error
Result(*ResultRequest, grpc.ServerStreamingServer[ResultResponse]) error
Attestation(*AttestationRequest, grpc.ServerStreamingServer[AttestationResponse]) error
IMAMeasurements(*IMAMeasurementsRequest, grpc.ServerStreamingServer[IMAMeasurementsResponse]) error
AttestationResult(context.Context, *AttestationResultRequest) (*AttestationResultResponse, error)
Algo(context.Context, *AlgoRequest) (*AlgoResponse, error)
Data(context.Context, *DataRequest) (*DataResponse, error)
Result(context.Context, *ResultRequest) (*ResultResponse, error)
Attestation(context.Context, *AttestationRequest) (*AttestationResponse, error)
mustEmbedUnimplementedAgentServiceServer()
}
// UnimplementedAgentServiceServer must be embedded to have
// forward compatible implementations.
//
// NOTE: this should be embedded by value instead of pointer to avoid a nil
// pointer dereference when methods are called.
type UnimplementedAgentServiceServer struct{}
// UnimplementedAgentServiceServer must be embedded to have forward compatible implementations.
type UnimplementedAgentServiceServer struct {
}
func (UnimplementedAgentServiceServer) Algo(grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]) error {
return status.Errorf(codes.Unimplemented, "method Algo not implemented")
func (UnimplementedAgentServiceServer) Algo(context.Context, *AlgoRequest) (*AlgoResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method Algo not implemented")
}
func (UnimplementedAgentServiceServer) Data(grpc.ClientStreamingServer[DataRequest, DataResponse]) error {
return status.Errorf(codes.Unimplemented, "method Data not implemented")
func (UnimplementedAgentServiceServer) Data(context.Context, *DataRequest) (*DataResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method Data not implemented")
}
func (UnimplementedAgentServiceServer) Result(*ResultRequest, grpc.ServerStreamingServer[ResultResponse]) error {
return status.Errorf(codes.Unimplemented, "method Result not implemented")
func (UnimplementedAgentServiceServer) Result(context.Context, *ResultRequest) (*ResultResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method Result not implemented")
}
func (UnimplementedAgentServiceServer) Attestation(*AttestationRequest, grpc.ServerStreamingServer[AttestationResponse]) error {
return status.Errorf(codes.Unimplemented, "method Attestation not implemented")
}
func (UnimplementedAgentServiceServer) IMAMeasurements(*IMAMeasurementsRequest, grpc.ServerStreamingServer[IMAMeasurementsResponse]) error {
return status.Errorf(codes.Unimplemented, "method IMAMeasurements not implemented")
}
func (UnimplementedAgentServiceServer) AttestationResult(context.Context, *AttestationResultRequest) (*AttestationResultResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method AttestationResult not implemented")
func (UnimplementedAgentServiceServer) Attestation(context.Context, *AttestationRequest) (*AttestationResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method Attestation not implemented")
}
func (UnimplementedAgentServiceServer) mustEmbedUnimplementedAgentServiceServer() {}
func (UnimplementedAgentServiceServer) testEmbeddedByValue() {}
// UnsafeAgentServiceServer may be embedded to opt out of forward compatibility for this service.
// Use of this interface is not recommended, as added methods to AgentServiceServer will
@@ -192,77 +119,77 @@ type UnsafeAgentServiceServer interface {
}
func RegisterAgentServiceServer(s grpc.ServiceRegistrar, srv AgentServiceServer) {
// If the following call pancis, it indicates UnimplementedAgentServiceServer was
// embedded by pointer and is nil. This will cause panics if an
// unimplemented method is ever invoked, so we test this at initialization
// time to prevent it from happening at runtime later due to I/O.
if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
t.testEmbeddedByValue()
}
s.RegisterService(&AgentService_ServiceDesc, srv)
}
func _AgentService_Algo_Handler(srv interface{}, stream grpc.ServerStream) error {
return srv.(AgentServiceServer).Algo(&grpc.GenericServerStream[AlgoRequest, AlgoResponse]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_AlgoServer = grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]
func _AgentService_Data_Handler(srv interface{}, stream grpc.ServerStream) error {
return srv.(AgentServiceServer).Data(&grpc.GenericServerStream[DataRequest, DataResponse]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_DataServer = grpc.ClientStreamingServer[DataRequest, DataResponse]
func _AgentService_Result_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(ResultRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(AgentServiceServer).Result(m, &grpc.GenericServerStream[ResultRequest, ResultResponse]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_ResultServer = grpc.ServerStreamingServer[ResultResponse]
func _AgentService_Attestation_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(AttestationRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(AgentServiceServer).Attestation(m, &grpc.GenericServerStream[AttestationRequest, AttestationResponse]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_AttestationServer = grpc.ServerStreamingServer[AttestationResponse]
func _AgentService_IMAMeasurements_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(IMAMeasurementsRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(AgentServiceServer).IMAMeasurements(m, &grpc.GenericServerStream[IMAMeasurementsRequest, IMAMeasurementsResponse]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type AgentService_IMAMeasurementsServer = grpc.ServerStreamingServer[IMAMeasurementsResponse]
func _AgentService_AttestationResult_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(AttestationResultRequest)
func _AgentService_Algo_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(AlgoRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(AgentServiceServer).AttestationResult(ctx, in)
return srv.(AgentServiceServer).Algo(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: AgentService_AttestationResult_FullMethodName,
FullMethod: AgentService_Algo_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(AgentServiceServer).AttestationResult(ctx, req.(*AttestationResultRequest))
return srv.(AgentServiceServer).Algo(ctx, req.(*AlgoRequest))
}
return interceptor(ctx, in, info, handler)
}
func _AgentService_Data_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(DataRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(AgentServiceServer).Data(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: AgentService_Data_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(AgentServiceServer).Data(ctx, req.(*DataRequest))
}
return interceptor(ctx, in, info, handler)
}
func _AgentService_Result_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(ResultRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(AgentServiceServer).Result(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: AgentService_Result_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(AgentServiceServer).Result(ctx, req.(*ResultRequest))
}
return interceptor(ctx, in, info, handler)
}
func _AgentService_Attestation_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(AttestationRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(AgentServiceServer).Attestation(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: AgentService_Attestation_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(AgentServiceServer).Attestation(ctx, req.(*AttestationRequest))
}
return interceptor(ctx, in, info, handler)
}
@@ -275,36 +202,22 @@ var AgentService_ServiceDesc = grpc.ServiceDesc{
HandlerType: (*AgentServiceServer)(nil),
Methods: []grpc.MethodDesc{
{
MethodName: "AttestationResult",
Handler: _AgentService_AttestationResult_Handler,
},
},
Streams: []grpc.StreamDesc{
{
StreamName: "Algo",
Handler: _AgentService_Algo_Handler,
ClientStreams: true,
},
{
StreamName: "Data",
Handler: _AgentService_Data_Handler,
ClientStreams: true,
},
{
StreamName: "Result",
Handler: _AgentService_Result_Handler,
ServerStreams: true,
},
{
StreamName: "Attestation",
Handler: _AgentService_Attestation_Handler,
ServerStreams: true,
},
{
StreamName: "IMAMeasurements",
Handler: _AgentService_IMAMeasurements_Handler,
ServerStreams: true,
MethodName: "Algo",
Handler: _AgentService_Algo_Handler,
},
{
MethodName: "Data",
Handler: _AgentService_Data_Handler,
},
{
MethodName: "Result",
Handler: _AgentService_Result_Handler,
},
{
MethodName: "Attestation",
Handler: _AgentService_Attestation_Handler,
},
},
Streams: []grpc.StreamDesc{},
Metadata: "agent/agent.proto",
}
-29
View File
@@ -1,29 +0,0 @@
// Code generated by "stringer -type=AgentEvent"; DO NOT EDIT.
package agent
import "strconv"
func _() {
// An "invalid array index" compiler error signifies that the constant values have changed.
// Re-run the stringer command to generate them again.
var x [1]struct{}
_ = x[Start-0]
_ = x[ManifestReceived-1]
_ = x[AlgorithmReceived-2]
_ = x[DataReceived-3]
_ = x[RunComplete-4]
_ = x[ResultsConsumed-5]
_ = x[RunFailed-6]
}
const _AgentEvent_name = "StartManifestReceivedAlgorithmReceivedDataReceivedRunCompleteResultsConsumedRunFailed"
var _AgentEvent_index = [...]uint8{0, 5, 21, 38, 50, 61, 76, 85}
func (i AgentEvent) String() string {
if i < 0 || i >= AgentEvent(len(_AgentEvent_index)-1) {
return "AgentEvent(" + strconv.FormatInt(int64(i), 10) + ")"
}
return _AgentEvent_name[_AgentEvent_index[i]:_AgentEvent_index[i+1]]
}
-30
View File
@@ -1,30 +0,0 @@
// Code generated by "stringer -type=AgentState"; DO NOT EDIT.
package agent
import "strconv"
func _() {
// An "invalid array index" compiler error signifies that the constant values have changed.
// Re-run the stringer command to generate them again.
var x [1]struct{}
_ = x[Idle-0]
_ = x[ReceivingManifest-1]
_ = x[ReceivingAlgorithm-2]
_ = x[ReceivingData-3]
_ = x[Running-4]
_ = x[ConsumingResults-5]
_ = x[Complete-6]
_ = x[Failed-7]
}
const _AgentState_name = "IdleReceivingManifestReceivingAlgorithmReceivingDataRunningConsumingResultsCompleteFailed"
var _AgentState_index = [...]uint8{0, 4, 21, 39, 52, 59, 75, 83, 89}
func (i AgentState) String() string {
if i < 0 || i >= AgentState(len(_AgentState_index)-1) {
return "AgentState(" + strconv.FormatInt(int64(i), 10) + ")"
}
return _AgentState_name[_AgentState_index[i]:_AgentState_index[i+1]]
}
-52
View File
@@ -1,52 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package algorithm
import (
"context"
"google.golang.org/grpc/metadata"
)
type AlgorithType string
const (
AlgoTypeBin AlgorithType = "bin"
AlgoTypePython AlgorithType = "python"
AlgoTypeWasm AlgorithType = "wasm"
AlgoTypeDocker AlgorithType = "docker"
AlgoTypeKey = "algo_type"
AlgoArgsKey = "algo_args"
ResultsDir = "results"
DatasetsDir = "datasets"
AlgoWorkingDir = "/cocos"
)
func AlgorithmTypeToContext(ctx context.Context, algoType string) context.Context {
return metadata.AppendToOutgoingContext(ctx, AlgoTypeKey, algoType)
}
func AlgorithmTypeFromContext(ctx context.Context) string {
return metadata.ValueFromIncomingContext(ctx, AlgoTypeKey)[0]
}
func AlgorithmArgsToContext(ctx context.Context, algoArgs []string) context.Context {
for _, arg := range algoArgs {
ctx = metadata.AppendToOutgoingContext(ctx, AlgoArgsKey, arg)
}
return ctx
}
func AlgorithmArgsFromContext(ctx context.Context) []string {
return metadata.ValueFromIncomingContext(ctx, AlgoArgsKey)
}
// Algorithm is an interface that specifies the API for an algorithm.
type Algorithm interface {
// Run executes the algorithm and returns the result.
Run() error
// Stop stops the algorithm.
Stop() error
}
-69
View File
@@ -1,69 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package binary
import (
"fmt"
"io"
"log/slog"
"os/exec"
"github.com/ultravioletrs/cocos/agent/algorithm"
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
"github.com/ultravioletrs/cocos/agent/events"
)
var _ algorithm.Algorithm = (*binary)(nil)
type binary struct {
algoFile string
stderr io.Writer
stdout io.Writer
args []string
cmd *exec.Cmd
}
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string, cmpID string) algorithm.Algorithm {
return &binary{
algoFile: algoFile,
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
stdout: &logging.Stdout{Logger: logger},
args: args,
}
}
func (b *binary) Run() error {
b.cmd = exec.Command(b.algoFile, b.args...)
b.cmd.Stderr = b.stderr
b.cmd.Stdout = b.stdout
if err := b.cmd.Start(); err != nil {
return fmt.Errorf("error starting algorithm: %v", err)
}
if err := b.cmd.Wait(); err != nil {
return fmt.Errorf("algorithm execution error: %v", err)
}
return nil
}
func (b *binary) Stop() error {
if b.cmd == nil {
return nil
}
if b.cmd.ProcessState != nil && b.cmd.ProcessState.Exited() {
return nil
}
if b.cmd.Process == nil {
return nil
}
if err := b.cmd.Process.Kill(); err != nil {
return fmt.Errorf("error stopping algorithm: %v", err)
}
return nil
}
-100
View File
@@ -1,100 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package binary
import (
"bytes"
"log/slog"
"os"
"testing"
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
"github.com/ultravioletrs/cocos/agent/events/mocks"
)
func TestNewAlgorithm(t *testing.T) {
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
eventsSvc := new(mocks.Service)
algoFile := "/path/to/algo"
args := []string{"arg1", "arg2"}
algo := NewAlgorithm(logger, eventsSvc, algoFile, args, "")
b, ok := algo.(*binary)
if !ok {
t.Fatalf("NewAlgorithm did not return a *binary")
}
if b.algoFile != algoFile {
t.Errorf("Expected algoFile to be %s, got %s", algoFile, b.algoFile)
}
if len(b.args) != len(args) {
t.Errorf("Expected %d args, got %d", len(args), len(b.args))
}
for i, arg := range args {
if b.args[i] != arg {
t.Errorf("Expected arg %d to be %s, got %s", i, arg, b.args[i])
}
}
if _, ok := b.stderr.(*logging.Stderr); !ok {
t.Errorf("Expected stderr to be *algorithm.Stderr")
}
if _, ok := b.stdout.(*logging.Stdout); !ok {
t.Errorf("Expected stdout to be *algorithm.Stdout")
}
}
func TestBinaryRun(t *testing.T) {
tests := []struct {
name string
algoFile string
args []string
expectedError bool
}{
{
name: "Successful execution",
algoFile: "echo",
args: []string{"Hello, World!"},
expectedError: false,
},
{
name: "Non-existent binary",
algoFile: "non_existent_binary",
args: []string{},
expectedError: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
eventsSvc := new(mocks.Service)
b := NewAlgorithm(logger, eventsSvc, tt.algoFile, tt.args, "").(*binary)
var stdout, stderr bytes.Buffer
b.stdout = &stdout
b.stderr = &stderr
err := b.Run()
if tt.expectedError && err == nil {
t.Errorf("Expected an error, but got none")
}
if !tt.expectedError && err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !tt.expectedError {
if stdout.Len() == 0 {
t.Errorf("Expected non-empty stdout")
}
}
})
}
}
-182
View File
@@ -1,182 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package docker
import (
"bufio"
"context"
"fmt"
"io"
"log/slog"
"os"
"path"
"github.com/docker/docker/api/types/container"
"github.com/docker/docker/api/types/image"
"github.com/docker/docker/api/types/mount"
"github.com/docker/docker/client"
"github.com/ultravioletrs/cocos/agent/algorithm"
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
"github.com/ultravioletrs/cocos/agent/events"
)
const (
containerName = "agent_container"
datasetsMountPath = "/cocos/datasets"
resultsMountPath = "/cocos/results"
)
var _ algorithm.Algorithm = (*docker)(nil)
type docker struct {
algoFile string
logger *slog.Logger
stderr io.Writer
stdout io.Writer
}
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile, cmpID string) algorithm.Algorithm {
d := &docker{
algoFile: algoFile,
logger: logger,
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
stdout: &logging.Stdout{Logger: logger},
}
return d
}
func (d *docker) Run() error {
// Create a new Docker client.
cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
if err != nil {
return fmt.Errorf("could not create a new Docker client: %v", err)
}
// Open the Docker image tar file.
imageFile, err := os.Open(d.algoFile)
if err != nil {
return fmt.Errorf("could not open Docker image: %v", err)
}
defer imageFile.Close()
ctx := context.Background()
// Load the Docker image from the tar file.
resp, err := cli.ImageLoad(ctx, imageFile, client.ImageLoadWithQuiet(true))
if err != nil {
return fmt.Errorf("could not load Docker image from file: %v", err)
}
defer resp.Body.Close()
// List the loaded images to get the image ID.
images, err := cli.ImageList(ctx, image.ListOptions{})
if err != nil {
return fmt.Errorf("could not get the Docker image list: %v", err)
}
var imageID string = ""
var dockerImageName string = ""
for _, image := range images {
for _, tag := range image.RepoTags {
imageID = image.ID
dockerImageName = tag
break
}
}
if imageID == "" {
return fmt.Errorf("could not find image ID")
}
// Create and start the container.
respContainer, err := cli.ContainerCreate(ctx, &container.Config{
Image: dockerImageName,
Tty: true,
AttachStdout: true,
AttachStderr: true,
}, &container.HostConfig{
Mounts: []mount.Mount{
{
Type: mount.TypeBind,
Source: path.Join(algorithm.AlgoWorkingDir, algorithm.DatasetsDir),
Target: datasetsMountPath,
},
{
Type: mount.TypeBind,
Source: path.Join(algorithm.AlgoWorkingDir, algorithm.ResultsDir),
Target: resultsMountPath,
},
},
}, nil, nil, containerName)
if err != nil {
return fmt.Errorf("could not create a Docker container: %v", err)
}
if err := cli.ContainerStart(ctx, respContainer.ID, container.StartOptions{}); err != nil {
return fmt.Errorf("could not start a Docker container: %v", err)
}
stdout, err := cli.ContainerLogs(ctx, respContainer.ID, container.LogsOptions{ShowStdout: true, Follow: true})
if err != nil {
return fmt.Errorf("could not read stdout from the container: %v", err)
}
defer stdout.Close()
go func() {
if err := writeToOut(stdout, d.stdout); err != nil {
d.logger.Warn(fmt.Sprintf("could not write to stdout: %v", err))
}
}()
stderr, err := cli.ContainerLogs(ctx, respContainer.ID, container.LogsOptions{ShowStderr: true, Follow: true})
if err != nil {
d.logger.Warn(fmt.Sprintf("could not read stderr from the container: %v", err))
}
defer stderr.Close()
go func() {
if err := writeToOut(stderr, d.stderr); err != nil {
d.logger.Warn(fmt.Sprintf("could not write to stderr: %v", err))
}
}()
statusCh, errCh := cli.ContainerWait(ctx, respContainer.ID, container.WaitConditionNotRunning)
select {
case err := <-errCh:
if err != nil {
return fmt.Errorf("could not wait for a Docker container: %v", err)
}
case <-statusCh:
}
defer func() {
if err = cli.ContainerRemove(ctx, respContainer.ID, container.RemoveOptions{Force: true}); err != nil {
d.logger.Warn(fmt.Sprintf("error could not remove container: %v", err))
}
if _, err := cli.ImageRemove(ctx, imageID, image.RemoveOptions{Force: true}); err != nil {
d.logger.Warn(fmt.Sprintf("error could not remove image: %v", err))
}
}()
return nil
}
func writeToOut(readCloser io.ReadCloser, ioWriter io.Writer) error {
scanner := bufio.NewScanner(readCloser)
for scanner.Scan() {
if _, err := ioWriter.Write(scanner.Bytes()); err != nil {
return fmt.Errorf("error writing to output: %v", err)
}
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("error reading container logs error: %v", err)
}
return nil
}
func (d *docker) Stop() error {
// To be supported later.
return nil
}
-29
View File
@@ -1,29 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package docker
import (
"log/slog"
"os"
"testing"
"github.com/stretchr/testify/assert"
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
"github.com/ultravioletrs/cocos/agent/events/mocks"
)
// TestNewAlgorithm tests the NewAlgorithm function.
func TestNewAlgorithm(t *testing.T) {
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
eventsSvc := new(mocks.Service)
algoFile := "/path/to/algo.tar"
algo := NewAlgorithm(logger, eventsSvc, algoFile, "")
d, ok := algo.(*docker)
assert.True(t, ok, "NewAlgorithm should return a *docker")
assert.Equal(t, algoFile, d.algoFile, "algoFile should be set correctly")
assert.NotNil(t, d.logger, "logger should be set")
assert.IsType(t, &logging.Stderr{}, d.stderr, "stderr should be of type *algorithm.Stderr")
assert.IsType(t, &logging.Stdout{}, d.stdout, "stdout should be of type *algorithm.Stdout")
}
-77
View File
@@ -1,77 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package logging
import (
"bytes"
"encoding/json"
"io"
"log/slog"
"github.com/ultravioletrs/cocos/agent/events"
)
var (
_ io.Writer = &Stdout{}
_ io.Writer = &Stderr{}
)
const (
bufSize = 1024
algorithmRun = "AlgorithmRun"
warningStatus = "Warning"
)
type Stdout struct {
Logger *slog.Logger
}
// Write implements io.Writer.
func (s *Stdout) Write(p []byte) (n int, err error) {
inBuf := bytes.NewBuffer(p)
buf := make([]byte, bufSize)
for {
n, err := inBuf.Read(buf)
if err != nil {
if err == io.EOF {
break
}
return len(p) - inBuf.Len(), err
}
s.Logger.Debug(string(buf[:n]))
}
return len(p), nil
}
type Stderr struct {
Logger *slog.Logger
EventSvc events.Service
CmpID string
}
// Write implements io.Writer.
func (s *Stderr) Write(p []byte) (n int, err error) {
inBuf := bytes.NewBuffer(p)
buf := make([]byte, bufSize)
for {
n, err := inBuf.Read(buf)
if err != nil {
if err == io.EOF {
break
}
return len(p) - inBuf.Len(), err
}
s.Logger.Error(string(buf[:n]))
}
s.EventSvc.SendEvent(s.CmpID, algorithmRun, warningStatus, json.RawMessage{})
return len(p), nil
}
-86
View File
@@ -1,86 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package logging
import (
"strings"
"testing"
mglog "github.com/absmach/magistrala/logger"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/ultravioletrs/cocos/agent/events/mocks"
"github.com/ultravioletrs/cocos/pkg/manager"
)
func TestStdoutWrite(t *testing.T) {
tests := []struct {
name string
input string
expected []string
}{
{
name: "Single line",
input: "Hello, World!",
expected: []string{"Hello, World!"},
},
{
name: "Multiple lines",
input: "Line 1\nLine 2\nLine 3",
expected: []string{"Line 1\nLine 2\nLine 3"},
},
{
name: "Long input",
input: strings.Repeat("a", bufSize+100),
expected: []string{strings.Repeat("a", bufSize), strings.Repeat("a", 100)},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
stdout := &Stdout{Logger: mglog.NewMock()}
n, err := stdout.Write([]byte(tt.input))
assert.NoError(t, err)
assert.Equal(t, len(tt.input), n)
})
}
}
func TestStderrWrite(t *testing.T) {
tests := []struct {
name string
input string
expected []string
}{
{
name: "Single line",
input: "Error: Something went wrong",
expected: []string{"Error: Something went wrong"},
},
{
name: "Multiple lines",
input: "Error 1\nError 2\nError 3",
expected: []string{"Error 1\nError 2\nError 3"},
},
{
name: "Long input",
input: strings.Repeat("e", bufSize+100),
expected: []string{strings.Repeat("e", bufSize), strings.Repeat("e", 100)},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mockEventService := mocks.NewService(t)
mockEventService.On("SendEvent", mock.Anything, "AlgorithmRun", manager.Warning.String(), mock.Anything).Return(nil)
stderr := &Stderr{Logger: mglog.NewMock(), EventSvc: mockEventService}
n, err := stderr.Write([]byte(tt.input))
assert.NoError(t, err)
assert.Equal(t, len(tt.input), n)
mockEventService.AssertExpectations(t)
})
}
}
-126
View File
@@ -1,126 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package python
import (
"context"
"fmt"
"io"
"log/slog"
"os"
"os/exec"
"path/filepath"
"github.com/ultravioletrs/cocos/agent/algorithm"
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
"github.com/ultravioletrs/cocos/agent/events"
"google.golang.org/grpc/metadata"
)
const (
PyRuntime = "python3"
PyRuntimeKey = "python_runtime"
)
func PythonRunTimeToContext(ctx context.Context, runtime string) context.Context {
return metadata.AppendToOutgoingContext(ctx, PyRuntimeKey, runtime)
}
func PythonRunTimeFromContext(ctx context.Context) string {
return metadata.ValueFromIncomingContext(ctx, PyRuntimeKey)[0]
}
var _ algorithm.Algorithm = (*python)(nil)
type python struct {
algoFile string
stderr io.Writer
stdout io.Writer
runtime string
requirementsFile string
args []string
cmd *exec.Cmd
}
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, runtime, requirementsFile, algoFile string, args []string, cmpID string) algorithm.Algorithm {
p := &python{
algoFile: algoFile,
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
stdout: &logging.Stdout{Logger: logger},
requirementsFile: requirementsFile,
args: args,
}
if runtime != "" {
p.runtime = runtime
} else {
p.runtime = PyRuntime
}
return p
}
func (p *python) Run() error {
venvPath := "venv"
createVenvCmd := exec.Command(p.runtime, "-m", "venv", venvPath)
createVenvCmd.Stderr = p.stderr
createVenvCmd.Stdout = p.stdout
if err := createVenvCmd.Run(); err != nil {
return fmt.Errorf("error creating virtual environment: %v", err)
}
pythonPath := filepath.Join(venvPath, "bin", "python")
updatePipCmd := exec.Command(pythonPath, "-m", "pip", "install", "--upgrade", "pip")
updatePipCmd.Stderr = p.stderr
updatePipCmd.Stdout = p.stdout
if err := updatePipCmd.Run(); err != nil {
return fmt.Errorf("error updating pip: %v", err)
}
if p.requirementsFile != "" {
rcmd := exec.Command(pythonPath, "-m", "pip", "install", "-r", p.requirementsFile)
rcmd.Stderr = p.stderr
rcmd.Stdout = p.stdout
if err := rcmd.Run(); err != nil {
return fmt.Errorf("error installing requirements: %v", err)
}
}
args := append([]string{p.algoFile}, p.args...)
p.cmd = exec.Command(pythonPath, args...)
p.cmd.Stderr = p.stderr
p.cmd.Stdout = p.stdout
if err := p.cmd.Start(); err != nil {
return fmt.Errorf("error starting algorithm: %v", err)
}
if err := p.cmd.Wait(); err != nil {
return fmt.Errorf("algorithm execution error: %v", err)
}
if err := os.RemoveAll(venvPath); err != nil {
return fmt.Errorf("error removing virtual environment: %v", err)
}
return nil
}
func (p *python) Stop() error {
if p.cmd == nil {
return nil
}
if p.cmd.ProcessState != nil && p.cmd.ProcessState.Exited() {
return nil
}
if p.cmd.Process == nil {
return nil
}
if err := p.cmd.Process.Kill(); err != nil {
return fmt.Errorf("error stopping algorithm: %v", err)
}
return nil
}
-148
View File
@@ -1,148 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package python
import (
"bytes"
"context"
"io"
"log/slog"
"os"
"path/filepath"
"strings"
"testing"
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
"github.com/ultravioletrs/cocos/agent/events/mocks"
"google.golang.org/grpc/metadata"
)
const runtime = "python3"
func TestPythonRunTimeToContext(t *testing.T) {
ctx := context.Background()
newCtx := PythonRunTimeToContext(ctx, runtime)
md, ok := metadata.FromOutgoingContext(newCtx)
if !ok {
t.Fatal("Expected metadata in context")
}
values := md.Get(PyRuntimeKey)
if len(values) != 1 || values[0] != runtime {
t.Errorf("Expected runtime %s, got %v", runtime, values)
}
}
func TestPythonRunTimeFromContext(t *testing.T) {
ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(PyRuntimeKey, runtime))
got := PythonRunTimeFromContext(ctx)
if got != runtime {
t.Errorf("Expected runtime %s, got %s", runtime, got)
}
}
func TestNewAlgorithm(t *testing.T) {
logger := &slog.Logger{}
eventsSvc := new(mocks.Service)
requirementsFile := "requirements.txt"
algoFile := "algorithm.py"
args := []string{"--arg1", "value1"}
algo := NewAlgorithm(logger, eventsSvc, runtime, requirementsFile, algoFile, args, "")
p, ok := algo.(*python)
if !ok {
t.Fatal("Expected *python type")
}
if p.runtime != runtime {
t.Errorf("Expected runtime %s, got %s", runtime, p.runtime)
}
if p.requirementsFile != requirementsFile {
t.Errorf("Expected requirementsFile %s, got %s", requirementsFile, p.requirementsFile)
}
if p.algoFile != algoFile {
t.Errorf("Expected algoFile %s, got %s", algoFile, p.algoFile)
}
if len(p.args) != len(args) {
t.Errorf("Expected %d args, got %d", len(args), len(p.args))
}
}
func TestRun(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "python-test")
if err != nil {
t.Fatal(err)
}
defer os.RemoveAll(tmpDir)
scriptContent := []byte("print('Hello, World!')")
scriptPath := filepath.Join(tmpDir, "test_script.py")
if err := os.WriteFile(scriptPath, scriptContent, 0o644); err != nil {
t.Fatal(err)
}
eventsSvc := new(mocks.Service)
var stdout, stderr bytes.Buffer
algo := &python{
algoFile: scriptPath,
stderr: io.MultiWriter(&stderr, &logging.Stderr{Logger: slog.Default(), EventSvc: eventsSvc}),
stdout: io.MultiWriter(&stdout, &logging.Stdout{Logger: slog.Default()}),
runtime: "python3",
}
err = algo.Run()
if err != nil {
t.Fatalf("Unexpected error: %v", err)
}
expectedOutput := "Hello, World!\n"
if !strings.Contains(stdout.String(), expectedOutput) {
t.Errorf("Expected output to contain %q, got %q", expectedOutput, stdout.String())
}
}
func TestRunWithRequirements(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "python-test")
if err != nil {
t.Fatal(err)
}
defer os.RemoveAll(tmpDir)
scriptContent := []byte("import requests\nprint(requests.__version__)")
scriptPath := filepath.Join(tmpDir, "test_script.py")
if err := os.WriteFile(scriptPath, scriptContent, 0o644); err != nil {
t.Fatal(err)
}
requirementsContent := []byte("requests==2.26.0")
requirementsPath := filepath.Join(tmpDir, "requirements.txt")
if err := os.WriteFile(requirementsPath, requirementsContent, 0o644); err != nil {
t.Fatal(err)
}
eventsSvc := new(mocks.Service)
var stdout, stderr bytes.Buffer
algo := &python{
algoFile: scriptPath,
requirementsFile: requirementsPath,
stderr: io.MultiWriter(&stderr, &logging.Stderr{Logger: slog.Default(), EventSvc: eventsSvc}),
stdout: io.MultiWriter(&stdout, &logging.Stdout{Logger: slog.Default()}),
runtime: "python3",
}
err = algo.Run()
if err != nil {
t.Fatalf("Unexpected error: %v", err)
}
if !strings.Contains(stdout.String(), "2.26.0") {
t.Errorf("Expected output to contain requests version 2.26.0, got %q", stdout.String())
}
}
-82
View File
@@ -1,82 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package algorithm_test
import (
"os"
"testing"
"github.com/ultravioletrs/cocos/agent/algorithm"
"github.com/ultravioletrs/cocos/internal"
)
func TestZipDirectory(t *testing.T) {
cases := []struct {
name string
directories []string
files []string
expected []string
}{
{
name: "empty directory",
directories: []string{"testdata"},
},
{
name: "single file",
files: []string{"file1.txt"},
},
{
name: "directory with single file",
directories: []string{"testdata"},
expected: []string{"testdata/file1.txt"},
},
{
name: "directory with multiple files",
directories: []string{"testdata"},
expected: []string{
"testdata/file1.txt",
"testdata/file2.txt",
"testdata/file3.txt",
},
},
{
name: "nested directories",
directories: []string{"testdata", "testdata/nested"},
expected: []string{
"testdata/nested/file1.txt",
"testdata/nested/file2.txt",
"testdata/nested/file3.txt",
},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if err := os.Mkdir(algorithm.ResultsDir, 0o755); err != nil {
t.Fatalf("error creating results directory: %s", err.Error())
}
defer func() {
if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
t.Fatalf("error removing results directory and its contents: %s", err.Error())
}
}()
for _, dir := range tc.directories {
if dir != "" {
if err := os.Mkdir(algorithm.ResultsDir+"/"+dir, 0o755); err != nil {
t.Fatalf("error creating test directory: %s", err.Error())
}
}
}
for _, file := range tc.files {
if _, err := os.Create(algorithm.ResultsDir + "/" + file); err != nil {
t.Fatalf("error creating test file: %s", err.Error())
}
}
if _, err := internal.ZipDirectoryToMemory(algorithm.ResultsDir); err != nil {
t.Errorf("ZipDirectory() error = %v", err)
}
})
}
}
-75
View File
@@ -1,75 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package wasm
import (
"fmt"
"io"
"log/slog"
"os/exec"
"github.com/ultravioletrs/cocos/agent/algorithm"
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
"github.com/ultravioletrs/cocos/agent/events"
)
const wasmRuntime = "wasmedge"
var mapDirOption = []string{"--dir", ".:" + algorithm.ResultsDir}
var _ algorithm.Algorithm = (*wasm)(nil)
type wasm struct {
algoFile string
stderr io.Writer
stdout io.Writer
args []string
cmd *exec.Cmd
}
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, args []string, algoFile, cmpID string) algorithm.Algorithm {
return &wasm{
algoFile: algoFile,
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
stdout: &logging.Stdout{Logger: logger},
args: args,
}
}
func (w *wasm) Run() error {
args := append(mapDirOption, w.algoFile)
args = append(args, w.args...)
w.cmd = exec.Command(wasmRuntime, args...)
w.cmd.Stderr = w.stderr
w.cmd.Stdout = w.stdout
if err := w.cmd.Start(); err != nil {
return fmt.Errorf("error starting algorithm: %v", err)
}
if err := w.cmd.Wait(); err != nil {
return fmt.Errorf("algorithm execution error: %v", err)
}
return nil
}
func (w *wasm) Stop() error {
if w.cmd == nil {
return nil
}
if w.cmd.ProcessState != nil && w.cmd.ProcessState.Exited() {
return nil
}
if w.cmd.Process == nil {
return nil
}
if err := w.cmd.Process.Kill(); err != nil {
return fmt.Errorf("error stopping algorithm: %v", err)
}
return nil
}
-89
View File
@@ -1,89 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package wasm
import (
"log/slog"
"os"
"os/exec"
"testing"
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
"github.com/ultravioletrs/cocos/agent/events/mocks"
)
func TestNewAlgorithm(t *testing.T) {
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
eventsSvc := new(mocks.Service)
algoFile := "test.wasm"
args := []string{"arg1", "arg2"}
algo := NewAlgorithm(logger, eventsSvc, args, algoFile, "")
w, ok := algo.(*wasm)
if !ok {
t.Fatalf("NewAlgorithm did not return a *wasm")
}
if w.algoFile != algoFile {
t.Errorf("Expected algoFile to be %s, got %s", algoFile, w.algoFile)
}
if len(w.args) != len(args) {
t.Errorf("Expected %d args, got %d", len(args), len(w.args))
}
_, ok = w.stderr.(*logging.Stderr)
if !ok {
t.Errorf("Expected stderr to be *algorithm.Stderr")
}
_, ok = w.stdout.(*logging.Stdout)
if !ok {
t.Errorf("Expected stdout to be *algorithm.Stdout")
}
}
func TestRunError(t *testing.T) {
// Mock exec.Command to return an error
execCommand = mockExecCommandError
defer func() { execCommand = exec.Command }()
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
eventsSvc := new(mocks.Service)
algoFile := "test.wasm"
args := []string{"arg1", "arg2"}
w := NewAlgorithm(logger, eventsSvc, args, algoFile, "").(*wasm)
err := w.Run()
if err == nil {
t.Errorf("Run() should have returned an error")
}
}
func mockExecCommand(command string, args ...string) *exec.Cmd {
cs := []string{"-test.run=TestHelperProcess", "--", command}
cs = append(cs, args...)
cmd := exec.Command(os.Args[0], cs...)
cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1"}
return cmd
}
func mockExecCommandError(command string, args ...string) *exec.Cmd {
cmd := mockExecCommand(command, args...)
cmd.Env = append(cmd.Env, "GO_WANT_HELPER_PROCESS_ERROR=1")
return cmd
}
func TestHelperProcess(t *testing.T) {
if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
return
}
if os.Getenv("GO_WANT_HELPER_PROCESS_ERROR") == "1" {
os.Exit(1)
}
os.Exit(0)
}
var execCommand = exec.Command
+218
View File
@@ -0,0 +1,218 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"fmt"
"time"
"github.com/go-kit/kit/endpoint"
kitgrpc "github.com/go-kit/kit/transport/grpc"
"github.com/ultravioletrs/cocos/agent"
"google.golang.org/grpc"
)
const svcName = "agent.AgentService"
type grpcClient struct {
algo endpoint.Endpoint
data endpoint.Endpoint
result endpoint.Endpoint
attestation endpoint.Endpoint
timeout time.Duration
}
// NewClient returns new gRPC client instance.
func NewClient(conn *grpc.ClientConn, timeout time.Duration) agent.AgentServiceClient {
return &grpcClient{
algo: kitgrpc.NewClient(
conn,
svcName,
"Algo",
encodeAlgoRequest,
decodeAlgoResponse,
agent.AlgoResponse{},
).Endpoint(),
data: kitgrpc.NewClient(
conn,
svcName,
"Data",
encodeDataRequest,
decodeDataResponse,
agent.DataResponse{},
).Endpoint(),
result: kitgrpc.NewClient(
conn,
svcName,
"Result",
encodeResultRequest,
decodeResultResponse,
agent.ResultResponse{},
).Endpoint(),
attestation: kitgrpc.NewClient(
conn,
svcName,
"Attestation",
encodeAttestationRequest,
decodeAttestationResponse,
agent.AttestationResponse{},
).Endpoint(),
timeout: timeout,
}
}
// encodeAlgoRequest is a transport/grpc.EncodeRequestFunc that
// converts a user-domain algoReq to a gRPC request.
func encodeAlgoRequest(_ context.Context, request interface{}) (interface{}, error) {
req, ok := request.(*algoReq)
if !ok {
return nil, fmt.Errorf("invalid request type: %T", request)
}
return &agent.AlgoRequest{
Algorithm: req.Algorithm,
Provider: req.Provider,
Id: req.Id,
}, nil
}
// decodeAlgoResponse is a transport/grpc.DecodeResponseFunc that
// converts a gRPC AlgoResponse to a user-domain response.
func decodeAlgoResponse(_ context.Context, grpcResponse interface{}) (interface{}, error) {
_, ok := grpcResponse.(*agent.AlgoResponse)
if !ok {
return nil, fmt.Errorf("invalid response type: %T", grpcResponse)
}
return algoRes{}, nil
}
// encodeDataRequest is a transport/grpc.EncodeRequestFunc that
// converts a user-domain dataReq to a gRPC request.
func encodeDataRequest(_ context.Context, request interface{}) (interface{}, error) {
req, ok := request.(*dataReq)
if !ok {
return nil, fmt.Errorf("invalid request type: %T", request)
}
return &agent.DataRequest{
Dataset: req.Dataset,
Provider: req.Provider,
Id: req.Id,
}, nil
}
// decodeDataResponse is a transport/grpc.DecodeResponseFunc that
// converts a gRPC DataResponse to a user-domain response.
func decodeDataResponse(_ context.Context, grpcResponse interface{}) (interface{}, error) {
_, ok := grpcResponse.(*agent.DataResponse)
if !ok {
return nil, fmt.Errorf("invalid response type: %T", grpcResponse)
}
return dataRes{}, nil
}
// encodeResultRequest is a transport/grpc.EncodeRequestFunc that
// converts a user-domain resultReq to a gRPC request.
func encodeResultRequest(_ context.Context, request interface{}) (interface{}, error) {
req, ok := request.(*resultReq)
if !ok {
return nil, fmt.Errorf("invalid request type: %T", request)
}
return &agent.ResultRequest{
Consumer: req.Consumer,
}, nil
}
// decodeResultResponse is a transport/grpc.DecodeResponseFunc that
// converts a gRPC ResultResponse to a user-domain response.
func decodeResultResponse(_ context.Context, grpcResponse interface{}) (interface{}, error) {
response, ok := grpcResponse.(*agent.ResultResponse)
if !ok {
return nil, fmt.Errorf("invalid response type: %T", grpcResponse)
}
return resultRes{
File: response.File,
}, nil
}
// encodeAttestationRequest is a transport/grpc.EncodeRequestFunc that
// converts a user-domain attestationReq to a gRPC request.
func encodeAttestationRequest(_ context.Context, request interface{}) (interface{}, error) {
req, ok := request.(*attestationReq)
if !ok {
return nil, fmt.Errorf("invalid request type: %T", request)
}
return &agent.AttestationRequest{ReportData: req.ReportData[:]}, nil
}
// decodeAttestationResponse is a transport/grpc.DecodeResponseFunc that
// converts a gRPC AttestationResponse to a user-domain response.
func decodeAttestationResponse(_ context.Context, grpcResponse interface{}) (interface{}, error) {
response, ok := grpcResponse.(*agent.AttestationResponse)
if !ok {
return nil, fmt.Errorf("invalid response type: %T", grpcResponse)
}
return attestationRes{
File: response.File,
}, nil
}
// Algo implements the Algo method of the agent.AgentServiceClient interface.
func (c grpcClient) Algo(ctx context.Context, request *agent.AlgoRequest, _ ...grpc.CallOption) (*agent.AlgoResponse, error) {
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
_, err := c.algo(ctx, &algoReq{Algorithm: request.Algorithm, Provider: request.Provider, Id: request.Id})
if err != nil {
return nil, err
}
return &agent.AlgoResponse{}, nil
}
// Data implements the Data method of the agent.AgentServiceClient interface.
func (c grpcClient) Data(ctx context.Context, request *agent.DataRequest, _ ...grpc.CallOption) (*agent.DataResponse, error) {
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
_, err := c.data(ctx, &dataReq{Dataset: request.Dataset, Provider: request.Provider, Id: request.Id})
if err != nil {
return nil, err
}
return &agent.DataResponse{}, nil
}
// Result implements the Result method of the agent.AgentServiceClient interface.
func (c grpcClient) Result(ctx context.Context, request *agent.ResultRequest, _ ...grpc.CallOption) (*agent.ResultResponse, error) {
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
res, err := c.result(ctx, &resultReq{Consumer: request.Consumer})
if err != nil {
return nil, err
}
resultRes := res.(resultRes)
return &agent.ResultResponse{File: resultRes.File}, nil
}
// Result implements the Result method of the agent.AgentServiceClient interface.
func (c grpcClient) Attestation(ctx context.Context, request *agent.AttestationRequest, _ ...grpc.CallOption) (*agent.AttestationResponse, error) {
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
res, err := c.attestation(ctx, &attestationReq{ReportData: [agent.ReportDataSize]byte(request.ReportData)})
if err != nil {
return nil, err
}
attestationRes := res.(attestationRes)
return &agent.AttestationResponse{File: attestationRes.File}, nil
}
+4 -35
View File
@@ -7,7 +7,6 @@ import (
"github.com/go-kit/kit/endpoint"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/pkg/attestation"
)
func algoEndpoint(svc agent.Service) endpoint.Endpoint {
@@ -18,7 +17,7 @@ func algoEndpoint(svc agent.Service) endpoint.Endpoint {
return algoRes{}, err
}
algo := agent.Algorithm{Algorithm: req.Algorithm, Requirements: req.Requirements}
algo := agent.Algorithm{Algorithm: req.Algorithm, Provider: req.Provider, ID: req.Id}
err := svc.Algo(ctx, algo)
if err != nil {
@@ -37,7 +36,7 @@ func dataEndpoint(svc agent.Service) endpoint.Endpoint {
return dataRes{}, err
}
dataset := agent.Dataset{Dataset: req.Dataset, Filename: req.Filename}
dataset := agent.Dataset{Dataset: req.Dataset, Provider: req.Provider, ID: req.Id}
err := svc.Data(ctx, dataset)
if err != nil {
@@ -55,7 +54,7 @@ func resultEndpoint(svc agent.Service) endpoint.Endpoint {
if err := req.validate(); err != nil {
return resultRes{}, err
}
file, err := svc.Result(ctx)
file, err := svc.Result(ctx, req.Consumer)
if err != nil {
return resultRes{}, err
}
@@ -71,7 +70,7 @@ func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
if err := req.validate(); err != nil {
return attestationRes{}, err
}
file, err := svc.Attestation(ctx, req.TeeNonce, req.VtpmNonce, attestation.PlatformType(req.AttType))
file, err := svc.Attestation(ctx, req.ReportData)
if err != nil {
return attestationRes{}, err
}
@@ -79,33 +78,3 @@ func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
return attestationRes{File: file}, nil
}
}
func imaMeasurementsEndpoint(svc agent.Service) endpoint.Endpoint {
return func(ctx context.Context, request interface{}) (interface{}, error) {
req := request.(imaMeasurementsReq)
if err := req.validate(); err != nil {
return imaMeasurementsRes{}, err
}
file, pcr10, err := svc.IMAMeasurements(ctx)
if err != nil {
return imaMeasurementsRes{}, err
}
return imaMeasurementsRes{File: file, PCR10: pcr10}, nil
}
}
func attestationResultEndpoint(svc agent.Service) endpoint.Endpoint {
return func(ctx context.Context, request interface{}) (interface{}, error) {
req := request.(FetchAttestationResultReq)
if err := req.validate(); err != nil {
return fetchAttestationResultRes{}, err
}
file, err := svc.AttestationResult(ctx, req.tokenNonce, attestation.PlatformType(req.AttType))
if err != nil {
return fetchAttestationResultRes{}, err
}
return fetchAttestationResultRes{File: file}, nil
}
}
-227
View File
@@ -1,227 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/mock"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/agent/mocks"
"github.com/ultravioletrs/cocos/pkg/attestation"
"golang.org/x/crypto/sha3"
)
const svcErr = "Service Error"
func TestAlgoEndpoint(t *testing.T) {
svc := new(mocks.Service)
tests := []struct {
name string
req algoReq
expectedErr bool
}{
{
name: "Success",
req: algoReq{Algorithm: []byte("algorithm")},
},
{
name: "Validation Error",
req: algoReq{},
expectedErr: true,
},
{
name: "Service Error",
req: algoReq{Algorithm: []byte("algorithm")},
expectedErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.name == svcErr {
svc.On("Algo", context.Background(), agent.Algorithm{Algorithm: tt.req.Algorithm}).Return(errors.New("")).Once()
} else {
svc.On("Algo", context.Background(), agent.Algorithm{Algorithm: tt.req.Algorithm}).Return(nil).Once()
}
endpoint := algoEndpoint(svc)
_, err := endpoint(context.Background(), tt.req)
if (err != nil) != tt.expectedErr {
t.Errorf("algoEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
}
})
}
}
func TestDataEndpoint(t *testing.T) {
svc := new(mocks.Service)
tests := []struct {
name string
req dataReq
expectedErr bool
}{
{
name: "Success",
req: dataReq{Dataset: []byte("dataset")},
},
{
name: "Validation Error",
req: dataReq{},
expectedErr: true,
},
{
name: "Service Error",
req: dataReq{Dataset: []byte("dataset")},
expectedErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.name == svcErr {
svc.On("Data", context.Background(), agent.Dataset{Dataset: tt.req.Dataset}).Return(errors.New("")).Once()
} else {
svc.On("Data", context.Background(), agent.Dataset{Dataset: tt.req.Dataset}).Return(nil).Once()
}
endpoint := dataEndpoint(svc)
_, err := endpoint(context.Background(), tt.req)
if (err != nil) != tt.expectedErr {
t.Errorf("dataEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
}
})
}
}
func TestResultEndpoint(t *testing.T) {
svc := new(mocks.Service)
tests := []struct {
name string
req resultReq
expectedErr bool
}{
{
name: "Success",
req: resultReq{},
},
{
name: "Service Error",
req: resultReq{},
expectedErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.name == svcErr {
svc.On("Result", context.Background()).Return([]byte{}, errors.New("")).Once()
} else {
svc.On("Result", context.Background()).Return([]byte{}, nil).Once()
}
endpoint := resultEndpoint(svc)
res, err := endpoint(context.Background(), tt.req)
if (err != nil) != tt.expectedErr {
t.Errorf("resultEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
}
if err == nil {
_, ok := res.(resultRes)
if !ok {
t.Errorf("resultEndpoint() returned unexpected type %T", res)
}
}
})
}
}
func TestAttestationEndpoint(t *testing.T) {
svc := new(mocks.Service)
tests := []struct {
name string
req attestationReq
expectedErr bool
}{
{
name: "Success",
req: attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.SNP},
},
{
name: "Service Error",
req: attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.SNP},
expectedErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.name == svcErr {
svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, errors.New("")).Once()
} else {
svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, nil).Once()
}
endpoint := attestationEndpoint(svc)
res, err := endpoint(context.Background(), tt.req)
if (err != nil) != tt.expectedErr {
t.Errorf("attestationEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
}
if err == nil {
_, ok := res.(attestationRes)
if !ok {
t.Errorf("attestationEndpoint() returned unexpected type %T", res)
}
}
})
}
}
func TestAttestationResultEndpoint(t *testing.T) {
svc := new(mocks.Service)
tests := []struct {
name string
req FetchAttestationResultReq
mockErr error
expectedErr bool
}{
{
name: "Success",
req: FetchAttestationResultReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.AzureToken},
mockErr: nil,
expectedErr: false,
},
{
name: "Service Error",
req: FetchAttestationResultReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.AzureToken},
mockErr: errors.New("mock failure"),
expectedErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Only call service mock if validation is expected to pass
if err := tt.req.validate(); err == nil {
svc.On("AttestationResult", mock.Anything, tt.req.tokenNonce, attestation.PlatformType(tt.req.AttType)).
Return([]byte("mock file"), tt.mockErr).Once()
}
endpoint := attestationResultEndpoint(svc)
res, err := endpoint(context.Background(), tt.req)
if (err != nil) != tt.expectedErr {
t.Errorf("attestationResultEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
}
if !tt.expectedErr {
r, ok := res.(fetchAttestationResultRes)
if !ok {
t.Errorf("attestationResultEndpoint() returned unexpected type %T", res)
}
if string(r.File) != "mock file" {
t.Errorf("expected file content 'mock file', got %s", r.File)
}
}
svc.AssertExpectations(t)
})
}
}
-74
View File
@@ -1,74 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/agent/auth"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
type authInterceptor struct {
auth auth.Authenticator
}
type wrappedServerStream struct {
grpc.ServerStream
ctx context.Context
}
func (s *wrappedServerStream) Context() context.Context {
return s.ctx
}
func NewAuthInterceptor(authSvc auth.Authenticator) (grpc.UnaryServerInterceptor, grpc.StreamServerInterceptor) {
ai := &authInterceptor{auth: authSvc}
return ai.AuthUnaryInterceptor(), ai.AuthStreamInterceptor()
}
func (s *authInterceptor) AuthStreamInterceptor() grpc.StreamServerInterceptor {
return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
switch info.FullMethod {
case agent.AgentService_Algo_FullMethodName:
if _, err := s.auth.AuthenticateUser(stream.Context(), auth.AlgorithmProviderRole); err != nil {
return status.Errorf(codes.Unauthenticated, "%v", err.Error())
}
return handler(srv, stream)
case agent.AgentService_Data_FullMethodName:
ctx, err := s.auth.AuthenticateUser(stream.Context(), auth.DataProviderRole)
if err != nil {
return status.Errorf(codes.Unauthenticated, "%s", err.Error())
}
wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
return handler(srv, wrapped)
case agent.AgentService_Result_FullMethodName:
ctx, err := s.auth.AuthenticateUser(stream.Context(), auth.ConsumerRole)
if err != nil {
return status.Errorf(codes.Unauthenticated, "%v", err.Error())
}
wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
return handler(srv, wrapped)
default:
return handler(srv, stream)
}
}
}
func (s *authInterceptor) AuthUnaryInterceptor() grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
switch info.FullMethod {
case agent.AgentService_Result_FullMethodName:
ctx, err := s.auth.AuthenticateUser(ctx, auth.ConsumerRole)
if err != nil {
return nil, status.Errorf(codes.Unauthenticated, "%v", err.Error())
}
return handler(ctx, req)
default:
return handler(ctx, req)
}
}
}
-152
View File
@@ -1,152 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"testing"
"github.com/stretchr/testify/mock"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/agent/auth"
"github.com/ultravioletrs/cocos/agent/auth/mocks"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
func TestAuthUnaryInterceptor(t *testing.T) {
authmock := new(mocks.Authenticator)
tests := []struct {
name string
authorized bool
method string
role auth.UserRole
wantErr bool
}{
{
name: "authorized result method",
authorized: true,
method: agent.AgentService_Result_FullMethodName,
role: auth.ConsumerRole,
wantErr: false,
},
{
name: "unauthorized result method",
authorized: false,
method: agent.AgentService_Result_FullMethodName,
role: auth.ConsumerRole,
wantErr: true,
},
{
name: "other method",
authorized: false,
method: "OtherMethod",
role: auth.ConsumerRole,
wantErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
switch tt.authorized {
case true:
mockCall := authmock.On("AuthenticateUser", context.Background(), tt.role).Return(context.Background(), nil)
mockCall.Once()
case false:
mockCall := authmock.On("AuthenticateUser", context.Background(), tt.role).Return(context.Background(), auth.ErrMissingMetadata)
mockCall.Once()
}
unaryInt, _ := NewAuthInterceptor(authmock)
_, err := unaryInt(context.Background(), nil, &grpc.UnaryServerInfo{FullMethod: tt.method}, func(ctx context.Context, req interface{}) (interface{}, error) {
return nil, nil
})
if tt.wantErr && err == nil {
t.Errorf("expected error, got none")
} else if !tt.wantErr && err != nil {
t.Errorf("unexpected error: %v", err)
}
})
}
}
func TestAuthStreamInterceptor(t *testing.T) {
authmock := new(mocks.Authenticator)
tests := []struct {
name string
authorized bool
method string
role auth.UserRole
wantErr bool
}{
{
name: "authorized algo method",
authorized: true,
method: agent.AgentService_Algo_FullMethodName,
role: auth.AlgorithmProviderRole,
wantErr: false,
},
{
name: "unauthorized algo method",
authorized: false,
method: agent.AgentService_Algo_FullMethodName,
role: auth.AlgorithmProviderRole,
wantErr: true,
},
{
name: "authorized data method",
authorized: true,
method: agent.AgentService_Data_FullMethodName,
role: auth.DataProviderRole,
wantErr: false,
},
{
name: "unauthorized data method",
authorized: false,
method: agent.AgentService_Data_FullMethodName,
role: auth.DataProviderRole,
wantErr: true,
},
{
name: "other method",
authorized: false,
method: "OtherMethod",
role: auth.DataProviderRole,
wantErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
switch tt.authorized {
case true:
mockCall := authmock.On("AuthenticateUser", mock.Anything, tt.role).Return(context.Background(), nil)
mockCall.Once()
case false:
mockCall := authmock.On("AuthenticateUser", mock.Anything, tt.role).Return(context.Background(), auth.ErrMissingMetadata)
mockCall.Once()
}
_, streamInt := NewAuthInterceptor(authmock)
err := streamInt(nil, &mockServerStream{ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs())}, &grpc.StreamServerInfo{FullMethod: tt.method}, func(srv interface{}, stream grpc.ServerStream) error {
return nil
})
if tt.wantErr && err == nil {
t.Errorf("expected error, got none")
} else if !tt.wantErr && err != nil {
t.Errorf("unexpected error: %v", err)
}
})
}
}
type mockServerStream struct {
grpc.ServerStream
ctx context.Context
}
func (m *mockServerStream) Context() context.Context {
return m.ctx
}
+21 -36
View File
@@ -4,37 +4,49 @@ package grpc
import (
"errors"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
)
type algoReq struct {
Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
Requirements []byte
Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
Provider string `protobuf:"bytes,2,opt,name=provider,proto3" json:"provider,omitempty"`
Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"`
}
func (req algoReq) validate() error {
if len(req.Algorithm) == 0 {
return errors.New("algorithm binary is required")
}
if req.Id == "" {
return errors.New("malformed entity")
}
if req.Provider == "" {
return errors.New("malformed entity")
}
return nil
}
type dataReq struct {
Dataset []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
Filename string
Provider string `protobuf:"bytes,2,opt,name=provider,proto3" json:"provider,omitempty"`
Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"`
}
func (req dataReq) validate() error {
if len(req.Dataset) == 0 {
return errors.New("dataset CSV file is required")
}
if req.Id == "" {
return errors.New("malformed entity")
}
if req.Provider == "" {
return errors.New("malformed entity")
}
return nil
}
type resultReq struct{}
type resultReq struct {
Consumer string `protobuf:"bytes,1,opt,name=consumer,proto3" json:"consumer,omitempty"`
}
func (req resultReq) validate() error {
// No request parameters to validate, so no validation logic needed
@@ -42,36 +54,9 @@ func (req resultReq) validate() error {
}
type attestationReq struct {
TeeNonce [quoteprovider.Nonce]byte
VtpmNonce [vtpm.Nonce]byte
AttType attestation.PlatformType
}
type FetchAttestationResultReq struct {
tokenNonce [vtpm.Nonce]byte
AttType attestation.PlatformType
ReportData [64]byte
}
func (req attestationReq) validate() error {
return validateAttestationType(req.AttType)
}
func (req FetchAttestationResultReq) validate() error {
return validateAttestationType(req.AttType)
}
func validateAttestationType(attType attestation.PlatformType) error {
switch attType {
case attestation.SNP, attestation.VTPM, attestation.SNPvTPM, attestation.AzureToken:
return nil
default:
return errors.New("invalid attestation type")
}
}
type imaMeasurementsReq struct{}
func (req imaMeasurementsReq) validate() error {
// No request parameters to validate, so no validation logic needed
return nil
}
-9
View File
@@ -13,12 +13,3 @@ type resultRes struct {
type attestationRes struct {
File []byte
}
type imaMeasurementsRes struct {
File []byte
PCR10 []byte
}
type fetchAttestationResultRes struct {
File []byte `protobuf:"bytes,1,opt,name=AttestationResult,proto3" json:"AttestationResult,omitempty"`
}
+41 -242
View File
@@ -3,43 +3,18 @@
package grpc
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"strconv"
"github.com/go-kit/kit/transport/grpc"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
)
const (
bufferSize = 1024 * 1024
FileSizeKey = "file-size"
)
var (
ErrTEENonceLength = errors.New("malformed report data, expect less or equal to 64 bytes")
ErrVTPMNonceLength = errors.New("malformed vTPM nonce, expect less or equal to 32 bytes")
ErrTokenNonceLength = errors.New("malformed token nonce, expect less or equal to 32 bytes")
)
var _ agent.AgentServiceServer = (*grpcServer)(nil)
type grpcServer struct {
algo grpc.Handler
data grpc.Handler
result grpc.Handler
attestation grpc.Handler
imaMeasurements grpc.Handler
attestationResult grpc.Handler
algo grpc.Handler
data grpc.Handler
result grpc.Handler
attestation grpc.Handler
agent.UnimplementedAgentServiceServer
}
@@ -66,16 +41,6 @@ func NewServer(svc agent.Service) agent.AgentServiceServer {
decodeAttestationRequest,
encodeAttestationResponse,
),
imaMeasurements: grpc.NewServer(
imaMeasurementsEndpoint(svc),
decodeIMAMeasurementsRequest,
encodeIMAMeasurementsResponse,
),
attestationResult: grpc.NewServer(
attestationResultEndpoint(svc),
decodeAttestationResultRequest,
encodeAttestationResultResponse,
),
}
}
@@ -83,8 +48,9 @@ func decodeAlgoRequest(_ context.Context, grpcReq interface{}) (interface{}, err
req := grpcReq.(*agent.AlgoRequest)
return algoReq{
Algorithm: req.Algorithm,
Requirements: req.Requirements,
Algorithm: req.Algorithm,
Provider: req.Provider,
Id: req.Id,
}, nil
}
@@ -97,7 +63,8 @@ func decodeDataRequest(_ context.Context, grpcReq interface{}) (interface{}, err
return dataReq{
Dataset: req.Dataset,
Filename: req.Filename,
Provider: req.Provider,
Id: req.Id,
}, nil
}
@@ -106,7 +73,8 @@ func encodeDataResponse(_ context.Context, response interface{}) (interface{}, e
}
func decodeResultRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
return resultReq{}, nil
req := grpcReq.(*agent.ResultRequest)
return resultReq{Consumer: req.Consumer}, nil
}
func encodeResultResponse(_ context.Context, response interface{}) (interface{}, error) {
@@ -118,20 +86,10 @@ func encodeResultResponse(_ context.Context, response interface{}) (interface{},
func decodeAttestationRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
req := grpcReq.(*agent.AttestationRequest)
var reportData [quoteprovider.Nonce]byte
var nonce [vtpm.Nonce]byte
if len(req.TeeNonce) > quoteprovider.Nonce {
return nil, ErrTEENonceLength
if len(req.ReportData) != agent.ReportDataSize {
return nil, errors.New("malformed report data, expect 64 bytes")
}
if len(req.VtpmNonce) > vtpm.Nonce {
return nil, ErrVTPMNonceLength
}
copy(reportData[:], req.TeeNonce)
copy(nonce[:], req.VtpmNonce)
return attestationReq{TeeNonce: reportData, VtpmNonce: nonce, AttType: attestation.PlatformType(req.Type)}, nil
return attestationReq{ReportData: [agent.ReportDataSize]byte(req.ReportData)}, nil
}
func encodeAttestationResponse(_ context.Context, response interface{}) (interface{}, error) {
@@ -141,197 +99,38 @@ func encodeAttestationResponse(_ context.Context, response interface{}) (interfa
}, nil
}
func encodeAttestationResultResponse(_ context.Context, response interface{}) (interface{}, error) {
res := response.(fetchAttestationResultRes)
return &agent.AttestationResultResponse{
File: res.File,
}, nil
}
func decodeAttestationResultRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
req := grpcReq.(*agent.AttestationResultRequest)
var nonce [vtpm.Nonce]byte
if len(req.TokenNonce) > vtpm.Nonce {
return nil, ErrVTPMNonceLength
}
copy(nonce[:], req.TokenNonce)
return FetchAttestationResultReq{tokenNonce: nonce, AttType: attestation.PlatformType(req.Type)}, nil
}
// Algo implements agent.AgentServiceServer.
func (s *grpcServer) Algo(stream agent.AgentService_AlgoServer) error {
var algoFile, reqFile []byte
for {
algoChunk, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
return status.Error(codes.Internal, err.Error())
}
algoFile = append(algoFile, algoChunk.Algorithm...)
reqFile = append(reqFile, algoChunk.Requirements...)
}
_, res, err := s.algo.ServeGRPC(stream.Context(), &agent.AlgoRequest{Algorithm: algoFile, Requirements: reqFile})
if err != nil {
return err
}
ar := res.(*agent.AlgoResponse)
return stream.SendAndClose(ar)
}
// Data implements agent.AgentServiceServer.
func (s *grpcServer) Data(stream agent.AgentService_DataServer) error {
var dataFile []byte
var filename string
for {
dataChunk, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
return status.Error(codes.Internal, err.Error())
}
dataFile = append(dataFile, dataChunk.Dataset...)
filename = dataChunk.Filename
}
_, res, err := s.data.ServeGRPC(stream.Context(), &agent.DataRequest{Dataset: dataFile, Filename: filename})
if err != nil {
return err
}
ar := res.(*agent.DataResponse)
return stream.SendAndClose(ar)
}
func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_ResultServer) error {
_, res, err := s.result.ServeGRPC(stream.Context(), req)
if err != nil {
return err
}
rr := res.(*agent.ResultResponse)
if err := stream.SetHeader(metadata.New(map[string]string{FileSizeKey: fmt.Sprint(len(rr.File))})); err != nil {
return status.Error(codes.Internal, err.Error())
}
resultBuffer := bytes.NewBuffer(rr.File)
buf := make([]byte, bufferSize)
for {
n, err := resultBuffer.Read(buf)
if err == io.EOF {
break
}
if err != nil {
return status.Error(codes.Internal, err.Error())
}
if err := stream.Send(&agent.ResultResponse{File: buf[:n]}); err != nil {
return status.Error(codes.Internal, err.Error())
}
}
return nil
}
func (s *grpcServer) Attestation(req *agent.AttestationRequest, stream agent.AgentService_AttestationServer) error {
_, res, err := s.attestation.ServeGRPC(stream.Context(), req)
if err != nil {
return err
}
rr := res.(*agent.AttestationResponse)
if err := stream.SetHeader(metadata.New(map[string]string{FileSizeKey: fmt.Sprint(len(rr.File))})); err != nil {
return status.Error(codes.Internal, err.Error())
}
attestationBuffer := bytes.NewBuffer(rr.File)
buf := make([]byte, bufferSize)
for {
n, err := attestationBuffer.Read(buf)
if err == io.EOF {
break
}
if err != nil {
return status.Error(codes.Internal, err.Error())
}
if err := stream.Send(&agent.AttestationResponse{File: buf[:n]}); err != nil {
return status.Error(codes.Internal, err.Error())
}
}
return nil
}
func decodeIMAMeasurementsRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
return imaMeasurementsReq{}, nil
}
func encodeIMAMeasurementsResponse(_ context.Context, response interface{}) (interface{}, error) {
res := response.(imaMeasurementsRes)
return &agent.IMAMeasurementsResponse{
File: res.File,
Pcr10: res.PCR10,
}, nil
}
func (s *grpcServer) IMAMeasurements(req *agent.IMAMeasurementsRequest, stream agent.AgentService_IMAMeasurementsServer) error {
_, res, err := s.imaMeasurements.ServeGRPC(stream.Context(), req)
if err != nil {
return err
}
rr := res.(*agent.IMAMeasurementsResponse)
if err := stream.SetHeader(metadata.New(map[string]string{FileSizeKey: strconv.Itoa(len(rr.File))})); err != nil {
return status.Error(codes.Internal, err.Error())
}
imaBuff := bytes.NewBuffer(rr.File)
pcr10Buff := bytes.NewBuffer(rr.Pcr10)
imaResBuff := make([]byte, bufferSize)
pcr10ResBuff := make([]byte, bufferSize)
for {
nIma, errIma := imaBuff.Read(imaResBuff)
if errIma != nil && errIma != io.EOF {
return status.Error(codes.Internal, errIma.Error())
}
nPcr, errPcr := pcr10Buff.Read(pcr10ResBuff)
if errPcr != nil && errPcr != io.EOF {
return status.Error(codes.Internal, errPcr.Error())
}
if nIma == 0 && errIma == io.EOF &&
nPcr == 0 && errPcr == io.EOF {
break
}
if err := stream.Send(&agent.IMAMeasurementsResponse{File: imaResBuff[:nIma], Pcr10: pcr10ResBuff[:nPcr]}); err != nil {
return status.Error(codes.Internal, err.Error())
}
}
return nil
}
func (s *grpcServer) AttestationResult(ctx context.Context, req *agent.AttestationResultRequest) (*agent.AttestationResultResponse, error) {
_, res, err := s.attestationResult.ServeGRPC(ctx, req)
func (s *grpcServer) Algo(ctx context.Context, req *agent.AlgoRequest) (*agent.AlgoResponse, error) {
_, res, err := s.algo.ServeGRPC(ctx, req)
if err != nil {
return nil, err
}
rr, ok := res.(*agent.AttestationResultResponse)
ar := res.(*agent.AlgoResponse)
return ar, nil
}
if !ok {
return nil, status.Error(codes.Internal, "failed to cast response to FetchAttestationResultResponse")
func (s *grpcServer) Data(ctx context.Context, req *agent.DataRequest) (*agent.DataResponse, error) {
_, res, err := s.data.ServeGRPC(ctx, req)
if err != nil {
return nil, err
}
dr := res.(*agent.DataResponse)
return dr, nil
}
func (s *grpcServer) Result(ctx context.Context, req *agent.ResultRequest) (*agent.ResultResponse, error) {
_, res, err := s.result.ServeGRPC(ctx, req)
if err != nil {
return nil, err
}
rr := res.(*agent.ResultResponse)
return rr, nil
}
func (s *grpcServer) Attestation(ctx context.Context, req *agent.AttestationRequest) (*agent.AttestationResponse, error) {
_, res, err := s.attestation.ServeGRPC(ctx, req)
if err != nil {
return nil, err
}
rr := res.(*agent.AttestationResponse)
return rr, nil
}
-247
View File
@@ -1,247 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"io"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/agent/mocks"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
type MockAgentService_AlgoServer struct {
grpc.ServerStream
mock.Mock
ctx context.Context
}
func (m *MockAgentService_AlgoServer) Context() context.Context {
return m.ctx
}
func (m *MockAgentService_AlgoServer) Recv() (*agent.AlgoRequest, error) {
args := m.Called()
return args.Get(0).(*agent.AlgoRequest), args.Error(1)
}
func (m *MockAgentService_AlgoServer) SendAndClose(resp *agent.AlgoResponse) error {
args := m.Called(resp)
return args.Error(0)
}
type MockAgentService_DataServer struct {
grpc.ServerStream
mock.Mock
ctx context.Context
}
func (m *MockAgentService_DataServer) Context() context.Context {
return m.ctx
}
func (m *MockAgentService_DataServer) Recv() (*agent.DataRequest, error) {
args := m.Called()
return args.Get(0).(*agent.DataRequest), args.Error(1)
}
func (m *MockAgentService_DataServer) SendAndClose(resp *agent.DataResponse) error {
args := m.Called(resp)
return args.Error(0)
}
type MockAgentService_ResultServer struct {
grpc.ServerStream
mock.Mock
ctx context.Context
}
func (m *MockAgentService_ResultServer) Context() context.Context {
return m.ctx
}
func (m *MockAgentService_ResultServer) SetHeader(metadata.MD) error {
return nil
}
func (m *MockAgentService_ResultServer) Send(resp *agent.ResultResponse) error {
args := m.Called(resp)
return args.Error(0)
}
type MockAgentService_AttestationServer struct {
grpc.ServerStream
mock.Mock
ctx context.Context
}
func (m *MockAgentService_AttestationServer) Context() context.Context {
return m.ctx
}
func (m *MockAgentService_AttestationServer) Send(resp *agent.AttestationResponse) error {
args := m.Called(resp)
return args.Error(0)
}
func (m *MockAgentService_AttestationServer) SetHeader(metadata.MD) error {
return nil
}
func TestAlgo(t *testing.T) {
mockService := new(mocks.Service)
server := NewServer(mockService)
mockStream := &MockAgentService_AlgoServer{ctx: context.Background()}
mockStream.On("Recv").Return(&agent.AlgoRequest{Algorithm: []byte("algo"), Requirements: []byte("req")}, nil).Once()
mockStream.On("Recv").Return(&agent.AlgoRequest{}, io.EOF)
mockStream.On("SendAndClose", &agent.AlgoResponse{}).Return(nil)
mockService.On("Algo", context.Background(), agent.Algorithm{Algorithm: []byte("algo"), Requirements: []byte("req")}).Return(nil)
err := server.Algo(mockStream)
assert.NoError(t, err)
mockStream.AssertExpectations(t)
mockService.AssertExpectations(t)
}
func TestData(t *testing.T) {
mockService := new(mocks.Service)
server := NewServer(mockService)
mockStream := &MockAgentService_DataServer{ctx: context.Background()}
mockStream.On("Recv").Return(&agent.DataRequest{Dataset: []byte("data"), Filename: "test.txt"}, nil).Once()
mockStream.On("Recv").Return(&agent.DataRequest{}, io.EOF)
mockStream.On("SendAndClose", &agent.DataResponse{}).Return(nil)
mockService.On("Data", context.Background(), agent.Dataset{Dataset: []byte("data"), Filename: "test.txt"}).Return(nil)
err := server.Data(mockStream)
assert.NoError(t, err)
mockStream.AssertExpectations(t)
mockService.AssertExpectations(t)
}
func TestResult(t *testing.T) {
mockService := new(mocks.Service)
server := NewServer(mockService)
mockStream := &MockAgentService_ResultServer{ctx: context.Background()}
mockService.On("Result", mock.Anything).Return([]byte("result data"), nil)
mockStream.On("Send", mock.AnythingOfType("*agent.ResultResponse")).Return(nil)
err := server.Result(&agent.ResultRequest{}, mockStream)
assert.NoError(t, err)
mockStream.AssertExpectations(t)
mockService.AssertExpectations(t)
}
func TestAttestation(t *testing.T) {
mockService := new(mocks.Service)
server := NewServer(mockService)
mockStream := &MockAgentService_AttestationServer{ctx: context.Background()}
mockStream.On("Send", mock.AnythingOfType("*agent.AttestationResponse")).Return(nil)
reportData := [quoteprovider.Nonce]byte{}
vtpmNonce := [vtpm.Nonce]byte{}
attestationType := attestation.SNP
mockService.On("Attestation", mock.Anything, reportData, vtpmNonce, attestationType).Return([]byte("attestation data"), nil)
err := server.Attestation(&agent.AttestationRequest{TeeNonce: reportData[:]}, mockStream)
assert.NoError(t, err)
mockService.AssertExpectations(t)
}
func TestAttestationResult(t *testing.T) {
mockService := new(mocks.Service)
server := NewServer(mockService)
vtpmNonce := [vtpm.Nonce]byte{}
attestationType := attestation.SNP
mockService.On("AttestationResult", mock.Anything, vtpmNonce, attestationType).Return([]byte("attestation data"), nil)
resp, err := server.AttestationResult(context.Background(), &agent.AttestationResultRequest{TokenNonce: vtpmNonce[:]})
assert.NoError(t, err)
assert.Equal(t, []byte("attestation data"), resp.File)
mockService.AssertExpectations(t)
}
func TestDecodeAlgoRequest(t *testing.T) {
req := &agent.AlgoRequest{Algorithm: []byte("algo"), Requirements: []byte("req")}
decoded, err := decodeAlgoRequest(context.Background(), req)
assert.NoError(t, err)
assert.Equal(t, algoReq{Algorithm: []byte("algo"), Requirements: []byte("req")}, decoded)
}
func TestEncodeAlgoResponse(t *testing.T) {
encoded, err := encodeAlgoResponse(context.Background(), algoRes{})
assert.NoError(t, err)
assert.Equal(t, &agent.AlgoResponse{}, encoded)
}
func TestDecodeDataRequest(t *testing.T) {
req := &agent.DataRequest{Dataset: []byte("data"), Filename: "test.txt"}
decoded, err := decodeDataRequest(context.Background(), req)
assert.NoError(t, err)
assert.Equal(t, dataReq{Dataset: []byte("data"), Filename: "test.txt"}, decoded)
}
func TestEncodeDataResponse(t *testing.T) {
encoded, err := encodeDataResponse(context.Background(), dataRes{})
assert.NoError(t, err)
assert.Equal(t, &agent.DataResponse{}, encoded)
}
func TestDecodeResultRequest(t *testing.T) {
decoded, err := decodeResultRequest(context.Background(), &agent.ResultRequest{})
assert.NoError(t, err)
assert.Equal(t, resultReq{}, decoded)
}
func TestEncodeResultResponse(t *testing.T) {
encoded, err := encodeResultResponse(context.Background(), resultRes{File: []byte("result")})
assert.NoError(t, err)
assert.Equal(t, &agent.ResultResponse{File: []byte("result")}, encoded)
}
func TestDecodeAttestationRequest(t *testing.T) {
nonce := [quoteprovider.Nonce]byte{}
req := &agent.AttestationRequest{TeeNonce: nonce[:]}
decoded, err := decodeAttestationRequest(context.Background(), req)
assert.NoError(t, err)
assert.Equal(t, attestationReq{TeeNonce: nonce}, decoded)
}
func TestEncodeAttestationResponse(t *testing.T) {
encoded, err := encodeAttestationResponse(context.Background(), attestationRes{File: []byte("attestation")})
assert.NoError(t, err)
assert.Equal(t, &agent.AttestationResponse{File: []byte("attestation")}, encoded)
}
func TestEncodeAttestationResultResponse(t *testing.T) {
encoded, err := encodeAttestationResultResponse(context.Background(), fetchAttestationResultRes{File: []byte("attestation")})
assert.NoError(t, err)
assert.Equal(t, &agent.AttestationResultResponse{File: []byte("attestation")}, encoded)
}
func TestDecodeAttestationResultRequest(t *testing.T) {
nonce := [vtpm.Nonce]byte{}
req := &agent.AttestationResultRequest{TokenNonce: nonce[:]}
decoded, err := decodeAttestationResultRequest(context.Background(), req)
assert.NoError(t, err)
assert.Equal(t, FetchAttestationResultReq{tokenNonce: nonce}, decoded)
}
+4 -70
View File
@@ -13,9 +13,6 @@ import (
"time"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
)
var _ agent.Service = (*loggingMiddleware)(nil)
@@ -30,43 +27,6 @@ func LoggingMiddleware(svc agent.Service, logger *slog.Logger) agent.Service {
return &loggingMiddleware{logger, svc}
}
// State implements agent.Service.
func (lm *loggingMiddleware) State() (state string) {
defer func(begin time.Time) {
message := fmt.Sprintf("Method State took %s to complete with state %s", time.Since(begin), state)
lm.logger.Debug(message)
}(time.Now())
return lm.svc.State()
}
// InitComputation implements agent.Service.
func (lm *loggingMiddleware) InitComputation(ctx context.Context, cmp agent.Computation) (err error) {
defer func(begin time.Time) {
message := fmt.Sprintf("Method InitComputation for computation id %s took %s to complete", cmp.ID, time.Since(begin))
if err != nil {
lm.logger.WithGroup(cmp.ID).Warn(fmt.Sprintf("%s with error: %s", message, err))
return
}
lm.logger.WithGroup(cmp.ID).Info(fmt.Sprintf("%s without errors", message))
}(time.Now())
return lm.svc.InitComputation(ctx, cmp)
}
// StopComputation implements agent.Service.
func (lm *loggingMiddleware) StopComputation(ctx context.Context) (err error) {
defer func(begin time.Time) {
message := fmt.Sprintf("Method StopComputation took %s to complete", time.Since(begin))
if err != nil {
lm.logger.Warn(fmt.Sprintf("%s with error: %s", message, err))
return
}
lm.logger.Info(fmt.Sprintf("%s without errors", message))
}(time.Now())
return lm.svc.StopComputation(ctx)
}
func (lm *loggingMiddleware) Algo(ctx context.Context, algorithm agent.Algorithm) (err error) {
defer func(begin time.Time) {
message := fmt.Sprintf("Method Algo took %s to complete", time.Since(begin))
@@ -93,7 +53,7 @@ func (lm *loggingMiddleware) Data(ctx context.Context, dataset agent.Dataset) (e
return lm.svc.Data(ctx, dataset)
}
func (lm *loggingMiddleware) Result(ctx context.Context) (response []byte, err error) {
func (lm *loggingMiddleware) Result(ctx context.Context, consumer string) (response []byte, err error) {
defer func(begin time.Time) {
message := fmt.Sprintf("Method Result took %s to complete", time.Since(begin))
if err != nil {
@@ -103,10 +63,10 @@ func (lm *loggingMiddleware) Result(ctx context.Context) (response []byte, err e
lm.logger.Info(fmt.Sprintf("%s without errors", message))
}(time.Now())
return lm.svc.Result(ctx)
return lm.svc.Result(ctx, consumer)
}
func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) (response []byte, err error) {
func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [agent.ReportDataSize]byte) (response []byte, err error) {
defer func(begin time.Time) {
message := fmt.Sprintf("Method Attestation took %s to complete", time.Since(begin))
if err != nil {
@@ -116,31 +76,5 @@ func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [quotep
lm.logger.Info(fmt.Sprintf("%s without errors", message))
}(time.Now())
return lm.svc.Attestation(ctx, reportData, nonce, attType)
}
func (lm *loggingMiddleware) IMAMeasurements(ctx context.Context) (file []byte, pcr10 []byte, err error) {
defer func(begin time.Time) {
message := fmt.Sprintf("Method IMAMeasurements took %s to complete", time.Since(begin))
if err != nil {
lm.logger.Warn(fmt.Sprintf("%s with error: %s", message, err))
return
}
lm.logger.Info(fmt.Sprintf("%s without errors", message))
}(time.Now())
return lm.svc.IMAMeasurements(ctx)
}
func (lm *loggingMiddleware) AttestationResult(ctx context.Context, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) (response []byte, err error) {
defer func(begin time.Time) {
message := fmt.Sprintf("Method AttestationResult took %s to complete", time.Since(begin))
if err != nil {
lm.logger.Warn(fmt.Sprintf("%s with error: %s", message, err))
return
}
lm.logger.Info(fmt.Sprintf("%s without errors", message))
}(time.Now())
return lm.svc.AttestationResult(ctx, nonce, attType)
return lm.svc.Attestation(ctx, reportData)
}
+4 -55
View File
@@ -12,9 +12,6 @@ import (
"github.com/go-kit/kit/metrics"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
)
var _ agent.Service = (*metricsMiddleware)(nil)
@@ -35,36 +32,6 @@ func MetricsMiddleware(svc agent.Service, counter metrics.Counter, latency metri
}
}
// State implements agent.Service.
func (ms *metricsMiddleware) State() string {
defer func(begin time.Time) {
ms.counter.With("method", "state").Add(1)
ms.latency.With("method", "state").Observe(time.Since(begin).Seconds())
}(time.Now())
return ms.svc.State()
}
// InitComputation implements agent.Service.
func (ms *metricsMiddleware) InitComputation(ctx context.Context, cmp agent.Computation) error {
defer func(begin time.Time) {
ms.counter.With("method", "init_computation").Add(1)
ms.latency.With("method", "init_computation").Observe(time.Since(begin).Seconds())
}(time.Now())
return ms.svc.InitComputation(ctx, cmp)
}
// StopComputation implements agent.Service.
func (ms *metricsMiddleware) StopComputation(ctx context.Context) error {
defer func(begin time.Time) {
ms.counter.With("method", "stop_computation").Add(1)
ms.latency.With("method", "stop_computation").Observe(time.Since(begin).Seconds())
}(time.Now())
return ms.svc.StopComputation(ctx)
}
func (ms *metricsMiddleware) Algo(ctx context.Context, algorithm agent.Algorithm) error {
defer func(begin time.Time) {
ms.counter.With("method", "algo").Add(1)
@@ -83,38 +50,20 @@ func (ms *metricsMiddleware) Data(ctx context.Context, dataset agent.Dataset) er
return ms.svc.Data(ctx, dataset)
}
func (ms *metricsMiddleware) Result(ctx context.Context) ([]byte, error) {
func (ms *metricsMiddleware) Result(ctx context.Context, consumer string) ([]byte, error) {
defer func(begin time.Time) {
ms.counter.With("method", "result").Add(1)
ms.latency.With("method", "result").Observe(time.Since(begin).Seconds())
}(time.Now())
return ms.svc.Result(ctx)
return ms.svc.Result(ctx, consumer)
}
func (ms *metricsMiddleware) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error) {
func (ms *metricsMiddleware) Attestation(ctx context.Context, reportData [agent.ReportDataSize]byte) ([]byte, error) {
defer func(begin time.Time) {
ms.counter.With("method", "attestation").Add(1)
ms.latency.With("method", "attestation").Observe(time.Since(begin).Seconds())
}(time.Now())
return ms.svc.Attestation(ctx, reportData, nonce, attType)
}
func (ms *metricsMiddleware) AttestationResult(ctx context.Context, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error) {
defer func(begin time.Time) {
ms.counter.With("method", "attestation_result").Add(1)
ms.latency.With("method", "attestation_result").Observe(time.Since(begin).Seconds())
}(time.Now())
return ms.svc.AttestationResult(ctx, nonce, attType)
}
func (ms *metricsMiddleware) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
defer func(begin time.Time) {
ms.counter.With("method", "imameasurements").Add(1)
ms.latency.With("method", "imameasurements").Observe(time.Since(begin).Seconds())
}(time.Now())
return ms.svc.IMAMeasurements(ctx)
return ms.svc.Attestation(ctx, reportData)
}
-173
View File
@@ -1,173 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package auth
import (
"context"
"crypto"
"crypto/ecdsa"
"crypto/ed25519"
"crypto/rsa"
"crypto/sha256"
"crypto/x509"
"encoding/base64"
"github.com/absmach/magistrala/pkg/errors"
"github.com/ultravioletrs/cocos/agent"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
)
type UserRole string
const (
UserMetadataKey = "user-id"
SignatureMetadataKey = "signature"
ConsumerRole UserRole = "consumer"
DataProviderRole UserRole = "data-provider"
AlgorithmProviderRole UserRole = "algorithm-provider"
)
var (
ErrMissingMetadata = errors.New("missing metadata")
ErrInvalidMetadata = errors.New("invalid metadata")
ErrSignatureVerificationFailed = errors.New("signature verification failed")
)
type Authenticator interface {
AuthenticateUser(ctx context.Context, role UserRole) (context.Context, error)
}
type service struct {
resultConsumers []interface{}
datasetProviders []interface{}
algorithmProvider interface{}
}
func New(manifest agent.Computation) (Authenticator, error) {
s := &service{}
for _, rc := range manifest.ResultConsumers {
pubKey, err := x509.ParsePKIXPublicKey(rc.UserKey)
if err != nil {
return nil, err
}
pKey, err := decodePublicKey(pubKey)
if err != nil {
return nil, err
}
s.resultConsumers = append(s.resultConsumers, pKey)
}
for _, dp := range manifest.Datasets {
pubKey, err := x509.ParsePKIXPublicKey(dp.UserKey)
if err != nil {
return nil, err
}
pKey, err := decodePublicKey(pubKey)
if err != nil {
return nil, err
}
s.datasetProviders = append(s.datasetProviders, pKey)
}
pubKey, err := x509.ParsePKIXPublicKey(manifest.Algorithm.UserKey)
if err != nil {
return nil, err
}
pKey, err := decodePublicKey(pubKey)
if err != nil {
return nil, err
}
s.algorithmProvider = pKey
return s, nil
}
func extractSignature(md metadata.MD) (string, error) {
signature := md.Get(SignatureMetadataKey)
if len(signature) != 1 {
return "", status.Errorf(codes.Unauthenticated, "invalid metadata")
}
return signature[0], nil
}
func verifySignature(role UserRole, signature string, publicKey any) error {
hash := sha256.Sum256([]byte(role))
sigByte, err := base64.StdEncoding.DecodeString(signature)
if err != nil {
return err
}
var ok bool
switch publicKey := publicKey.(type) {
case *rsa.PublicKey:
if err = rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, hash[:], sigByte); err != nil {
return err
}
return nil
case *ecdsa.PublicKey:
ok = ecdsa.VerifyASN1(publicKey, hash[:], sigByte)
case ed25519.PublicKey:
ok = ed25519.Verify(publicKey, []byte(role), sigByte)
}
if !ok {
return ErrSignatureVerificationFailed
}
return nil
}
func (s *service) AuthenticateUser(ctx context.Context, role UserRole) (context.Context, error) {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
return nil, ErrMissingMetadata
}
signature, err := extractSignature(md)
if err != nil {
return nil, errors.Wrap(err, ErrInvalidMetadata)
}
switch role {
case ConsumerRole:
for i, rc := range s.resultConsumers {
if err := verifySignature(role, signature, rc); err == nil {
return agent.IndexToContext(ctx, i), nil
}
}
case DataProviderRole:
for _, dp := range s.datasetProviders {
if err := verifySignature(role, signature, dp); err == nil {
return ctx, nil
}
}
case AlgorithmProviderRole:
if err := verifySignature(role, signature, s.algorithmProvider); err == nil {
return ctx, nil
}
}
return ctx, ErrSignatureVerificationFailed
}
func decodePublicKey(key any) (pubKey any, err error) {
switch key := key.(type) {
case *rsa.PublicKey:
return key, nil
case *ecdsa.PublicKey:
return key, nil
case ed25519.PublicKey:
return key, nil
default:
return nil, errors.New("unsupported public key type")
}
}
-158
View File
@@ -1,158 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package auth
import (
"context"
"crypto"
"crypto/ecdsa"
"crypto/ed25519"
"crypto/elliptic"
"crypto/rand"
"crypto/rsa"
"crypto/sha256"
"crypto/x509"
"encoding/base64"
"testing"
"github.com/absmach/magistrala/pkg/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/ultravioletrs/cocos/agent"
"google.golang.org/grpc/metadata"
)
func TestAuthenticateUser(t *testing.T) {
resultConsumerKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
dataProviderKey, err := rsa.GenerateKey(rand.Reader, 2048)
require.NoError(t, err)
pubEd25519Key, algorithmProviderKey, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
resultConsumerPubKey, err := x509.MarshalPKIXPublicKey(&resultConsumerKey.PublicKey)
require.NoError(t, err)
dataProviderPubKey, err := x509.MarshalPKIXPublicKey(&dataProviderKey.PublicKey)
require.NoError(t, err)
algorithmProviderPubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
require.NoError(t, err)
manifest := agent.Computation{
ResultConsumers: []agent.ResultConsumer{{UserKey: resultConsumerPubKey}},
Datasets: []agent.Dataset{{UserKey: dataProviderPubKey}},
Algorithm: agent.Algorithm{UserKey: algorithmProviderPubKey},
}
auth, err := New(manifest)
if err != nil {
t.Fatalf("failed to create authenticator: %v", err)
}
testCases := []struct {
name string
role UserRole
key any
expectedErr error
}{
{
name: "valid result consumer",
role: ConsumerRole,
key: resultConsumerKey,
expectedErr: nil,
},
{
name: "valid data provider",
role: DataProviderRole,
key: dataProviderKey,
expectedErr: nil,
},
{
name: "valid algorithm provider",
role: AlgorithmProviderRole,
key: algorithmProviderKey,
expectedErr: nil,
},
{
name: "invalid role",
role: "invalid-role",
key: resultConsumerKey,
expectedErr: ErrSignatureVerificationFailed,
},
{
name: "invalid key",
role: ConsumerRole,
key: dataProviderKey,
expectedErr: ErrSignatureVerificationFailed,
},
{
name: "missing signature",
role: ConsumerRole,
key: resultConsumerKey,
expectedErr: ErrInvalidMetadata,
},
{
name: "missing metadata",
role: ConsumerRole,
key: resultConsumerKey,
expectedErr: ErrMissingMetadata,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
signature, err := signRole(tc.role, tc.key)
if err != nil {
t.Fatalf("failed to sign role: %v", err)
}
ctx := context.Background()
switch tc.name {
case "missing signature":
ctx = metadata.NewIncomingContext(ctx, metadata.Pairs())
case "missing metadata":
default:
ctx = metadata.NewIncomingContext(ctx, metadata.Pairs(SignatureMetadataKey, signature))
}
ctx, err = auth.AuthenticateUser(ctx, tc.role)
assert.True(t, errors.Contains(err, tc.expectedErr), "expected error %v, got %v", tc.expectedErr, err)
if err == nil {
switch id, ok := agent.IndexFromContext(ctx); {
case tc.role == ConsumerRole:
assert.True(t, ok, "expected index in context")
assert.Equal(t, 0, id, "expected index 0 in context")
default:
assert.False(t, ok, "expected no index in context")
}
}
})
}
}
func signRole(role UserRole, key crypto.PrivateKey) (string, error) {
var signature []byte
var err error
switch k := key.(type) {
case ed25519.PrivateKey:
signature, err = k.Sign(rand.Reader, []byte(role), crypto.Hash(0))
case *rsa.PrivateKey, *ecdsa.PrivateKey:
hash := sha256.Sum256([]byte(role))
signer := key.(crypto.Signer)
signature, err = signer.Sign(rand.Reader, hash[:], crypto.SHA256)
default:
return "", errors.New("unsupported key type")
}
if err != nil {
return "", err
}
return base64.StdEncoding.EncodeToString(signature), nil
}
-100
View File
@@ -1,100 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mocks
import (
context "context"
auth "github.com/ultravioletrs/cocos/agent/auth"
mock "github.com/stretchr/testify/mock"
)
// Authenticator is an autogenerated mock type for the Authenticator type
type Authenticator struct {
mock.Mock
}
type Authenticator_Expecter struct {
mock *mock.Mock
}
func (_m *Authenticator) EXPECT() *Authenticator_Expecter {
return &Authenticator_Expecter{mock: &_m.Mock}
}
// AuthenticateUser provides a mock function with given fields: ctx, role
func (_m *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRole) (context.Context, error) {
ret := _m.Called(ctx, role)
if len(ret) == 0 {
panic("no return value specified for AuthenticateUser")
}
var r0 context.Context
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) (context.Context, error)); ok {
return rf(ctx, role)
}
if rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) context.Context); ok {
r0 = rf(ctx, role)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(context.Context)
}
}
if rf, ok := ret.Get(1).(func(context.Context, auth.UserRole) error); ok {
r1 = rf(ctx, role)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// Authenticator_AuthenticateUser_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AuthenticateUser'
type Authenticator_AuthenticateUser_Call struct {
*mock.Call
}
// AuthenticateUser is a helper method to define mock.On call
// - ctx context.Context
// - role auth.UserRole
func (_e *Authenticator_Expecter) AuthenticateUser(ctx interface{}, role interface{}) *Authenticator_AuthenticateUser_Call {
return &Authenticator_AuthenticateUser_Call{Call: _e.mock.On("AuthenticateUser", ctx, role)}
}
func (_c *Authenticator_AuthenticateUser_Call) Run(run func(ctx context.Context, role auth.UserRole)) *Authenticator_AuthenticateUser_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(auth.UserRole))
})
return _c
}
func (_c *Authenticator_AuthenticateUser_Call) Return(_a0 context.Context, _a1 error) *Authenticator_AuthenticateUser_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *Authenticator_AuthenticateUser_Call) RunAndReturn(run func(context.Context, auth.UserRole) (context.Context, error)) *Authenticator_AuthenticateUser_Call {
_c.Call.Return(run)
return _c
}
// NewAuthenticator creates a new instance of Authenticator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewAuthenticator(t interface {
mock.TestingT
Cleanup(func())
}) *Authenticator {
mock := &Authenticator{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
+45 -48
View File
@@ -3,35 +3,34 @@
package agent
import (
"context"
"encoding/json"
"fmt"
"google.golang.org/grpc/metadata"
"reflect"
)
var _ fmt.Stringer = (*Datasets)(nil)
var (
_ fmt.Stringer = (*Datasets)(nil)
_ fmt.Stringer = (*Algorithms)(nil)
)
type AgentConfig struct {
Port string `json:"port,omitempty"`
CertFile string `json:"cert_file,omitempty"`
KeyFile string `json:"server_key,omitempty"`
ServerCAFile string `json:"server_ca_file,omitempty"`
ClientCAFile string `json:"client_ca_file,omitempty"`
AttestedTls bool `json:"attested_tls,omitempty"`
LogLevel string `json:"log_level"`
Host string `json:"host"`
Port string `json:"port"`
CertFile string `json:"cert_file"`
KeyFile string `json:"server_key"`
ServerCAFile string `json:"server_ca_file"`
ClientCAFile string `json:"client_ca_file"`
}
type Computation struct {
ID string `json:"id,omitempty"`
Name string `json:"name,omitempty"`
Description string `json:"description,omitempty"`
Datasets Datasets `json:"datasets,omitempty"`
Algorithm Algorithm `json:"algorithm,omitempty"`
ResultConsumers []ResultConsumer `json:"result_consumers,omitempty"`
}
type ResultConsumer struct {
UserKey []byte `json:"user_key,omitempty"`
ID string `json:"id,omitempty"`
Name string `json:"name,omitempty"`
Description string `json:"description,omitempty"`
Datasets Datasets `json:"datasets,omitempty"`
Algorithms Algorithms `json:"algorithms,omitempty"`
ResultConsumers []string `json:"result_consumers,omitempty"`
AgentConfig AgentConfig `json:"agent_config,omitempty"`
}
func (d *Datasets) String() string {
@@ -42,44 +41,42 @@ func (d *Datasets) String() string {
return string(dat)
}
func (a *Algorithms) String() string {
dat, err := json.Marshal(a)
if err != nil {
return ""
}
return string(dat)
}
type Dataset struct {
Dataset []byte `json:"-"`
Hash [32]byte `json:"hash,omitempty"`
UserKey []byte `json:"user_key,omitempty"`
Filename string `json:"filename,omitempty"`
Provider string `json:"provider,omitempty"`
ID string `json:"id,omitempty"`
}
type Datasets []Dataset
type Algorithm struct {
Algorithm []byte `json:"-"`
Hash [32]byte `json:"hash,omitempty"`
UserKey []byte `json:"user_key,omitempty"`
Requirements []byte `json:"-"`
Algorithm []byte `json:"-"`
Hash [32]byte `json:"hash,omitempty"`
Provider string `json:"provider,omitempty"`
ID string `json:"id,omitempty"`
}
type ManifestIndexKey struct{}
type Algorithms []Algorithm
func IndexToContext(ctx context.Context, index int) context.Context {
return context.WithValue(ctx, ManifestIndexKey{}, index)
}
func IndexFromContext(ctx context.Context) (int, bool) {
index, ok := ctx.Value(ManifestIndexKey{}).(int)
return index, ok
}
const DecompressKey = "decompress"
func DecompressFromContext(ctx context.Context) bool {
vals := metadata.ValueFromIncomingContext(ctx, DecompressKey)
if len(vals) == 0 {
return false
func containsID(slice interface{}, id string) int {
rangeOnMe := reflect.ValueOf(slice)
for i := 0; i < rangeOnMe.Len(); i++ {
s := rangeOnMe.Index(i)
f := s.FieldByName("ID")
if f.IsValid() {
if f.Interface() == id {
return i
}
}
}
return vals[0] == "true"
}
func DecompressToContext(ctx context.Context, decompress bool) context.Context {
return metadata.AppendToOutgoingContext(ctx, DecompressKey, fmt.Sprintf("%t", decompress))
return -1
}
-131
View File
@@ -1,131 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package agent
import (
"context"
"encoding/json"
"reflect"
"testing"
"google.golang.org/grpc/metadata"
)
func TestDatasetsString(t *testing.T) {
datasets := Datasets{
{
Hash: [32]byte{1, 2, 3},
UserKey: []byte("user_key"),
Filename: "test.dat",
},
}
expected := `[{"hash":[1,2,3,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],"user_key":"dXNlcl9rZXk=","filename":"test.dat"}]`
result := datasets.String()
if result != expected {
t.Errorf("Datasets.String() = %v, want %v", result, expected)
}
}
func TestIndexToContext(t *testing.T) {
ctx := context.Background()
index := 5
newCtx := IndexToContext(ctx, index)
result, ok := IndexFromContext(newCtx)
if !ok {
t.Errorf("IndexFromContext() ok = false, want true")
}
if result != index {
t.Errorf("IndexFromContext() = %v, want %v", result, index)
}
}
func TestDecompressFromContext(t *testing.T) {
tests := []struct {
name string
ctx context.Context
expected bool
}{
{
name: "No decompress metadata",
ctx: context.Background(),
expected: false,
},
{
name: "Decompress true",
ctx: metadata.NewIncomingContext(
context.Background(),
metadata.Pairs(DecompressKey, "true"),
),
expected: true,
},
{
name: "Decompress false",
ctx: metadata.NewIncomingContext(
context.Background(),
metadata.Pairs(DecompressKey, "false"),
),
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := DecompressFromContext(tt.ctx)
if result != tt.expected {
t.Errorf("DecompressFromContext() = %v, want %v", result, tt.expected)
}
})
}
}
func TestDecompressToContext(t *testing.T) {
ctx := context.Background()
decompress := true
newCtx := DecompressToContext(ctx, decompress)
md, ok := metadata.FromOutgoingContext(newCtx)
if !ok {
t.Errorf("metadata.FromOutgoingContext() ok = false, want true")
}
vals := md.Get(DecompressKey)
if len(vals) != 1 {
t.Errorf("len(md.Get(DecompressKey)) = %v, want 1", len(vals))
}
if vals[0] != "true" {
t.Errorf("md.Get(DecompressKey)[0] = %v, want 'true'", vals[0])
}
}
func TestAgentConfigJSON(t *testing.T) {
config := AgentConfig{
Port: "8080",
CertFile: "cert.pem",
KeyFile: "key.pem",
ServerCAFile: "server_ca.pem",
ClientCAFile: "client_ca.pem",
AttestedTls: true,
}
data, err := json.Marshal(config)
if err != nil {
t.Fatalf("Failed to marshal AgentConfig: %v", err)
}
var unmarshaledConfig AgentConfig
err = json.Unmarshal(data, &unmarshaledConfig)
if err != nil {
t.Fatalf("Failed to unmarshal AgentConfig: %v", err)
}
if !reflect.DeepEqual(config, unmarshaledConfig) {
t.Errorf("Unmarshaled config does not match original. Got %+v, want %+v", unmarshaledConfig, config)
}
}
-360
View File
@@ -1,360 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"log/slog"
"sync"
"time"
"github.com/absmach/magistrala/pkg/errors"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/agent/cvms"
"github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
"github.com/ultravioletrs/cocos/agent/cvms/server"
pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
"golang.org/x/sync/errgroup"
"google.golang.org/protobuf/proto"
)
const (
reconnectInterval = 5 * time.Second
sendTimeout = 5 * time.Second
pendingMsgFile = "pending_messages.json"
)
var (
errCorruptedManifest = errors.New("received manifest may be corrupted")
errUnknonwMessageType = errors.New("unknown message type")
)
type PendingMessage struct {
Message *cvms.ClientStreamMessage
Time time.Time
}
type CVMSClient struct {
mu sync.Mutex
stream cvms.Service_ProcessClient
svc agent.Service
messageQueue chan *cvms.ClientStreamMessage
logger *slog.Logger
runReqManager *runRequestManager
sp server.AgentServer
storage storage.Storage
reconnectFn func(context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error)
grpcClient pkggrpc.Client
}
// NewClient returns new gRPC client instance.
func NewClient(stream cvms.Service_ProcessClient, svc agent.Service, messageQueue chan *cvms.ClientStreamMessage, logger *slog.Logger, sp server.AgentServer, storageDir string, reconnectFn func(context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error), grpcClient pkggrpc.Client) (*CVMSClient, error) {
store, err := storage.NewFileStorage(storageDir)
if err != nil {
return nil, err
}
return &CVMSClient{
stream: stream,
svc: svc,
messageQueue: messageQueue,
logger: logger,
runReqManager: newRunRequestManager(),
sp: sp,
storage: store,
reconnectFn: reconnectFn,
grpcClient: grpcClient,
}, nil
}
func (client *CVMSClient) Process(ctx context.Context, cancel context.CancelFunc) error {
for {
err := client.processWithRetry(ctx)
if ctx.Err() != nil {
return ctx.Err()
}
client.logger.Info("Connection lost, attempting to reconnect...", "error", err)
time.Sleep(reconnectInterval)
grpcClient, stream, err := client.reconnectFn(ctx)
if err != nil {
client.logger.Error("Failed to reconnect", "error", err)
continue
}
client.mu.Lock()
client.stream = stream
client.grpcClient = grpcClient
client.mu.Unlock()
}
}
func (client *CVMSClient) processWithRetry(ctx context.Context) error {
eg, ctx := errgroup.WithContext(ctx)
eg.Go(func() error {
return client.handleIncomingMessages(ctx)
})
eg.Go(func() error {
return client.handleOutgoingMessages(ctx)
})
return eg.Wait()
}
func (client *CVMSClient) handleIncomingMessages(ctx context.Context) error {
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
req, err := client.stream.Recv()
if err != nil {
return err
}
if err := client.processIncomingMessage(ctx, req); err != nil {
return err
}
}
}
}
func (client *CVMSClient) handleOutgoingMessages(ctx context.Context) error {
pendingMsgs, err := client.storage.Load()
if err != nil {
client.logger.Error("Failed to load pending messages", "error", err)
} else {
client.sendPendingMessages(pendingMsgs)
}
for {
select {
case <-ctx.Done():
return ctx.Err()
case msg := <-client.messageQueue:
if err := client.sendStreamMessage(msg); err != nil {
if err := client.storage.Add(msg); err != nil {
client.logger.Error("Failed to store pending message", "error", err)
}
client.logger.Error("Failed to send message, stored for retry", "error", err)
}
}
}
}
func (client *CVMSClient) sendStreamMessage(msg *cvms.ClientStreamMessage) error {
client.mu.Lock()
defer client.mu.Unlock()
return client.stream.Send(msg)
}
func (client *CVMSClient) sendPendingMessages(pending []storage.Message) {
for _, pm := range pending {
if err := client.sendStreamMessage(pm.Message); err != nil {
if err := client.storage.Add(pm.Message); err != nil {
client.logger.Error("Failed to store pending message", "error", err)
}
client.logger.Error("Failed to resend pending message", "error", err)
} else {
client.logger.Info("Successfully resent pending message")
}
}
if err := client.storage.Clear(); err != nil {
client.logger.Error("Failed to clear pending messages", "error", err)
}
}
func (client *CVMSClient) processIncomingMessage(ctx context.Context, req *cvms.ServerStreamMessage) error {
switch mes := req.Message.(type) {
case *cvms.ServerStreamMessage_RunReqChunks:
return client.handleRunReqChunks(ctx, mes)
case *cvms.ServerStreamMessage_StopComputation:
go client.handleStopComputation(ctx, mes)
case *cvms.ServerStreamMessage_AgentStateReq:
client.handleAgentStateReq(mes)
case *cvms.ServerStreamMessage_DisconnectReq:
client.logger.Info("Received disconnect request")
client.mu.Lock()
if err := client.grpcClient.Close(); err != nil {
client.logger.Error("Failed to close gRPC client", "error", err)
}
client.mu.Unlock()
default:
return errUnknonwMessageType
}
return nil
}
func (client *CVMSClient) handleAgentStateReq(mes *cvms.ServerStreamMessage_AgentStateReq) {
state := client.svc.State()
msg := &cvms.ClientStreamMessage_AgentStateRes{
AgentStateRes: &cvms.AgentStateRes{
State: state,
Id: mes.AgentStateReq.Id,
},
}
client.sendMessage(&cvms.ClientStreamMessage{Message: msg})
}
func (client *CVMSClient) handleRunReqChunks(ctx context.Context, msg *cvms.ServerStreamMessage_RunReqChunks) error {
buffer, complete := client.runReqManager.addChunk(msg.RunReqChunks.Id, msg.RunReqChunks.Data, msg.RunReqChunks.IsLast)
if complete {
var runReq cvms.ComputationRunReq
if err := proto.Unmarshal(buffer, &runReq); err != nil {
return errors.Wrap(err, errCorruptedManifest)
}
go client.executeRun(ctx, &runReq)
}
return nil
}
func (client *CVMSClient) executeRun(ctx context.Context, runReq *cvms.ComputationRunReq) {
ac := agent.Computation{
ID: runReq.Id,
Name: runReq.Name,
Description: runReq.Description,
}
if runReq.Algorithm != nil {
ac.Algorithm = agent.Algorithm{
Hash: [32]byte(runReq.Algorithm.Hash),
UserKey: runReq.Algorithm.UserKey,
}
}
for _, ds := range runReq.Datasets {
ac.Datasets = append(ac.Datasets, agent.Dataset{
Hash: [32]byte(ds.Hash),
UserKey: ds.UserKey,
})
}
for _, rc := range runReq.ResultConsumers {
ac.ResultConsumers = append(ac.ResultConsumers, agent.ResultConsumer{
UserKey: rc.UserKey,
})
}
if err := client.svc.InitComputation(ctx, ac); err != nil {
client.logger.Warn(err.Error())
return
}
client.mu.Lock()
defer client.mu.Unlock()
if runReq.AgentConfig == nil {
runReq.AgentConfig = &cvms.AgentConfig{}
}
runRes := &cvms.ClientStreamMessage_RunRes{
RunRes: &cvms.RunResponse{
ComputationId: runReq.Id,
},
}
if err := client.sp.Start(agent.AgentConfig{
Port: runReq.AgentConfig.Port,
CertFile: runReq.AgentConfig.CertFile,
KeyFile: runReq.AgentConfig.KeyFile,
ServerCAFile: runReq.AgentConfig.ServerCaFile,
ClientCAFile: runReq.AgentConfig.ClientCaFile,
AttestedTls: runReq.AgentConfig.AttestedTls,
}, ac); err != nil {
client.logger.Warn(err.Error())
runRes.RunRes.Error = err.Error()
}
client.sendMessage(&cvms.ClientStreamMessage{Message: runRes})
}
func (client *CVMSClient) handleStopComputation(ctx context.Context, mes *cvms.ServerStreamMessage_StopComputation) {
msg := &cvms.ClientStreamMessage_StopComputationRes{
StopComputationRes: &cvms.StopComputationResponse{
ComputationId: mes.StopComputation.ComputationId,
},
}
if err := client.svc.StopComputation(ctx); err != nil {
msg.StopComputationRes.Message = err.Error()
}
client.mu.Lock()
if err := client.sp.Stop(); err != nil {
msg.StopComputationRes.Message = err.Error()
}
client.mu.Unlock()
client.sendMessage(&cvms.ClientStreamMessage{Message: msg})
}
func (client *CVMSClient) sendMessage(mes *cvms.ClientStreamMessage) {
ctx, cancel := context.WithTimeout(context.Background(), sendTimeout)
defer cancel()
select {
case client.messageQueue <- mes:
case <-ctx.Done():
client.logger.Warn("Failed to send message: timeout exceeded")
}
}
type runRequestManager struct {
requests map[string]*runRequest
mu sync.Mutex
}
type runRequest struct {
buffer []byte
lastChunk time.Time
timer *time.Timer
}
func newRunRequestManager() *runRequestManager {
return &runRequestManager{
requests: make(map[string]*runRequest),
}
}
func (m *runRequestManager) addChunk(id string, chunk []byte, isLast bool) ([]byte, bool) {
m.mu.Lock()
defer m.mu.Unlock()
req, exists := m.requests[id]
if !exists {
req = &runRequest{
buffer: make([]byte, 0),
lastChunk: time.Now(),
timer: time.AfterFunc(runReqTimeout, func() { m.timeoutRequest(id) }),
}
m.requests[id] = req
}
req.buffer = append(req.buffer, chunk...)
req.lastChunk = time.Now()
req.timer.Reset(runReqTimeout)
if isLast {
delete(m.requests, id)
req.timer.Stop()
return req.buffer, true
}
return nil, false
}
func (m *runRequestManager) timeoutRequest(id string) {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.requests, id)
// Log timeout or handle it as needed
}
-215
View File
@@ -1,215 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"testing"
"time"
mglog "github.com/absmach/magistrala/logger"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/ultravioletrs/cocos/agent/cvms"
servermocks "github.com/ultravioletrs/cocos/agent/cvms/server/mocks"
"github.com/ultravioletrs/cocos/agent/mocks"
pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
clientmocks "github.com/ultravioletrs/cocos/pkg/clients/grpc/mocks"
"google.golang.org/grpc"
"google.golang.org/protobuf/proto"
)
type mockStream struct {
mock.Mock
grpc.ClientStream
}
func (m *mockStream) Recv() (*cvms.ServerStreamMessage, error) {
args := m.Called()
return args.Get(0).(*cvms.ServerStreamMessage), args.Error(1)
}
func (m *mockStream) Send(msg *cvms.ClientStreamMessage) error {
args := m.Called(msg)
return args.Error(0)
}
func TestManagerClient_Process(t *testing.T) {
tests := []struct {
name string
setupMocks func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer)
expectError bool
errorMsg string
}{
{
name: "Stop computation",
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer) {
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
Message: &cvms.ServerStreamMessage_StopComputation{
StopComputation: &cvms.StopComputation{},
},
}, nil)
mockStream.On("Send", mock.Anything).Return(nil)
mockSvc.On("StopComputation", mock.Anything).Return(nil)
mockServerSvc.On("Stop").Return(nil)
},
expectError: true,
errorMsg: "context deadline exceeded",
},
{
name: "Run request chunks",
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer) {
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
Message: &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{},
},
}, nil)
mockStream.On("Send", mock.Anything).Return(nil).Once()
mockSvc.On("Run", mock.Anything, mock.Anything).Return("", assert.AnError).Once()
},
expectError: true,
},
{
name: "Receive error",
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer) {
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{}, assert.AnError)
},
expectError: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage)
logger := mglog.NewMock()
go func() {
<-messageQueue
}()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
tc.setupMocks(mockStream, mockSvc, mockServerSvc)
err = client.Process(ctx, cancel)
if tc.expectError {
assert.Error(t, err)
if tc.errorMsg != "" {
assert.Contains(t, err.Error(), tc.errorMsg)
}
} else {
assert.NoError(t, err)
}
})
}
}
func TestManagerClient_handleRunReqChunks(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
runReq := &cvms.ComputationRunReq{
Id: "test-id",
}
runReqBytes, _ := proto.Marshal(runReq)
chunk1 := &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{
Id: "chunk-1",
Data: runReqBytes[:len(runReqBytes)/2],
IsLast: false,
},
}
chunk2 := &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{
Id: "chunk-1",
Data: runReqBytes[len(runReqBytes)/2:],
IsLast: true,
},
}
mockSvc.On("InitComputation", mock.Anything, mock.Anything).Return(nil)
mockServerSvc.On("Start", mock.Anything, mock.Anything, mock.Anything).Return(nil)
err = client.handleRunReqChunks(context.Background(), chunk1)
assert.NoError(t, err)
err = client.handleRunReqChunks(context.Background(), chunk2)
assert.NoError(t, err)
// Wait for the goroutine to finish
time.Sleep(50 * time.Millisecond)
mockSvc.AssertExpectations(t)
assert.Len(t, messageQueue, 1)
msg := <-messageQueue
runRes, ok := msg.Message.(*cvms.ClientStreamMessage_RunRes)
assert.True(t, ok)
assert.Equal(t, "test-id", runRes.RunRes.ComputationId)
}
func TestManagerClient_handleStopComputation(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
stopReq := &cvms.ServerStreamMessage_StopComputation{
StopComputation: &cvms.StopComputation{
ComputationId: "test-comp-id",
},
}
mockSvc.On("StopComputation", mock.Anything).Return(nil)
mockServerSvc.On("Stop").Return(nil)
client.handleStopComputation(context.Background(), stopReq)
// Wait for the goroutine to finish
time.Sleep(50 * time.Millisecond)
mockSvc.AssertExpectations(t)
assert.Len(t, messageQueue, 1)
msg := <-messageQueue
stopRes, ok := msg.Message.(*cvms.ClientStreamMessage_StopComputationRes)
assert.True(t, ok)
assert.Equal(t, "test-comp-id", stopRes.StopComputationRes.ComputationId)
assert.Empty(t, stopRes.StopComputationRes.Message)
}
func TestManagerClient_timeoutRequest(t *testing.T) {
rm := newRunRequestManager()
rm.requests["test-id"] = &runRequest{
timer: time.NewTimer(100 * time.Millisecond),
buffer: []byte("test-data"),
lastChunk: time.Now(),
}
rm.timeoutRequest("test-id")
assert.Len(t, rm.requests, 0)
}
-5
View File
@@ -1,5 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Package grpc contains implementation of kit service gRPC API.
package grpc
-133
View File
@@ -1,133 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"bytes"
"context"
"errors"
"io"
"time"
"github.com/ultravioletrs/cocos/agent/cvms"
"golang.org/x/sync/errgroup"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/peer"
"google.golang.org/protobuf/proto"
)
var (
_ cvms.ServiceServer = (*grpcServer)(nil)
ErrUnexpectedMsg = errors.New("unknown message type")
)
const (
bufferSize = 1024 * 1024 // 1 MB
runReqTimeout = 30 * time.Second
)
type SendFunc func(*cvms.ServerStreamMessage) error
type grpcServer struct {
cvms.UnimplementedServiceServer
incoming chan *cvms.ClientStreamMessage
svc Service
}
type Service interface {
Run(ctx context.Context, ipAddress string, sendMessage SendFunc, authInfo credentials.AuthInfo)
}
// NewServer returns new AuthServiceServer instance.
func NewServer(incoming chan *cvms.ClientStreamMessage, svc Service) cvms.ServiceServer {
return &grpcServer{
incoming: incoming,
svc: svc,
}
}
func (s *grpcServer) Process(stream cvms.Service_ProcessServer) error {
client, ok := peer.FromContext(stream.Context())
if !ok {
return errors.New("failed to get peer info")
}
eg, ctx := errgroup.WithContext(stream.Context())
eg.Go(func() error {
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
req, err := stream.Recv()
if err != nil {
return err
}
s.incoming <- req
}
}
})
eg.Go(func() error {
sendMessage := func(msg *cvms.ServerStreamMessage) error {
select {
case <-ctx.Done():
return ctx.Err()
default:
switch m := msg.Message.(type) {
case *cvms.ServerStreamMessage_RunReq:
return s.sendRunReqInChunks(stream, m.RunReq)
default:
return stream.Send(msg)
}
}
}
s.svc.Run(ctx, client.Addr.String(), sendMessage, client.AuthInfo)
return nil
})
return eg.Wait()
}
func (s *grpcServer) sendRunReqInChunks(stream cvms.Service_ProcessServer, runReq *cvms.ComputationRunReq) error {
data, err := proto.Marshal(runReq)
if err != nil {
return err
}
dataBuffer := bytes.NewBuffer(data)
buf := make([]byte, bufferSize)
for {
n, err := dataBuffer.Read(buf)
isLast := false
if err == io.EOF {
isLast = true
} else if err != nil {
return err
}
chunk := &cvms.ServerStreamMessage{
Message: &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{
Id: runReq.Id,
Data: buf[:n],
IsLast: isLast,
},
},
}
if err := stream.Send(chunk); err != nil {
return err
}
if isLast {
break
}
}
return nil
}
-273
View File
@@ -1,273 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"testing"
"time"
"github.com/absmach/magistrala/pkg/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/ultravioletrs/cocos/agent/cvms"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/peer"
)
type mockServerStream struct {
mock.Mock
cvms.Service_ProcessServer
}
func (m *mockServerStream) Send(msg *cvms.ServerStreamMessage) error {
args := m.Called(msg)
return args.Error(0)
}
func (m *mockServerStream) Recv() (*cvms.ClientStreamMessage, error) {
args := m.Called()
return args.Get(0).(*cvms.ClientStreamMessage), args.Error(1)
}
func (m *mockServerStream) Context() context.Context {
args := m.Called()
return args.Get(0).(context.Context)
}
type mockService struct {
mock.Mock
}
func (m *mockService) Run(ctx context.Context, ipAddress string, sendMessage SendFunc, authInfo credentials.AuthInfo) {
m.Called(ctx, ipAddress, sendMessage, authInfo)
}
func TestNewServer(t *testing.T) {
incoming := make(chan *cvms.ClientStreamMessage)
mockSvc := new(mockService)
server := NewServer(incoming, mockSvc)
assert.NotNil(t, server)
assert.IsType(t, &grpcServer{}, server)
}
func TestGrpcServer_Process(t *testing.T) {
tests := []struct {
name string
recvReturn *cvms.ClientStreamMessage
recvError error
expectedError string
}{
{
name: "Process with context deadline exceeded",
recvReturn: &cvms.ClientStreamMessage{},
recvError: nil,
expectedError: "context deadline exceeded",
},
{
name: "Process with Recv error",
recvReturn: &cvms.ClientStreamMessage{},
recvError: errors.New("recv error"),
expectedError: "recv error",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
incoming := make(chan *cvms.ClientStreamMessage, 1)
mockSvc := new(mockService)
server := NewServer(incoming, mockSvc).(*grpcServer)
mockStream := new(mockServerStream)
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
mockStream.On("Context").Return(peer.NewContext(ctx, &peer.Peer{
Addr: mockAddr{},
AuthInfo: mockAuthInfo{},
}))
if tt.recvError == nil {
go func() {
for mes := range incoming {
assert.NotNil(t, mes)
}
}()
}
mockStream.On("Recv").Return(tt.recvReturn, tt.recvError)
mockSvc.On("Run", mock.Anything, "test", mock.Anything, mock.AnythingOfType("mockAuthInfo")).Return()
err := server.Process(mockStream)
assert.Error(t, err)
assert.Contains(t, err.Error(), tt.expectedError)
mockStream.AssertExpectations(t)
mockSvc.AssertExpectations(t)
})
}
}
func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
incoming := make(chan *cvms.ClientStreamMessage)
mockSvc := new(mockService)
server := NewServer(incoming, mockSvc).(*grpcServer)
mockStream := new(mockServerStream)
runReq := &cvms.ComputationRunReq{
Id: "test-id",
}
largePayload := make([]byte, bufferSize*2)
for i := range largePayload {
largePayload[i] = byte(i % 256)
}
runReq.Algorithm = &cvms.Algorithm{}
runReq.Algorithm.UserKey = largePayload
mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(nil).Times(4)
err := server.sendRunReqInChunks(mockStream, runReq)
assert.NoError(t, err)
mockStream.AssertExpectations(t)
calls := mockStream.Calls
assert.Equal(t, 4, len(calls))
for i, call := range calls {
msg := call.Arguments[0].(*cvms.ServerStreamMessage)
chunk := msg.GetRunReqChunks()
assert.NotNil(t, chunk)
assert.Equal(t, "test-id", chunk.Id)
if i < 3 {
assert.False(t, chunk.IsLast)
} else {
assert.Equal(t, 0, len(chunk.Data))
assert.True(t, chunk.IsLast)
}
}
}
type mockAddr struct{}
func (mockAddr) Network() string { return "test network" }
func (mockAddr) String() string { return "test" }
type mockAuthInfo struct{}
func (mockAuthInfo) AuthType() string { return "test auth" }
func TestGrpcServer_ProcessWithMockService(t *testing.T) {
tests := []struct {
name string
setupMockFn func(*mockService, *mockServerStream)
}{
{
name: "Run Request Test",
setupMockFn: func(mockSvc *mockService, mockStream *mockServerStream) {
mockSvc.On("Run", mock.Anything, "test", mock.Anything, mock.AnythingOfType("mockAuthInfo")).
Run(func(args mock.Arguments) {
sendFunc := args.Get(2).(SendFunc)
runReq := &cvms.ComputationRunReq{Id: "test-run-id"}
err := sendFunc(&cvms.ServerStreamMessage{
Message: &cvms.ServerStreamMessage_RunReq{
RunReq: runReq,
},
})
assert.NoError(t, err)
}).
Return()
mockStream.On("Send", mock.MatchedBy(func(msg *cvms.ServerStreamMessage) bool {
chunks := msg.GetRunReqChunks()
return chunks != nil && chunks.Id == "test-run-id"
})).Return(nil)
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
incoming := make(chan *cvms.ClientStreamMessage, 10)
mockSvc := new(mockService)
server := NewServer(incoming, mockSvc).(*grpcServer)
go func() {
for mes := range incoming {
assert.NotNil(t, mes)
}
}()
mockStream := new(mockServerStream)
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
defer cancel()
peerCtx := peer.NewContext(ctx, &peer.Peer{
Addr: mockAddr{},
AuthInfo: mockAuthInfo{},
})
mockStream.On("Context").Return(peerCtx)
mockStream.On("Recv").Return(&cvms.ClientStreamMessage{}, nil).Maybe()
tt.setupMockFn(mockSvc, mockStream)
go func() {
time.Sleep(150 * time.Millisecond)
cancel()
}()
err := server.Process(mockStream)
assert.Error(t, err)
assert.Contains(t, err.Error(), "context canceled")
mockStream.AssertExpectations(t)
mockSvc.AssertExpectations(t)
})
}
}
func TestGrpcServer_sendRunReqInChunksError(t *testing.T) {
incoming := make(chan *cvms.ClientStreamMessage)
mockSvc := new(mockService)
server := NewServer(incoming, mockSvc).(*grpcServer)
mockStream := new(mockServerStream)
runReq := &cvms.ComputationRunReq{
Id: "test-id",
}
// Simulate an error when sending
mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(errors.New("send error")).Once()
err := server.sendRunReqInChunks(mockStream, runReq)
assert.Error(t, err)
assert.Contains(t, err.Error(), "send error")
mockStream.AssertExpectations(t)
}
func TestGrpcServer_ProcessMissingPeerInfo(t *testing.T) {
incoming := make(chan *cvms.ClientStreamMessage)
mockSvc := new(mockService)
server := NewServer(incoming, mockSvc).(*grpcServer)
mockStream := new(mockServerStream)
ctx := context.Background()
// Return a context without peer info
mockStream.On("Context").Return(ctx)
err := server.Process(mockStream)
assert.Error(t, err)
assert.Contains(t, err.Error(), "failed to get peer info")
mockStream.AssertExpectations(t)
}
@@ -1,234 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mocks
import (
mock "github.com/stretchr/testify/mock"
cvms "github.com/ultravioletrs/cocos/agent/cvms"
storage "github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
)
// Storage is an autogenerated mock type for the Storage type
type Storage struct {
mock.Mock
}
type Storage_Expecter struct {
mock *mock.Mock
}
func (_m *Storage) EXPECT() *Storage_Expecter {
return &Storage_Expecter{mock: &_m.Mock}
}
// Add provides a mock function with given fields: msg
func (_m *Storage) Add(msg *cvms.ClientStreamMessage) error {
ret := _m.Called(msg)
if len(ret) == 0 {
panic("no return value specified for Add")
}
var r0 error
if rf, ok := ret.Get(0).(func(*cvms.ClientStreamMessage) error); ok {
r0 = rf(msg)
} else {
r0 = ret.Error(0)
}
return r0
}
// Storage_Add_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Add'
type Storage_Add_Call struct {
*mock.Call
}
// Add is a helper method to define mock.On call
// - msg *cvms.ClientStreamMessage
func (_e *Storage_Expecter) Add(msg interface{}) *Storage_Add_Call {
return &Storage_Add_Call{Call: _e.mock.On("Add", msg)}
}
func (_c *Storage_Add_Call) Run(run func(msg *cvms.ClientStreamMessage)) *Storage_Add_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*cvms.ClientStreamMessage))
})
return _c
}
func (_c *Storage_Add_Call) Return(_a0 error) *Storage_Add_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *Storage_Add_Call) RunAndReturn(run func(*cvms.ClientStreamMessage) error) *Storage_Add_Call {
_c.Call.Return(run)
return _c
}
// Clear provides a mock function with no fields
func (_m *Storage) Clear() error {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Clear")
}
var r0 error
if rf, ok := ret.Get(0).(func() error); ok {
r0 = rf()
} else {
r0 = ret.Error(0)
}
return r0
}
// Storage_Clear_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Clear'
type Storage_Clear_Call struct {
*mock.Call
}
// Clear is a helper method to define mock.On call
func (_e *Storage_Expecter) Clear() *Storage_Clear_Call {
return &Storage_Clear_Call{Call: _e.mock.On("Clear")}
}
func (_c *Storage_Clear_Call) Run(run func()) *Storage_Clear_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *Storage_Clear_Call) Return(_a0 error) *Storage_Clear_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *Storage_Clear_Call) RunAndReturn(run func() error) *Storage_Clear_Call {
_c.Call.Return(run)
return _c
}
// Load provides a mock function with no fields
func (_m *Storage) Load() ([]storage.Message, error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Load")
}
var r0 []storage.Message
var r1 error
if rf, ok := ret.Get(0).(func() ([]storage.Message, error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() []storage.Message); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]storage.Message)
}
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// Storage_Load_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Load'
type Storage_Load_Call struct {
*mock.Call
}
// Load is a helper method to define mock.On call
func (_e *Storage_Expecter) Load() *Storage_Load_Call {
return &Storage_Load_Call{Call: _e.mock.On("Load")}
}
func (_c *Storage_Load_Call) Run(run func()) *Storage_Load_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *Storage_Load_Call) Return(_a0 []storage.Message, _a1 error) *Storage_Load_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *Storage_Load_Call) RunAndReturn(run func() ([]storage.Message, error)) *Storage_Load_Call {
_c.Call.Return(run)
return _c
}
// Save provides a mock function with given fields: messages
func (_m *Storage) Save(messages []storage.Message) error {
ret := _m.Called(messages)
if len(ret) == 0 {
panic("no return value specified for Save")
}
var r0 error
if rf, ok := ret.Get(0).(func([]storage.Message) error); ok {
r0 = rf(messages)
} else {
r0 = ret.Error(0)
}
return r0
}
// Storage_Save_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Save'
type Storage_Save_Call struct {
*mock.Call
}
// Save is a helper method to define mock.On call
// - messages []storage.Message
func (_e *Storage_Expecter) Save(messages interface{}) *Storage_Save_Call {
return &Storage_Save_Call{Call: _e.mock.On("Save", messages)}
}
func (_c *Storage_Save_Call) Run(run func(messages []storage.Message)) *Storage_Save_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].([]storage.Message))
})
return _c
}
func (_c *Storage_Save_Call) Return(_a0 error) *Storage_Save_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *Storage_Save_Call) RunAndReturn(run func([]storage.Message) error) *Storage_Save_Call {
_c.Call.Return(run)
return _c
}
// NewStorage creates a new instance of Storage. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewStorage(t interface {
mock.TestingT
Cleanup(func())
}) *Storage {
mock := &Storage{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
-111
View File
@@ -1,111 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package storage
import (
"encoding/json"
"os"
"path/filepath"
"sync"
"time"
"github.com/ultravioletrs/cocos/agent/cvms"
)
// Message represents a pending message with its timestamp.
type Message struct {
Message *cvms.ClientStreamMessage
Time time.Time
}
// Storage defines the interface for message persistence operations.
type Storage interface {
// Load retrieves all pending messages from storage.
Load() ([]Message, error)
// Save persists the given messages to storage.
Save(messages []Message) error
// Add appends a new message to storage.
Add(msg *cvms.ClientStreamMessage) error
// Clear removes all messages from storage.
Clear() error
}
// FileStorage implements Storage interface using file-based persistence.
type FileStorage struct {
mu sync.Mutex
path string
msgs []Message
}
// NewFileStorage creates a new file-based storage instance.
func NewFileStorage(storageDir string) (*FileStorage, error) {
if err := os.MkdirAll(storageDir, 0o755); err != nil {
return nil, err
}
return &FileStorage{
path: filepath.Join(storageDir, "pending_messages.json"),
msgs: make([]Message, 0),
}, nil
}
func (fs *FileStorage) Load() ([]Message, error) {
fs.mu.Lock()
defer fs.mu.Unlock()
data, err := os.ReadFile(fs.path)
if os.IsNotExist(err) {
return nil, nil
}
if err != nil {
return nil, err
}
if err := json.Unmarshal(data, &fs.msgs); err != nil {
return nil, err
}
return fs.msgs, nil
}
func (fs *FileStorage) Save(messages []Message) error {
fs.mu.Lock()
defer fs.mu.Unlock()
fs.msgs = messages
data, err := json.Marshal(messages)
if err != nil {
return err
}
return os.WriteFile(fs.path, data, 0o644)
}
func (fs *FileStorage) Add(msg *cvms.ClientStreamMessage) error {
fs.mu.Lock()
defer fs.mu.Unlock()
fs.msgs = append(fs.msgs, Message{
Message: msg,
Time: time.Now(),
})
data, err := json.Marshal(fs.msgs)
if err != nil {
return err
}
return os.WriteFile(fs.path, data, 0o644)
}
func (fs *FileStorage) Clear() error {
fs.mu.Lock()
defer fs.mu.Unlock()
fs.msgs = make([]Message, 0)
return os.WriteFile(fs.path, []byte("[]"), 0o644)
}
File diff suppressed because it is too large Load Diff
-130
View File
@@ -1,130 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
syntax = "proto3";
import "google/protobuf/timestamp.proto";
package cvms;
option go_package = "./cvms";
service Service {
rpc Process(stream ClientStreamMessage) returns (stream ServerStreamMessage) {}
}
message AgentStateReq {
string id = 1;
}
message AgentStateRes {
string id = 1;
string state = 2;
}
message StopComputation {
string computation_id = 1;
}
message StopComputationResponse {
string computation_id = 1;
string message = 2;
}
message RunResponse{
string computation_id = 1;
string error = 2;
}
message AgentEvent {
string event_type = 1;
google.protobuf.Timestamp timestamp = 2;
string computation_id = 3;
bytes details = 4;
string originator = 5;
string status = 6;
}
message AgentLog {
string message = 1;
string computation_id = 2;
string level = 3;
google.protobuf.Timestamp timestamp = 4;
}
message ClientStreamMessage {
oneof message {
AgentLog agent_log = 1;
AgentEvent agent_event = 2;
RunResponse run_res = 3;
StopComputationResponse stopComputationRes = 4;
AgentStateRes agentStateRes = 5;
AttestationResponse vTPMattestationReport = 6;
azureAttestationResponse azureAttestationResult = 7;
}
}
message ServerStreamMessage {
oneof message {
RunReqChunks runReqChunks = 1;
ComputationRunReq runReq = 2;
StopComputation stopComputation = 3;
AgentStateReq agentStateReq = 4;
DisconnectReq disconnectReq = 5;
}
}
message DisconnectReq {
string id = 1;
}
message RunReqChunks {
bytes data = 1;
string id = 2;
bool is_last = 3;
}
message ComputationRunReq {
string id = 1;
string name = 2;
string description = 3;
repeated Dataset datasets = 4;
Algorithm algorithm = 5;
repeated ResultConsumer result_consumers = 6;
AgentConfig agent_config = 7;
}
message ResultConsumer {
bytes userKey = 1;
}
message Dataset {
bytes hash = 1; // should be sha3.Sum256, 32 byte length.
bytes userKey = 2;
string filename = 3;
}
message Algorithm {
bytes hash = 1; // should be sha3.Sum256, 32 byte length.
bytes userKey = 2;
}
message AgentConfig {
string port = 1;
string cert_file = 2;
string key_file = 3;
string client_ca_file = 4;
string server_ca_file = 5;
string log_level = 6;
bool attested_tls = 7;
}
message AttestationResponse {
bytes file = 1;
string certSerialNumber = 2;
}
message azureAttestationResponse {
bytes file = 1;
string certSerialNumber = 2;
}
-118
View File
@@ -1,118 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
// versions:
// - protoc-gen-go-grpc v1.5.1
// - protoc v5.29.0
// source: agent/cvms/cvms.proto
package cvms
import (
context "context"
grpc "google.golang.org/grpc"
codes "google.golang.org/grpc/codes"
status "google.golang.org/grpc/status"
)
// This is a compile-time assertion to ensure that this generated file
// is compatible with the grpc package it is being compiled against.
// Requires gRPC-Go v1.64.0 or later.
const _ = grpc.SupportPackageIsVersion9
const (
Service_Process_FullMethodName = "/cvms.Service/Process"
)
// ServiceClient is the client API for Service service.
//
// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
type ServiceClient interface {
Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage], error)
}
type serviceClient struct {
cc grpc.ClientConnInterface
}
func NewServiceClient(cc grpc.ClientConnInterface) ServiceClient {
return &serviceClient{cc}
}
func (c *serviceClient) Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &Service_ServiceDesc.Streams[0], Service_Process_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[ClientStreamMessage, ServerStreamMessage]{ClientStream: stream}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Service_ProcessClient = grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage]
// ServiceServer is the server API for Service service.
// All implementations must embed UnimplementedServiceServer
// for forward compatibility.
type ServiceServer interface {
Process(grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]) error
mustEmbedUnimplementedServiceServer()
}
// UnimplementedServiceServer must be embedded to have
// forward compatible implementations.
//
// NOTE: this should be embedded by value instead of pointer to avoid a nil
// pointer dereference when methods are called.
type UnimplementedServiceServer struct{}
func (UnimplementedServiceServer) Process(grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]) error {
return status.Errorf(codes.Unimplemented, "method Process not implemented")
}
func (UnimplementedServiceServer) mustEmbedUnimplementedServiceServer() {}
func (UnimplementedServiceServer) testEmbeddedByValue() {}
// UnsafeServiceServer may be embedded to opt out of forward compatibility for this service.
// Use of this interface is not recommended, as added methods to ServiceServer will
// result in compilation errors.
type UnsafeServiceServer interface {
mustEmbedUnimplementedServiceServer()
}
func RegisterServiceServer(s grpc.ServiceRegistrar, srv ServiceServer) {
// If the following call pancis, it indicates UnimplementedServiceServer was
// embedded by pointer and is nil. This will cause panics if an
// unimplemented method is ever invoked, so we test this at initialization
// time to prevent it from happening at runtime later due to I/O.
if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
t.testEmbeddedByValue()
}
s.RegisterService(&Service_ServiceDesc, srv)
}
func _Service_Process_Handler(srv interface{}, stream grpc.ServerStream) error {
return srv.(ServiceServer).Process(&grpc.GenericServerStream[ClientStreamMessage, ServerStreamMessage]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Service_ProcessServer = grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]
// Service_ServiceDesc is the grpc.ServiceDesc for Service service.
// It's only intended for direct use with grpc.RegisterService,
// and not to be introspected or modified (even as a copy)
var Service_ServiceDesc = grpc.ServiceDesc{
ServiceName: "cvms.Service",
HandlerType: (*ServiceServer)(nil),
Methods: []grpc.MethodDesc{},
Streams: []grpc.StreamDesc{
{
StreamName: "Process",
Handler: _Service_Process_Handler,
ServerStreams: true,
ClientStreams: true,
},
},
Metadata: "agent/cvms/cvms.proto",
}
-98
View File
@@ -1,98 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package server
import (
context "context"
"fmt"
"log/slog"
"github.com/ultravioletrs/cocos/agent"
agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
"github.com/ultravioletrs/cocos/agent/auth"
"github.com/ultravioletrs/cocos/internal/server"
grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
"google.golang.org/grpc"
"google.golang.org/grpc/reflection"
)
const (
svcName = "agent"
defSvcGRPCPort = "7002"
)
type AgentServer interface {
Start(cfg agent.AgentConfig, cmp agent.Computation) error
Stop() error
}
type agentServer struct {
gs server.Server
logger *slog.Logger
svc agent.Service
host string
caUrl string
cvmId string
}
func NewServer(logger *slog.Logger, svc agent.Service, host string, caUrl string, cvmId string) AgentServer {
return &agentServer{
logger: logger,
svc: svc,
host: host,
caUrl: caUrl,
cvmId: cvmId,
}
}
func (as *agentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
if cfg.Port == "" {
cfg.Port = defSvcGRPCPort
}
agentGrpcServerConfig := server.AgentConfig{
ServerConfig: server.ServerConfig{
BaseConfig: server.BaseConfig{
Host: as.host,
Port: cfg.Port,
CertFile: cfg.CertFile,
KeyFile: cfg.KeyFile,
ServerCAFile: cfg.ServerCAFile,
ClientCAFile: cfg.ClientCAFile,
},
},
AttestedTLS: cfg.AttestedTls,
}
registerAgentServiceServer := func(srv *grpc.Server) {
reflection.Register(srv)
agent.RegisterAgentServiceServer(srv, agentgrpc.NewServer(as.svc))
}
authSvc, err := auth.New(cmp)
if err != nil {
as.logger.WithGroup(cmp.ID).Error(fmt.Sprintf("failed to create auth service %s", err.Error()))
return err
}
ctx, cancel := context.WithCancel(context.Background())
as.gs = grpcserver.New(ctx, cancel, svcName, agentGrpcServerConfig, registerAgentServiceServer, as.logger, authSvc, as.caUrl, as.cvmId)
go func() {
err := as.gs.Start()
if err != nil {
as.logger.Error(fmt.Sprintf("failed to start grpc server %s", err.Error()))
}
}()
return nil
}
func (as *agentServer) Stop() error {
if as.gs == nil {
return nil
}
return as.gs.Stop()
}
-130
View File
@@ -1,130 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mocks
import (
mock "github.com/stretchr/testify/mock"
agent "github.com/ultravioletrs/cocos/agent"
)
// AgentServer is an autogenerated mock type for the AgentServer type
type AgentServer struct {
mock.Mock
}
type AgentServer_Expecter struct {
mock *mock.Mock
}
func (_m *AgentServer) EXPECT() *AgentServer_Expecter {
return &AgentServer_Expecter{mock: &_m.Mock}
}
// Start provides a mock function with given fields: cfg, cmp
func (_m *AgentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
ret := _m.Called(cfg, cmp)
if len(ret) == 0 {
panic("no return value specified for Start")
}
var r0 error
if rf, ok := ret.Get(0).(func(agent.AgentConfig, agent.Computation) error); ok {
r0 = rf(cfg, cmp)
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentServer_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
type AgentServer_Start_Call struct {
*mock.Call
}
// Start is a helper method to define mock.On call
// - cfg agent.AgentConfig
// - cmp agent.Computation
func (_e *AgentServer_Expecter) Start(cfg interface{}, cmp interface{}) *AgentServer_Start_Call {
return &AgentServer_Start_Call{Call: _e.mock.On("Start", cfg, cmp)}
}
func (_c *AgentServer_Start_Call) Run(run func(cfg agent.AgentConfig, cmp agent.Computation)) *AgentServer_Start_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(agent.AgentConfig), args[1].(agent.Computation))
})
return _c
}
func (_c *AgentServer_Start_Call) Return(_a0 error) *AgentServer_Start_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentServer_Start_Call) RunAndReturn(run func(agent.AgentConfig, agent.Computation) error) *AgentServer_Start_Call {
_c.Call.Return(run)
return _c
}
// Stop provides a mock function with no fields
func (_m *AgentServer) Stop() error {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Stop")
}
var r0 error
if rf, ok := ret.Get(0).(func() error); ok {
r0 = rf()
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentServer_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
type AgentServer_Stop_Call struct {
*mock.Call
}
// Stop is a helper method to define mock.On call
func (_e *AgentServer_Expecter) Stop() *AgentServer_Stop_Call {
return &AgentServer_Stop_Call{Call: _e.mock.On("Stop")}
}
func (_c *AgentServer_Stop_Call) Run(run func()) *AgentServer_Stop_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentServer_Stop_Call) Return(_a0 error) *AgentServer_Stop_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentServer_Stop_Call) RunAndReturn(run func() error) *AgentServer_Stop_Call {
_c.Call.Return(run)
return _c
}
// NewAgentServer creates a new instance of AgentServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewAgentServer(t interface {
mock.TestingT
Cleanup(func())
}) *AgentServer {
mock := &AgentServer{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
+46 -19
View File
@@ -4,38 +4,65 @@ package events
import (
"encoding/json"
"time"
"github.com/ultravioletrs/cocos/agent/cvms"
"github.com/mdlayher/vsock"
"github.com/ultravioletrs/cocos/pkg/manager"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/timestamppb"
)
type service struct {
service string
queue chan *cvms.ClientStreamMessage
service string
computationID string
conn *vsock.Conn
}
type AgentEvent struct {
EventType string `json:"event_type"`
Timestamp time.Time `json:"timestamp"`
ComputationID string `json:"computation_id,omitempty"`
Details json.RawMessage `json:"details,omitempty"`
Originator string `json:"originator"`
Status string `json:"status,omitempty"`
}
type Service interface {
SendEvent(cmpID, event, status string, details json.RawMessage)
SendEvent(event, status string, details json.RawMessage) error
Close() error
}
func New(svc string, queue chan *cvms.ClientStreamMessage) (Service, error) {
func New(svc, computationID string, sockPort uint32) (Service, error) {
conn, err := vsock.Dial(vsock.Host, sockPort, nil)
if err != nil {
return nil, err
}
return &service{
service: svc,
queue: queue,
service: svc,
computationID: computationID,
conn: conn,
}, nil
}
func (s *service) SendEvent(cmpID, event, status string, details json.RawMessage) {
s.queue <- &cvms.ClientStreamMessage{
Message: &cvms.ClientStreamMessage_AgentEvent{
AgentEvent: &cvms.AgentEvent{
EventType: event,
Timestamp: timestamppb.Now(),
ComputationId: cmpID,
Originator: s.service,
Status: status,
Details: details,
},
},
func (s *service) SendEvent(event, status string, details json.RawMessage) error {
body := manager.ClientStreamMessage{Message: &manager.ClientStreamMessage_AgentEvent{AgentEvent: &manager.AgentEvent{
EventType: event,
Timestamp: timestamppb.Now(),
ComputationId: s.computationID,
Originator: s.service,
Status: status,
Details: details,
}}}
protoBody, err := proto.Marshal(&body)
if err != nil {
return err
}
if _, err := s.conn.Write(protoBody); err != nil {
return err
}
return nil
}
func (s *service) Close() error {
return s.conn.Close()
}
-362
View File
@@ -1,362 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.5
// protoc v5.29.0
// source: agent/events/events.proto
package events
import (
protoreflect "google.golang.org/protobuf/reflect/protoreflect"
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
timestamppb "google.golang.org/protobuf/types/known/timestamppb"
reflect "reflect"
sync "sync"
unsafe "unsafe"
)
const (
// Verify that this generated code is sufficiently up-to-date.
_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
// Verify that runtime/protoimpl is sufficiently up-to-date.
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
)
type AgentEvent struct {
state protoimpl.MessageState `protogen:"open.v1"`
EventType string `protobuf:"bytes,1,opt,name=event_type,json=eventType,proto3" json:"event_type,omitempty"`
Timestamp *timestamppb.Timestamp `protobuf:"bytes,2,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
ComputationId string `protobuf:"bytes,3,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
Details []byte `protobuf:"bytes,4,opt,name=details,proto3" json:"details,omitempty"`
Originator string `protobuf:"bytes,5,opt,name=originator,proto3" json:"originator,omitempty"`
Status string `protobuf:"bytes,6,opt,name=status,proto3" json:"status,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *AgentEvent) Reset() {
*x = AgentEvent{}
mi := &file_agent_events_events_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *AgentEvent) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*AgentEvent) ProtoMessage() {}
func (x *AgentEvent) ProtoReflect() protoreflect.Message {
mi := &file_agent_events_events_proto_msgTypes[0]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use AgentEvent.ProtoReflect.Descriptor instead.
func (*AgentEvent) Descriptor() ([]byte, []int) {
return file_agent_events_events_proto_rawDescGZIP(), []int{0}
}
func (x *AgentEvent) GetEventType() string {
if x != nil {
return x.EventType
}
return ""
}
func (x *AgentEvent) GetTimestamp() *timestamppb.Timestamp {
if x != nil {
return x.Timestamp
}
return nil
}
func (x *AgentEvent) GetComputationId() string {
if x != nil {
return x.ComputationId
}
return ""
}
func (x *AgentEvent) GetDetails() []byte {
if x != nil {
return x.Details
}
return nil
}
func (x *AgentEvent) GetOriginator() string {
if x != nil {
return x.Originator
}
return ""
}
func (x *AgentEvent) GetStatus() string {
if x != nil {
return x.Status
}
return ""
}
type AgentLog struct {
state protoimpl.MessageState `protogen:"open.v1"`
Message string `protobuf:"bytes,1,opt,name=message,proto3" json:"message,omitempty"`
ComputationId string `protobuf:"bytes,2,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
Level string `protobuf:"bytes,3,opt,name=level,proto3" json:"level,omitempty"`
Timestamp *timestamppb.Timestamp `protobuf:"bytes,4,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *AgentLog) Reset() {
*x = AgentLog{}
mi := &file_agent_events_events_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *AgentLog) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*AgentLog) ProtoMessage() {}
func (x *AgentLog) ProtoReflect() protoreflect.Message {
mi := &file_agent_events_events_proto_msgTypes[1]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use AgentLog.ProtoReflect.Descriptor instead.
func (*AgentLog) Descriptor() ([]byte, []int) {
return file_agent_events_events_proto_rawDescGZIP(), []int{1}
}
func (x *AgentLog) GetMessage() string {
if x != nil {
return x.Message
}
return ""
}
func (x *AgentLog) GetComputationId() string {
if x != nil {
return x.ComputationId
}
return ""
}
func (x *AgentLog) GetLevel() string {
if x != nil {
return x.Level
}
return ""
}
func (x *AgentLog) GetTimestamp() *timestamppb.Timestamp {
if x != nil {
return x.Timestamp
}
return nil
}
type EventsLogs struct {
state protoimpl.MessageState `protogen:"open.v1"`
// Types that are valid to be assigned to Message:
//
// *EventsLogs_AgentLog
// *EventsLogs_AgentEvent
Message isEventsLogs_Message `protobuf_oneof:"message"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *EventsLogs) Reset() {
*x = EventsLogs{}
mi := &file_agent_events_events_proto_msgTypes[2]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *EventsLogs) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*EventsLogs) ProtoMessage() {}
func (x *EventsLogs) ProtoReflect() protoreflect.Message {
mi := &file_agent_events_events_proto_msgTypes[2]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use EventsLogs.ProtoReflect.Descriptor instead.
func (*EventsLogs) Descriptor() ([]byte, []int) {
return file_agent_events_events_proto_rawDescGZIP(), []int{2}
}
func (x *EventsLogs) GetMessage() isEventsLogs_Message {
if x != nil {
return x.Message
}
return nil
}
func (x *EventsLogs) GetAgentLog() *AgentLog {
if x != nil {
if x, ok := x.Message.(*EventsLogs_AgentLog); ok {
return x.AgentLog
}
}
return nil
}
func (x *EventsLogs) GetAgentEvent() *AgentEvent {
if x != nil {
if x, ok := x.Message.(*EventsLogs_AgentEvent); ok {
return x.AgentEvent
}
}
return nil
}
type isEventsLogs_Message interface {
isEventsLogs_Message()
}
type EventsLogs_AgentLog struct {
AgentLog *AgentLog `protobuf:"bytes,1,opt,name=agent_log,json=agentLog,proto3,oneof"`
}
type EventsLogs_AgentEvent struct {
AgentEvent *AgentEvent `protobuf:"bytes,2,opt,name=agent_event,json=agentEvent,proto3,oneof"`
}
func (*EventsLogs_AgentLog) isEventsLogs_Message() {}
func (*EventsLogs_AgentEvent) isEventsLogs_Message() {}
var File_agent_events_events_proto protoreflect.FileDescriptor
var file_agent_events_events_proto_rawDesc = string([]byte{
0x0a, 0x19, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2f, 0x65,
0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x06, 0x65, 0x76, 0x65,
0x6e, 0x74, 0x73, 0x1a, 0x1f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74,
0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70,
0x72, 0x6f, 0x74, 0x6f, 0x22, 0xde, 0x01, 0x0a, 0x0a, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76,
0x65, 0x6e, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x5f, 0x74, 0x79, 0x70,
0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x54, 0x79,
0x70, 0x65, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18,
0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70,
0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d,
0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x12, 0x25, 0x0a, 0x0e,
0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x03,
0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f,
0x6e, 0x49, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x73, 0x18, 0x04,
0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x73, 0x12, 0x1e, 0x0a,
0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x18, 0x05, 0x20, 0x01, 0x28,
0x09, 0x52, 0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x16, 0x0a,
0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x73,
0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x9b, 0x01, 0x0a, 0x08, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x4c,
0x6f, 0x67, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, 0x01, 0x20,
0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x25, 0x0a, 0x0e,
0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x02,
0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f,
0x6e, 0x49, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x18, 0x03, 0x20, 0x01,
0x28, 0x09, 0x52, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d,
0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67,
0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54,
0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74,
0x61, 0x6d, 0x70, 0x22, 0x7f, 0x0a, 0x0a, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x4c, 0x6f, 0x67,
0x73, 0x12, 0x2f, 0x0a, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x18, 0x01,
0x20, 0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x41, 0x67,
0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x48, 0x00, 0x52, 0x08, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x4c,
0x6f, 0x67, 0x12, 0x35, 0x0a, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x65, 0x76, 0x65, 0x6e,
0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73,
0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x48, 0x00, 0x52, 0x0a, 0x61,
0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73,
0x73, 0x61, 0x67, 0x65, 0x42, 0x0a, 0x5a, 0x08, 0x2e, 0x2f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73,
0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
})
var (
file_agent_events_events_proto_rawDescOnce sync.Once
file_agent_events_events_proto_rawDescData []byte
)
func file_agent_events_events_proto_rawDescGZIP() []byte {
file_agent_events_events_proto_rawDescOnce.Do(func() {
file_agent_events_events_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_events_events_proto_rawDesc), len(file_agent_events_events_proto_rawDesc)))
})
return file_agent_events_events_proto_rawDescData
}
var file_agent_events_events_proto_msgTypes = make([]protoimpl.MessageInfo, 3)
var file_agent_events_events_proto_goTypes = []any{
(*AgentEvent)(nil), // 0: events.AgentEvent
(*AgentLog)(nil), // 1: events.AgentLog
(*EventsLogs)(nil), // 2: events.EventsLogs
(*timestamppb.Timestamp)(nil), // 3: google.protobuf.Timestamp
}
var file_agent_events_events_proto_depIdxs = []int32{
3, // 0: events.AgentEvent.timestamp:type_name -> google.protobuf.Timestamp
3, // 1: events.AgentLog.timestamp:type_name -> google.protobuf.Timestamp
1, // 2: events.EventsLogs.agent_log:type_name -> events.AgentLog
0, // 3: events.EventsLogs.agent_event:type_name -> events.AgentEvent
4, // [4:4] is the sub-list for method output_type
4, // [4:4] is the sub-list for method input_type
4, // [4:4] is the sub-list for extension type_name
4, // [4:4] is the sub-list for extension extendee
0, // [0:4] is the sub-list for field type_name
}
func init() { file_agent_events_events_proto_init() }
func file_agent_events_events_proto_init() {
if File_agent_events_events_proto != nil {
return
}
file_agent_events_events_proto_msgTypes[2].OneofWrappers = []any{
(*EventsLogs_AgentLog)(nil),
(*EventsLogs_AgentEvent)(nil),
}
type x struct{}
out := protoimpl.TypeBuilder{
File: protoimpl.DescBuilder{
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_events_events_proto_rawDesc), len(file_agent_events_events_proto_rawDesc)),
NumEnums: 0,
NumMessages: 3,
NumExtensions: 0,
NumServices: 0,
},
GoTypes: file_agent_events_events_proto_goTypes,
DependencyIndexes: file_agent_events_events_proto_depIdxs,
MessageInfos: file_agent_events_events_proto_msgTypes,
}.Build()
File_agent_events_events_proto = out.File
file_agent_events_events_proto_goTypes = nil
file_agent_events_events_proto_depIdxs = nil
}
-33
View File
@@ -1,33 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
syntax = "proto3";
package events;
import "google/protobuf/timestamp.proto";
option go_package = "./events";
message AgentEvent {
string event_type = 1;
google.protobuf.Timestamp timestamp = 2;
string computation_id = 3;
bytes details = 4;
string originator = 5;
string status = 6;
}
message AgentLog {
string message = 1;
string computation_id = 2;
string level = 3;
google.protobuf.Timestamp timestamp = 4;
}
message EventsLogs {
oneof message {
AgentLog agent_log = 1;
AgentEvent agent_event = 2;
}
}
-38
View File
@@ -1,38 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package events
import (
"encoding/json"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/ultravioletrs/cocos/agent/cvms"
)
func TestSendEventSuccess(t *testing.T) {
queue := make(chan *cvms.ClientStreamMessage, 1)
svc, err := New("test_service", queue)
assert.NoError(t, err)
details := json.RawMessage(`{"key": "value"}`)
go func() {
msg := <-queue
assert.NotNil(t, msg)
assert.NotNil(t, msg.GetAgentEvent())
assert.Equal(t, "test_event", msg.GetAgentEvent().EventType)
assert.Equal(t, "testid", msg.GetAgentEvent().ComputationId)
assert.Equal(t, "test_service", msg.GetAgentEvent().Originator)
assert.Equal(t, "success", msg.GetAgentEvent().Status)
now := time.Now()
eventTimestamp := msg.GetAgentEvent().GetTimestamp().AsTime()
assert.WithinDuration(t, now, eventTimestamp, 1*time.Second)
}()
svc.SendEvent("testid", "test_event", "success", details)
time.Sleep(1 * time.Second)
}
-75
View File
@@ -1,75 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mocks
import (
json "encoding/json"
mock "github.com/stretchr/testify/mock"
)
// Service is an autogenerated mock type for the Service type
type Service struct {
mock.Mock
}
type Service_Expecter struct {
mock *mock.Mock
}
func (_m *Service) EXPECT() *Service_Expecter {
return &Service_Expecter{mock: &_m.Mock}
}
// SendEvent provides a mock function with given fields: cmpID, event, status, details
func (_m *Service) SendEvent(cmpID string, event string, status string, details json.RawMessage) {
_m.Called(cmpID, event, status, details)
}
// Service_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
type Service_SendEvent_Call struct {
*mock.Call
}
// SendEvent is a helper method to define mock.On call
// - cmpID string
// - event string
// - status string
// - details json.RawMessage
func (_e *Service_Expecter) SendEvent(cmpID interface{}, event interface{}, status interface{}, details interface{}) *Service_SendEvent_Call {
return &Service_SendEvent_Call{Call: _e.mock.On("SendEvent", cmpID, event, status, details)}
}
func (_c *Service_SendEvent_Call) Run(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(string), args[1].(string), args[2].(string), args[3].(json.RawMessage))
})
return _c
}
func (_c *Service_SendEvent_Call) Return() *Service_SendEvent_Call {
_c.Call.Return()
return _c
}
func (_c *Service_SendEvent_Call) RunAndReturn(run func(string, string, string, json.RawMessage)) *Service_SendEvent_Call {
_c.Run(run)
return _c
}
// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewService(t interface {
mock.TestingT
Cleanup(func())
}) *Service {
mock := &Service{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
-520
View File
@@ -1,520 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mocks
import (
agent "github.com/ultravioletrs/cocos/agent"
attestation "github.com/ultravioletrs/cocos/pkg/attestation"
context "context"
mock "github.com/stretchr/testify/mock"
)
// Service is an autogenerated mock type for the Service type
type Service struct {
mock.Mock
}
type Service_Expecter struct {
mock *mock.Mock
}
func (_m *Service) EXPECT() *Service_Expecter {
return &Service_Expecter{mock: &_m.Mock}
}
// Algo provides a mock function with given fields: ctx, algorithm
func (_m *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
ret := _m.Called(ctx, algorithm)
if len(ret) == 0 {
panic("no return value specified for Algo")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, agent.Algorithm) error); ok {
r0 = rf(ctx, algorithm)
} else {
r0 = ret.Error(0)
}
return r0
}
// Service_Algo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Algo'
type Service_Algo_Call struct {
*mock.Call
}
// Algo is a helper method to define mock.On call
// - ctx context.Context
// - algorithm agent.Algorithm
func (_e *Service_Expecter) Algo(ctx interface{}, algorithm interface{}) *Service_Algo_Call {
return &Service_Algo_Call{Call: _e.mock.On("Algo", ctx, algorithm)}
}
func (_c *Service_Algo_Call) Run(run func(ctx context.Context, algorithm agent.Algorithm)) *Service_Algo_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(agent.Algorithm))
})
return _c
}
func (_c *Service_Algo_Call) Return(_a0 error) *Service_Algo_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *Service_Algo_Call) RunAndReturn(run func(context.Context, agent.Algorithm) error) *Service_Algo_Call {
_c.Call.Return(run)
return _c
}
// Attestation provides a mock function with given fields: ctx, reportData, nonce, attType
func (_m *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
ret := _m.Called(ctx, reportData, nonce, attType)
if len(ret) == 0 {
panic("no return value specified for Attestation")
}
var r0 []byte
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)); ok {
return rf(ctx, reportData, nonce, attType)
}
if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) []byte); ok {
r0 = rf(ctx, reportData, nonce, attType)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]byte)
}
}
if rf, ok := ret.Get(1).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) error); ok {
r1 = rf(ctx, reportData, nonce, attType)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// Service_Attestation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Attestation'
type Service_Attestation_Call struct {
*mock.Call
}
// Attestation is a helper method to define mock.On call
// - ctx context.Context
// - reportData [64]byte
// - nonce [32]byte
// - attType attestation.PlatformType
func (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{}, nonce interface{}, attType interface{}) *Service_Attestation_Call {
return &Service_Attestation_Call{Call: _e.mock.On("Attestation", ctx, reportData, nonce, attType)}
}
func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType)) *Service_Attestation_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([64]byte), args[2].([32]byte), args[3].(attestation.PlatformType))
})
return _c
}
func (_c *Service_Attestation_Call) Return(_a0 []byte, _a1 error) *Service_Attestation_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *Service_Attestation_Call) RunAndReturn(run func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)) *Service_Attestation_Call {
_c.Call.Return(run)
return _c
}
// AttestationResult provides a mock function with given fields: ctx, nonce, attType
func (_m *Service) AttestationResult(ctx context.Context, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
ret := _m.Called(ctx, nonce, attType)
if len(ret) == 0 {
panic("no return value specified for AttestationResult")
}
var r0 []byte
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, [32]byte, attestation.PlatformType) ([]byte, error)); ok {
return rf(ctx, nonce, attType)
}
if rf, ok := ret.Get(0).(func(context.Context, [32]byte, attestation.PlatformType) []byte); ok {
r0 = rf(ctx, nonce, attType)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]byte)
}
}
if rf, ok := ret.Get(1).(func(context.Context, [32]byte, attestation.PlatformType) error); ok {
r1 = rf(ctx, nonce, attType)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// Service_AttestationResult_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AttestationResult'
type Service_AttestationResult_Call struct {
*mock.Call
}
// AttestationResult is a helper method to define mock.On call
// - ctx context.Context
// - nonce [32]byte
// - attType attestation.PlatformType
func (_e *Service_Expecter) AttestationResult(ctx interface{}, nonce interface{}, attType interface{}) *Service_AttestationResult_Call {
return &Service_AttestationResult_Call{Call: _e.mock.On("AttestationResult", ctx, nonce, attType)}
}
func (_c *Service_AttestationResult_Call) Run(run func(ctx context.Context, nonce [32]byte, attType attestation.PlatformType)) *Service_AttestationResult_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([32]byte), args[2].(attestation.PlatformType))
})
return _c
}
func (_c *Service_AttestationResult_Call) Return(_a0 []byte, _a1 error) *Service_AttestationResult_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *Service_AttestationResult_Call) RunAndReturn(run func(context.Context, [32]byte, attestation.PlatformType) ([]byte, error)) *Service_AttestationResult_Call {
_c.Call.Return(run)
return _c
}
// Data provides a mock function with given fields: ctx, dataset
func (_m *Service) Data(ctx context.Context, dataset agent.Dataset) error {
ret := _m.Called(ctx, dataset)
if len(ret) == 0 {
panic("no return value specified for Data")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, agent.Dataset) error); ok {
r0 = rf(ctx, dataset)
} else {
r0 = ret.Error(0)
}
return r0
}
// Service_Data_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Data'
type Service_Data_Call struct {
*mock.Call
}
// Data is a helper method to define mock.On call
// - ctx context.Context
// - dataset agent.Dataset
func (_e *Service_Expecter) Data(ctx interface{}, dataset interface{}) *Service_Data_Call {
return &Service_Data_Call{Call: _e.mock.On("Data", ctx, dataset)}
}
func (_c *Service_Data_Call) Run(run func(ctx context.Context, dataset agent.Dataset)) *Service_Data_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(agent.Dataset))
})
return _c
}
func (_c *Service_Data_Call) Return(_a0 error) *Service_Data_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *Service_Data_Call) RunAndReturn(run func(context.Context, agent.Dataset) error) *Service_Data_Call {
_c.Call.Return(run)
return _c
}
// IMAMeasurements provides a mock function with given fields: ctx
func (_m *Service) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for IMAMeasurements")
}
var r0 []byte
var r1 []byte
var r2 error
if rf, ok := ret.Get(0).(func(context.Context) ([]byte, []byte, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) []byte); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]byte)
}
}
if rf, ok := ret.Get(1).(func(context.Context) []byte); ok {
r1 = rf(ctx)
} else {
if ret.Get(1) != nil {
r1 = ret.Get(1).([]byte)
}
}
if rf, ok := ret.Get(2).(func(context.Context) error); ok {
r2 = rf(ctx)
} else {
r2 = ret.Error(2)
}
return r0, r1, r2
}
// Service_IMAMeasurements_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IMAMeasurements'
type Service_IMAMeasurements_Call struct {
*mock.Call
}
// IMAMeasurements is a helper method to define mock.On call
// - ctx context.Context
func (_e *Service_Expecter) IMAMeasurements(ctx interface{}) *Service_IMAMeasurements_Call {
return &Service_IMAMeasurements_Call{Call: _e.mock.On("IMAMeasurements", ctx)}
}
func (_c *Service_IMAMeasurements_Call) Run(run func(ctx context.Context)) *Service_IMAMeasurements_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *Service_IMAMeasurements_Call) Return(_a0 []byte, _a1 []byte, _a2 error) *Service_IMAMeasurements_Call {
_c.Call.Return(_a0, _a1, _a2)
return _c
}
func (_c *Service_IMAMeasurements_Call) RunAndReturn(run func(context.Context) ([]byte, []byte, error)) *Service_IMAMeasurements_Call {
_c.Call.Return(run)
return _c
}
// InitComputation provides a mock function with given fields: ctx, cmp
func (_m *Service) InitComputation(ctx context.Context, cmp agent.Computation) error {
ret := _m.Called(ctx, cmp)
if len(ret) == 0 {
panic("no return value specified for InitComputation")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, agent.Computation) error); ok {
r0 = rf(ctx, cmp)
} else {
r0 = ret.Error(0)
}
return r0
}
// Service_InitComputation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'InitComputation'
type Service_InitComputation_Call struct {
*mock.Call
}
// InitComputation is a helper method to define mock.On call
// - ctx context.Context
// - cmp agent.Computation
func (_e *Service_Expecter) InitComputation(ctx interface{}, cmp interface{}) *Service_InitComputation_Call {
return &Service_InitComputation_Call{Call: _e.mock.On("InitComputation", ctx, cmp)}
}
func (_c *Service_InitComputation_Call) Run(run func(ctx context.Context, cmp agent.Computation)) *Service_InitComputation_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(agent.Computation))
})
return _c
}
func (_c *Service_InitComputation_Call) Return(_a0 error) *Service_InitComputation_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *Service_InitComputation_Call) RunAndReturn(run func(context.Context, agent.Computation) error) *Service_InitComputation_Call {
_c.Call.Return(run)
return _c
}
// Result provides a mock function with given fields: ctx
func (_m *Service) Result(ctx context.Context) ([]byte, error) {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for Result")
}
var r0 []byte
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) []byte); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]byte)
}
}
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// Service_Result_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Result'
type Service_Result_Call struct {
*mock.Call
}
// Result is a helper method to define mock.On call
// - ctx context.Context
func (_e *Service_Expecter) Result(ctx interface{}) *Service_Result_Call {
return &Service_Result_Call{Call: _e.mock.On("Result", ctx)}
}
func (_c *Service_Result_Call) Run(run func(ctx context.Context)) *Service_Result_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *Service_Result_Call) Return(_a0 []byte, _a1 error) *Service_Result_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *Service_Result_Call) RunAndReturn(run func(context.Context) ([]byte, error)) *Service_Result_Call {
_c.Call.Return(run)
return _c
}
// State provides a mock function with no fields
func (_m *Service) State() string {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for State")
}
var r0 string
if rf, ok := ret.Get(0).(func() string); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(string)
}
return r0
}
// Service_State_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'State'
type Service_State_Call struct {
*mock.Call
}
// State is a helper method to define mock.On call
func (_e *Service_Expecter) State() *Service_State_Call {
return &Service_State_Call{Call: _e.mock.On("State")}
}
func (_c *Service_State_Call) Run(run func()) *Service_State_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *Service_State_Call) Return(_a0 string) *Service_State_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *Service_State_Call) RunAndReturn(run func() string) *Service_State_Call {
_c.Call.Return(run)
return _c
}
// StopComputation provides a mock function with given fields: ctx
func (_m *Service) StopComputation(ctx context.Context) error {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for StopComputation")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context) error); ok {
r0 = rf(ctx)
} else {
r0 = ret.Error(0)
}
return r0
}
// Service_StopComputation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'StopComputation'
type Service_StopComputation_Call struct {
*mock.Call
}
// StopComputation is a helper method to define mock.On call
// - ctx context.Context
func (_e *Service_Expecter) StopComputation(ctx interface{}) *Service_StopComputation_Call {
return &Service_StopComputation_Call{Call: _e.mock.On("StopComputation", ctx)}
}
func (_c *Service_StopComputation_Call) Run(run func(ctx context.Context)) *Service_StopComputation_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *Service_StopComputation_Call) Return(_a0 error) *Service_StopComputation_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *Service_StopComputation_Call) RunAndReturn(run func(context.Context) error) *Service_StopComputation_Call {
_c.Call.Return(run)
return _c
}
// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewService(t interface {
mock.TestingT
Cleanup(func())
}) *Service {
mock := &Service{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
-434
View File
@@ -1,434 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mocks
import (
context "context"
agent "github.com/ultravioletrs/cocos/agent"
metadata "google.golang.org/grpc/metadata"
mock "github.com/stretchr/testify/mock"
)
// AgentService_AlgoClient is an autogenerated mock type for the AgentService_AlgoClient type
type AgentService_AlgoClient[Req interface{}, Res interface{}] struct {
mock.Mock
}
type AgentService_AlgoClient_Expecter[Req interface{}, Res interface{}] struct {
mock *mock.Mock
}
func (_m *AgentService_AlgoClient[Req, Res]) EXPECT() *AgentService_AlgoClient_Expecter[Req, Res] {
return &AgentService_AlgoClient_Expecter[Req, Res]{mock: &_m.Mock}
}
// CloseAndRecv provides a mock function with no fields
func (_m *AgentService_AlgoClient[Req, Res]) CloseAndRecv() (*agent.AlgoResponse, error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for CloseAndRecv")
}
var r0 *agent.AlgoResponse
var r1 error
if rf, ok := ret.Get(0).(func() (*agent.AlgoResponse, error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() *agent.AlgoResponse); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*agent.AlgoResponse)
}
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// AgentService_AlgoClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
type AgentService_AlgoClient_CloseAndRecv_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// CloseAndRecv is a helper method to define mock.On call
func (_e *AgentService_AlgoClient_Expecter[Req, Res]) CloseAndRecv() *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
return &AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]{Call: _e.mock.On("CloseAndRecv")}
}
func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) Return(_a0 *agent.AlgoResponse, _a1 error) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) RunAndReturn(run func() (*agent.AlgoResponse, error)) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// CloseSend provides a mock function with no fields
func (_m *AgentService_AlgoClient[Req, Res]) CloseSend() error {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for CloseSend")
}
var r0 error
if rf, ok := ret.Get(0).(func() error); ok {
r0 = rf()
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentService_AlgoClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
type AgentService_AlgoClient_CloseSend_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// CloseSend is a helper method to define mock.On call
func (_e *AgentService_AlgoClient_Expecter[Req, Res]) CloseSend() *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
return &AgentService_AlgoClient_CloseSend_Call[Req, Res]{Call: _e.mock.On("CloseSend")}
}
func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) RunAndReturn(run func() error) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// Context provides a mock function with no fields
func (_m *AgentService_AlgoClient[Req, Res]) Context() context.Context {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Context")
}
var r0 context.Context
if rf, ok := ret.Get(0).(func() context.Context); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(context.Context)
}
}
return r0
}
// AgentService_AlgoClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
type AgentService_AlgoClient_Context_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// Context is a helper method to define mock.On call
func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Context() *AgentService_AlgoClient_Context_Call[Req, Res] {
return &AgentService_AlgoClient_Context_Call[Req, Res]{Call: _e.mock.On("Context")}
}
func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Context_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) Return(_a0 context.Context) *AgentService_AlgoClient_Context_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) RunAndReturn(run func() context.Context) *AgentService_AlgoClient_Context_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// Header provides a mock function with no fields
func (_m *AgentService_AlgoClient[Req, Res]) Header() (metadata.MD, error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Header")
}
var r0 metadata.MD
var r1 error
if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() metadata.MD); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(metadata.MD)
}
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// AgentService_AlgoClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
type AgentService_AlgoClient_Header_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// Header is a helper method to define mock.On call
func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Header() *AgentService_AlgoClient_Header_Call[Req, Res] {
return &AgentService_AlgoClient_Header_Call[Req, Res]{Call: _e.mock.On("Header")}
}
func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Header_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) Return(_a0 metadata.MD, _a1 error) *AgentService_AlgoClient_Header_Call[Req, Res] {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) RunAndReturn(run func() (metadata.MD, error)) *AgentService_AlgoClient_Header_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// RecvMsg provides a mock function with given fields: m
func (_m *AgentService_AlgoClient[Req, Res]) RecvMsg(m interface{}) error {
ret := _m.Called(m)
if len(ret) == 0 {
panic("no return value specified for RecvMsg")
}
var r0 error
if rf, ok := ret.Get(0).(func(interface{}) error); ok {
r0 = rf(m)
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentService_AlgoClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
type AgentService_AlgoClient_RecvMsg_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// RecvMsg is a helper method to define mock.On call
// - m interface{}
func (_e *AgentService_AlgoClient_Expecter[Req, Res]) RecvMsg(m interface{}) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
return &AgentService_AlgoClient_RecvMsg_Call[Req, Res]{Call: _e.mock.On("RecvMsg", m)}
}
func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(interface{}))
})
return _c
}
func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// Send provides a mock function with given fields: _a0
func (_m *AgentService_AlgoClient[Req, Res]) Send(_a0 *agent.AlgoRequest) error {
ret := _m.Called(_a0)
if len(ret) == 0 {
panic("no return value specified for Send")
}
var r0 error
if rf, ok := ret.Get(0).(func(*agent.AlgoRequest) error); ok {
r0 = rf(_a0)
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentService_AlgoClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
type AgentService_AlgoClient_Send_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// Send is a helper method to define mock.On call
// - _a0 *agent.AlgoRequest
func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Send(_a0 interface{}) *AgentService_AlgoClient_Send_Call[Req, Res] {
return &AgentService_AlgoClient_Send_Call[Req, Res]{Call: _e.mock.On("Send", _a0)}
}
func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) Run(run func(_a0 *agent.AlgoRequest)) *AgentService_AlgoClient_Send_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*agent.AlgoRequest))
})
return _c
}
func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_Send_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) RunAndReturn(run func(*agent.AlgoRequest) error) *AgentService_AlgoClient_Send_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// SendMsg provides a mock function with given fields: m
func (_m *AgentService_AlgoClient[Req, Res]) SendMsg(m interface{}) error {
ret := _m.Called(m)
if len(ret) == 0 {
panic("no return value specified for SendMsg")
}
var r0 error
if rf, ok := ret.Get(0).(func(interface{}) error); ok {
r0 = rf(m)
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentService_AlgoClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
type AgentService_AlgoClient_SendMsg_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// SendMsg is a helper method to define mock.On call
// - m interface{}
func (_e *AgentService_AlgoClient_Expecter[Req, Res]) SendMsg(m interface{}) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
return &AgentService_AlgoClient_SendMsg_Call[Req, Res]{Call: _e.mock.On("SendMsg", m)}
}
func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(interface{}))
})
return _c
}
func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// Trailer provides a mock function with no fields
func (_m *AgentService_AlgoClient[Req, Res]) Trailer() metadata.MD {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Trailer")
}
var r0 metadata.MD
if rf, ok := ret.Get(0).(func() metadata.MD); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(metadata.MD)
}
}
return r0
}
// AgentService_AlgoClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
type AgentService_AlgoClient_Trailer_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// Trailer is a helper method to define mock.On call
func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Trailer() *AgentService_AlgoClient_Trailer_Call[Req, Res] {
return &AgentService_AlgoClient_Trailer_Call[Req, Res]{Call: _e.mock.On("Trailer")}
}
func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) Return(_a0 metadata.MD) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) RunAndReturn(run func() metadata.MD) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// NewAgentService_AlgoClient creates a new instance of AgentService_AlgoClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewAgentService_AlgoClient[Req interface{}, Res interface{}](t interface {
mock.TestingT
Cleanup(func())
}) *AgentService_AlgoClient[Req, Res] {
mock := &AgentService_AlgoClient[Req, Res]{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
-434
View File
@@ -1,434 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mocks
import (
context "context"
agent "github.com/ultravioletrs/cocos/agent"
metadata "google.golang.org/grpc/metadata"
mock "github.com/stretchr/testify/mock"
)
// AgentService_DataClient is an autogenerated mock type for the AgentService_DataClient type
type AgentService_DataClient[Req interface{}, Res interface{}] struct {
mock.Mock
}
type AgentService_DataClient_Expecter[Req interface{}, Res interface{}] struct {
mock *mock.Mock
}
func (_m *AgentService_DataClient[Req, Res]) EXPECT() *AgentService_DataClient_Expecter[Req, Res] {
return &AgentService_DataClient_Expecter[Req, Res]{mock: &_m.Mock}
}
// CloseAndRecv provides a mock function with no fields
func (_m *AgentService_DataClient[Req, Res]) CloseAndRecv() (*agent.DataResponse, error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for CloseAndRecv")
}
var r0 *agent.DataResponse
var r1 error
if rf, ok := ret.Get(0).(func() (*agent.DataResponse, error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() *agent.DataResponse); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*agent.DataResponse)
}
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// AgentService_DataClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
type AgentService_DataClient_CloseAndRecv_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// CloseAndRecv is a helper method to define mock.On call
func (_e *AgentService_DataClient_Expecter[Req, Res]) CloseAndRecv() *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
return &AgentService_DataClient_CloseAndRecv_Call[Req, Res]{Call: _e.mock.On("CloseAndRecv")}
}
func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) Run(run func()) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) Return(_a0 *agent.DataResponse, _a1 error) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) RunAndReturn(run func() (*agent.DataResponse, error)) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// CloseSend provides a mock function with no fields
func (_m *AgentService_DataClient[Req, Res]) CloseSend() error {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for CloseSend")
}
var r0 error
if rf, ok := ret.Get(0).(func() error); ok {
r0 = rf()
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentService_DataClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
type AgentService_DataClient_CloseSend_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// CloseSend is a helper method to define mock.On call
func (_e *AgentService_DataClient_Expecter[Req, Res]) CloseSend() *AgentService_DataClient_CloseSend_Call[Req, Res] {
return &AgentService_DataClient_CloseSend_Call[Req, Res]{Call: _e.mock.On("CloseSend")}
}
func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) Run(run func()) *AgentService_DataClient_CloseSend_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_CloseSend_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) RunAndReturn(run func() error) *AgentService_DataClient_CloseSend_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// Context provides a mock function with no fields
func (_m *AgentService_DataClient[Req, Res]) Context() context.Context {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Context")
}
var r0 context.Context
if rf, ok := ret.Get(0).(func() context.Context); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(context.Context)
}
}
return r0
}
// AgentService_DataClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
type AgentService_DataClient_Context_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// Context is a helper method to define mock.On call
func (_e *AgentService_DataClient_Expecter[Req, Res]) Context() *AgentService_DataClient_Context_Call[Req, Res] {
return &AgentService_DataClient_Context_Call[Req, Res]{Call: _e.mock.On("Context")}
}
func (_c *AgentService_DataClient_Context_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Context_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_DataClient_Context_Call[Req, Res]) Return(_a0 context.Context) *AgentService_DataClient_Context_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_DataClient_Context_Call[Req, Res]) RunAndReturn(run func() context.Context) *AgentService_DataClient_Context_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// Header provides a mock function with no fields
func (_m *AgentService_DataClient[Req, Res]) Header() (metadata.MD, error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Header")
}
var r0 metadata.MD
var r1 error
if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() metadata.MD); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(metadata.MD)
}
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// AgentService_DataClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
type AgentService_DataClient_Header_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// Header is a helper method to define mock.On call
func (_e *AgentService_DataClient_Expecter[Req, Res]) Header() *AgentService_DataClient_Header_Call[Req, Res] {
return &AgentService_DataClient_Header_Call[Req, Res]{Call: _e.mock.On("Header")}
}
func (_c *AgentService_DataClient_Header_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Header_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_DataClient_Header_Call[Req, Res]) Return(_a0 metadata.MD, _a1 error) *AgentService_DataClient_Header_Call[Req, Res] {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *AgentService_DataClient_Header_Call[Req, Res]) RunAndReturn(run func() (metadata.MD, error)) *AgentService_DataClient_Header_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// RecvMsg provides a mock function with given fields: m
func (_m *AgentService_DataClient[Req, Res]) RecvMsg(m interface{}) error {
ret := _m.Called(m)
if len(ret) == 0 {
panic("no return value specified for RecvMsg")
}
var r0 error
if rf, ok := ret.Get(0).(func(interface{}) error); ok {
r0 = rf(m)
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentService_DataClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
type AgentService_DataClient_RecvMsg_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// RecvMsg is a helper method to define mock.On call
// - m interface{}
func (_e *AgentService_DataClient_Expecter[Req, Res]) RecvMsg(m interface{}) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
return &AgentService_DataClient_RecvMsg_Call[Req, Res]{Call: _e.mock.On("RecvMsg", m)}
}
func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(interface{}))
})
return _c
}
func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// Send provides a mock function with given fields: _a0
func (_m *AgentService_DataClient[Req, Res]) Send(_a0 *agent.DataRequest) error {
ret := _m.Called(_a0)
if len(ret) == 0 {
panic("no return value specified for Send")
}
var r0 error
if rf, ok := ret.Get(0).(func(*agent.DataRequest) error); ok {
r0 = rf(_a0)
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentService_DataClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
type AgentService_DataClient_Send_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// Send is a helper method to define mock.On call
// - _a0 *agent.DataRequest
func (_e *AgentService_DataClient_Expecter[Req, Res]) Send(_a0 interface{}) *AgentService_DataClient_Send_Call[Req, Res] {
return &AgentService_DataClient_Send_Call[Req, Res]{Call: _e.mock.On("Send", _a0)}
}
func (_c *AgentService_DataClient_Send_Call[Req, Res]) Run(run func(_a0 *agent.DataRequest)) *AgentService_DataClient_Send_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*agent.DataRequest))
})
return _c
}
func (_c *AgentService_DataClient_Send_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_Send_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_DataClient_Send_Call[Req, Res]) RunAndReturn(run func(*agent.DataRequest) error) *AgentService_DataClient_Send_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// SendMsg provides a mock function with given fields: m
func (_m *AgentService_DataClient[Req, Res]) SendMsg(m interface{}) error {
ret := _m.Called(m)
if len(ret) == 0 {
panic("no return value specified for SendMsg")
}
var r0 error
if rf, ok := ret.Get(0).(func(interface{}) error); ok {
r0 = rf(m)
} else {
r0 = ret.Error(0)
}
return r0
}
// AgentService_DataClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
type AgentService_DataClient_SendMsg_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// SendMsg is a helper method to define mock.On call
// - m interface{}
func (_e *AgentService_DataClient_Expecter[Req, Res]) SendMsg(m interface{}) *AgentService_DataClient_SendMsg_Call[Req, Res] {
return &AgentService_DataClient_SendMsg_Call[Req, Res]{Call: _e.mock.On("SendMsg", m)}
}
func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_DataClient_SendMsg_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(interface{}))
})
return _c
}
func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_SendMsg_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_DataClient_SendMsg_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// Trailer provides a mock function with no fields
func (_m *AgentService_DataClient[Req, Res]) Trailer() metadata.MD {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Trailer")
}
var r0 metadata.MD
if rf, ok := ret.Get(0).(func() metadata.MD); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(metadata.MD)
}
}
return r0
}
// AgentService_DataClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
type AgentService_DataClient_Trailer_Call[Req interface{}, Res interface{}] struct {
*mock.Call
}
// Trailer is a helper method to define mock.On call
func (_e *AgentService_DataClient_Expecter[Req, Res]) Trailer() *AgentService_DataClient_Trailer_Call[Req, Res] {
return &AgentService_DataClient_Trailer_Call[Req, Res]{Call: _e.mock.On("Trailer")}
}
func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Trailer_Call[Req, Res] {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) Return(_a0 metadata.MD) *AgentService_DataClient_Trailer_Call[Req, Res] {
_c.Call.Return(_a0)
return _c
}
func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) RunAndReturn(run func() metadata.MD) *AgentService_DataClient_Trailer_Call[Req, Res] {
_c.Call.Return(run)
return _c
}
// NewAgentService_DataClient creates a new instance of AgentService_DataClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewAgentService_DataClient[Req interface{}, Res interface{}](t interface {
mock.TestingT
Cleanup(func())
}) *AgentService_DataClient[Req, Res] {
mock := &AgentService_DataClient[Req, Res]{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
+153 -429
View File
@@ -6,78 +6,22 @@ package agent
import (
"context"
"encoding/json"
"errors"
"fmt"
"log/slog"
"os"
"path/filepath"
"os/exec"
"slices"
sync "sync"
"github.com/absmach/magistrala/pkg/errors"
"github.com/ultravioletrs/cocos/agent/algorithm"
"github.com/ultravioletrs/cocos/agent/algorithm/binary"
"github.com/ultravioletrs/cocos/agent/algorithm/docker"
"github.com/ultravioletrs/cocos/agent/algorithm/python"
"github.com/ultravioletrs/cocos/agent/algorithm/wasm"
"github.com/google/go-sev-guest/client"
"github.com/ultravioletrs/cocos/agent/events"
"github.com/ultravioletrs/cocos/agent/statemachine"
"github.com/ultravioletrs/cocos/internal"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
"github.com/ultravioletrs/cocos/pkg/socket"
"golang.org/x/crypto/sha3"
)
var _ Service = (*agentService)(nil)
//go:generate stringer -type=AgentState
type AgentState int
const (
Idle AgentState = iota
ReceivingManifest
ReceivingAlgorithm
ReceivingData
Running
ConsumingResults
Complete
Failed
)
//go:generate stringer -type=AgentEvent
type AgentEvent int
const (
Start AgentEvent = iota
ManifestReceived
AlgorithmReceived
DataReceived
RunComplete
ResultsConsumed
RunFailed
)
//go:generate stringer -type=Status
type Status uint8
const (
IdleState Status = iota
InProgress
Ready
Completed
Terminated
Warning
Starting
)
const (
algoFilePermission = 0o700
)
const (
ImaMeasurementsFilePath = "/sys/kernel/security/integrity/ima/ascii_runtime_measurements"
ImaPcrIndex = 10
)
// ReportDataSize is the size of the report data expected by the attestation service.
const ReportDataSize = 64
var (
// ErrMalformedEntity indicates malformed entity specification (e.g.
@@ -86,447 +30,227 @@ var (
// ErrUnauthorizedAccess indicates missing or invalid credentials provided
// when accessing a protected resource.
ErrUnauthorizedAccess = errors.New("missing or invalid credentials provided")
// ErrUndeclaredAlgorithm indicates algorithm was not declared in computation manifest.
ErrUndeclaredDataset = errors.New("dataset not declared in computation manifest")
// ErrAllManifestItemsReceived indicates no new computation manifest items expected.
ErrAllManifestItemsReceived = errors.New("all expected manifest Items have been received")
// ErrUndeclaredConsumer indicates the consumer requesting results in not declared in computation manifest.
ErrUndeclaredConsumer = errors.New("result consumer is undeclared in computation manifest")
// ErrResultsNotReady indicates the computation results are not ready.
ErrResultsNotReady = errors.New("computation results are not yet ready")
// ErrStateNotReady agent received a request in the wrong state.
ErrStateNotReady = errors.New("agent not expecting this operation in the current state")
// ErrHashMismatch provided algorithm/dataset does not match hash in manifest.
ErrHashMismatch = errors.New("malformed data, hash does not match manifest")
// ErrFileNameMismatch provided dataset filename does not match filename in manifest.
ErrFileNameMismatch = errors.New("malformed data, filename does not match manifest")
// ErrAllResultsConsumed indicates all results have been consumed.
ErrAllResultsConsumed = errors.New("all results have been consumed by declared consumers")
// ErrAttestationFailed attestation failed.
ErrAttestationFailed = errors.New("failed to get raw quote")
// ErrAttestationVTpmFailed vTPM attestation failed.
ErrAttestationVTpmFailed = errors.New("failed to get vTPM quote")
// ErrFetchAzureToken azure token fetch failed.
ErrFetchAzureToken = errors.New("failed to get azure token")
// ErrAttType indicates that the attestation type that is requested does not exist or is not supported.
ErrAttestationType = errors.New("attestation type does not exist or is not supported")
// errUndeclaredAlgorithm indicates algorithm was not declared in computation manifest.
errUndeclaredAlgorithm = errors.New("algorithm not declared in computation manifest")
// errUndeclaredAlgorithm indicates algorithm was not declared in computation manifest.
errUndeclaredDataset = errors.New("dataset not declared in computation manifest")
// errProviderMissmatch algorithm/dataset provider does not match computation manifest.
errProviderMissmatch = errors.New("provider does not match declaration on manifest")
// errAllManifestItemsReceived indicates no new computation manifest items expected.
errAllManifestItemsReceived = errors.New("all expected manifest Items have been received")
// errUndeclaredConsumer indicates the consumer requesting results in not declared in computation manifest.
errUndeclaredConsumer = errors.New("result consumer is undeclared in computation manifest")
// errResultsNotReady indicates the computation results are not ready.
errResultsNotReady = errors.New("computation results are not yet ready")
// errStateNotReady agent received a request in the wrong state.
errStateNotReady = errors.New("agent not expecting this operation in the current state")
// errHashMismatch provided algorithm/dataset does not match hash in manifest.
errHashMismatch = errors.New("malformed data, hash does not match manifest")
)
// Service specifies an API that must be fullfiled by the domain service
// implementation, and all of its decorators (e.g. logging & metrics).
type Service interface {
InitComputation(ctx context.Context, cmp Computation) error
StopComputation(ctx context.Context) error
Algo(ctx context.Context, algorithm Algorithm) error
Data(ctx context.Context, dataset Dataset) error
Result(ctx context.Context) ([]byte, error)
Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error)
IMAMeasurements(ctx context.Context) ([]byte, []byte, error)
AttestationResult(ctx context.Context, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error)
State() string
Result(ctx context.Context, consumer string) ([]byte, error)
Attestation(ctx context.Context, reportData [ReportDataSize]byte) ([]byte, error)
}
type agentService struct {
mu sync.Mutex
computation Computation // Holds the current computation request details.
algorithm algorithm.Algorithm // Filepath to the algorithm received for the computation.
result []byte // Stores the result of the computation.
sm statemachine.StateMachine // Manages the state transitions of the agent service.
runError error // Stores any error encountered during the computation run.
eventSvc events.Service // Service for publishing events related to computation.
provider attestation.Provider // Provider for generating attestation quotes.
logger *slog.Logger // Logger for the agent service.
resultsConsumed bool // Indicates if the results have been consumed.
cancel context.CancelFunc // Cancels the computation context.
vmpl int // VMPL at which the Agent is running.
computation Computation // Holds the current computation request details.
algorithms [][]byte // Stores the algorithms received for the computation.
datasets [][]byte // Stores the datasets received for the computation.
result []byte // Stores the result of the computation.
sm *StateMachine // Manages the state transitions of the agent service.
runError error // Stores any error encountered during the computation run.
eventSvc events.Service // Service for publishing events related to computation.
}
const (
socketPath = "unix_socket"
pyRuntime = "python3"
)
var _ Service = (*agentService)(nil)
// New instantiates the agent service implementation.
func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, provider attestation.Provider, vmlp int) Service {
sm := statemachine.NewStateMachine(Idle)
ctx, cancel := context.WithCancel(ctx)
func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp Computation) Service {
svc := &agentService{
sm: sm,
sm: NewStateMachine(logger),
eventSvc: eventSvc,
provider: provider,
logger: logger,
cancel: cancel,
vmpl: vmlp,
}
go svc.sm.Start(ctx)
svc.sm.SendEvent(start)
svc.sm.StateFunctions[idle] = svc.publishEvent("in-progress", json.RawMessage{})
svc.sm.StateFunctions[receivingManifests] = svc.publishEvent("in-progress", json.RawMessage{})
svc.sm.StateFunctions[receivingAlgorithms] = svc.publishEvent("in-progress", json.RawMessage{})
svc.sm.StateFunctions[receivingData] = svc.publishEvent("in-progress", json.RawMessage{})
svc.sm.StateFunctions[resultsReady] = svc.publishEvent("in-progress", json.RawMessage{})
svc.sm.StateFunctions[complete] = svc.publishEvent("in-progress", json.RawMessage{})
svc.sm.StateFunctions[running] = svc.runComputation
transitions := []statemachine.Transition{
{From: Idle, Event: Start, To: ReceivingManifest},
{From: ReceivingManifest, Event: ManifestReceived, To: ReceivingAlgorithm},
}
transitions = append(transitions, []statemachine.Transition{
{From: Running, Event: RunComplete, To: ConsumingResults},
{From: Running, Event: RunFailed, To: Failed},
{From: ConsumingResults, Event: ResultsConsumed, To: Complete},
}...)
for _, t := range transitions {
sm.AddTransition(t)
}
sm.SetAction(ReceivingAlgorithm, svc.publishEvent(InProgress.String()))
sm.SetAction(ReceivingData, svc.publishEvent(InProgress.String()))
sm.SetAction(Running, svc.runComputation)
sm.SetAction(ConsumingResults, svc.publishEvent(Ready.String()))
sm.SetAction(Complete, svc.publishEvent(Completed.String()))
sm.SetAction(Failed, svc.publishEvent(Failed.String()))
go func() {
if err := sm.Start(ctx); err != nil {
logger.Error(err.Error())
}
}()
sm.SendEvent(Start)
svc.computation = cmp
svc.sm.SendEvent(manifestsReceived)
return svc
}
func (as *agentService) State() string {
return as.sm.GetState().String()
}
func (as *agentService) InitComputation(ctx context.Context, cmp Computation) error {
if as.sm.GetState() != ReceivingManifest {
return ErrStateNotReady
func (as *agentService) Algo(ctx context.Context, algorithm Algorithm) error {
if as.sm.GetState() != receivingAlgorithms {
return errStateNotReady
}
defer as.sm.SendEvent(ManifestReceived)
as.mu.Lock()
defer as.mu.Unlock()
as.computation = cmp
transitions := []statemachine.Transition{}
if len(cmp.Datasets) == 0 {
transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: Running})
} else {
transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: ReceivingData})
transitions = append(transitions, statemachine.Transition{From: ReceivingData, Event: DataReceived, To: Running})
if len(as.computation.Algorithms) == 0 {
return errAllManifestItemsReceived
}
for _, t := range transitions {
as.sm.AddTransition(t)
}
hash := sha3.Sum256(algorithm.Algorithm)
return nil
}
func (as *agentService) StopComputation(ctx context.Context) error {
as.mu.Lock()
defer as.mu.Unlock()
as.eventSvc.SendEvent(as.computation.ID, "Stopped", "Stopped", json.RawMessage{})
as.cancel()
if as.algorithm != nil {
if err := as.algorithm.Stop(); err != nil {
return fmt.Errorf("error stopping computation: %v", err)
index := containsID(as.computation.Algorithms, algorithm.ID)
switch index {
case -1:
return errUndeclaredAlgorithm
default:
if as.computation.Algorithms[index].Provider != algorithm.Provider {
return errProviderMissmatch
}
}
if err := os.RemoveAll(algorithm.DatasetsDir); err != nil {
return fmt.Errorf("error removing datasets directory: %v", err)
}
if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
return fmt.Errorf("error removing results directory: %v", err)
}
as.sm.Reset(Idle)
as.computation = Computation{}
as.algorithm = nil
as.result = nil
as.runError = nil
as.resultsConsumed = false
ctx, cancel := context.WithCancel(ctx)
as.cancel = cancel
go func() {
if err := as.sm.Start(ctx); err != nil {
as.logger.Error(err.Error())
if hash != as.computation.Algorithms[index].Hash {
return errHashMismatch
}
}()
as.sm.SendEvent(Start)
return nil
}
func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
if as.sm.GetState() != ReceivingAlgorithm {
return ErrStateNotReady
}
as.mu.Lock()
defer as.mu.Unlock()
if as.algorithm != nil {
return ErrAllManifestItemsReceived
as.computation.Algorithms = slices.Delete(as.computation.Algorithms, index, index+1)
}
hash := sha3.Sum256(algo.Algorithm)
as.algorithms = append(as.algorithms, algorithm.Algorithm)
if hash != as.computation.Algorithm.Hash {
return ErrHashMismatch
}
currentDir, err := os.Getwd()
if err != nil {
return fmt.Errorf("error getting current directory: %v", err)
}
f, err := os.Create(filepath.Join(currentDir, "algo"))
if err != nil {
return fmt.Errorf("error creating algorithm file: %v", err)
}
if _, err := f.Write(algo.Algorithm); err != nil {
return fmt.Errorf("error writing algorithm to file: %v", err)
}
if err := os.Chmod(f.Name(), algoFilePermission); err != nil {
return fmt.Errorf("error changing file permissions: %v", err)
}
if err := f.Close(); err != nil {
return fmt.Errorf("error closing file: %v", err)
}
algoType := algorithm.AlgorithmTypeFromContext(ctx)
if algoType == "" {
algoType = string(algorithm.AlgoTypeBin)
}
args := algorithm.AlgorithmArgsFromContext(ctx)
switch algoType {
case string(algorithm.AlgoTypeBin):
as.algorithm = binary.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args, as.computation.ID)
case string(algorithm.AlgoTypePython):
var requirementsFile string
if len(algo.Requirements) > 0 {
fr, err := os.CreateTemp("", "requirements.txt")
if err != nil {
return fmt.Errorf("error creating requirments file: %v", err)
}
if _, err := fr.Write(algo.Requirements); err != nil {
return fmt.Errorf("error writing requirements to file: %v", err)
}
if err := fr.Close(); err != nil {
return fmt.Errorf("error closing file: %v", err)
}
requirementsFile = fr.Name()
}
runtime := python.PythonRunTimeFromContext(ctx)
as.algorithm = python.NewAlgorithm(as.logger, as.eventSvc, runtime, requirementsFile, f.Name(), args, as.computation.ID)
case string(algorithm.AlgoTypeWasm):
as.algorithm = wasm.NewAlgorithm(as.logger, as.eventSvc, args, f.Name(), as.computation.ID)
case string(algorithm.AlgoTypeDocker):
as.algorithm = docker.NewAlgorithm(as.logger, as.eventSvc, f.Name(), as.computation.ID)
}
if err := os.Mkdir(algorithm.DatasetsDir, 0o755); err != nil {
return fmt.Errorf("error creating datasets directory: %v", err)
}
if as.algorithm != nil {
as.sm.SendEvent(AlgorithmReceived)
if len(as.computation.Algorithms) == 0 {
as.sm.SendEvent(algorithmsReceived)
}
return nil
}
func (as *agentService) Data(ctx context.Context, dataset Dataset) error {
if as.sm.GetState() != ReceivingData {
return ErrStateNotReady
if as.sm.GetState() != receivingData {
return errStateNotReady
}
as.mu.Lock()
defer as.mu.Unlock()
if len(as.computation.Datasets) == 0 {
return ErrAllManifestItemsReceived
return errAllManifestItemsReceived
}
hash := sha3.Sum256(dataset.Dataset)
matched := false
for i, d := range as.computation.Datasets {
if hash == d.Hash {
if d.Filename != "" && d.Filename != dataset.Filename {
return ErrFileNameMismatch
}
as.computation.Datasets = slices.Delete(as.computation.Datasets, i, i+1)
if DecompressFromContext(ctx) {
if err := internal.UnzipFromMemory(dataset.Dataset, algorithm.DatasetsDir); err != nil {
return fmt.Errorf("error decompressing dataset: %v", err)
}
} else {
f, err := os.Create(fmt.Sprintf("%s/%s", algorithm.DatasetsDir, dataset.Filename))
if err != nil {
return fmt.Errorf("error creating dataset file: %v", err)
}
if _, err := f.Write(dataset.Dataset); err != nil {
return fmt.Errorf("error writing dataset to file: %v", err)
}
if err := f.Close(); err != nil {
return fmt.Errorf("error closing file: %v", err)
}
}
matched = true
break
index := containsID(as.computation.Datasets, dataset.ID)
switch index {
case -1:
return errUndeclaredDataset
default:
if as.computation.Datasets[index].Provider != dataset.Provider {
return errProviderMissmatch
}
if hash != as.computation.Datasets[index].Hash {
return errHashMismatch
}
as.computation.Datasets = slices.Delete(as.computation.Datasets, index, index+1)
}
if !matched {
return ErrUndeclaredDataset
}
as.datasets = append(as.datasets, dataset.Dataset)
if len(as.computation.Datasets) == 0 {
defer as.sm.SendEvent(DataReceived)
as.sm.SendEvent(dataReceived)
}
return nil
}
func (as *agentService) Result(ctx context.Context) ([]byte, error) {
currentState := as.sm.GetState()
if currentState != ConsumingResults && currentState != Complete && currentState != Failed {
return []byte{}, ErrResultsNotReady
func (as *agentService) Result(ctx context.Context, consumer string) ([]byte, error) {
if as.sm.GetState() != resultsReady {
return []byte{}, errResultsNotReady
}
if len(as.computation.ResultConsumers) == 0 {
return []byte{}, errAllManifestItemsReceived
}
index := slices.Index(as.computation.ResultConsumers, consumer)
switch index {
case -1:
return []byte{}, errUndeclaredConsumer
default:
as.computation.ResultConsumers = slices.Delete(as.computation.ResultConsumers, index, index+1)
}
index, ok := IndexFromContext(ctx)
if !ok {
return []byte{}, ErrUndeclaredConsumer
if len(as.computation.ResultConsumers) == 0 {
as.sm.SendEvent(resultsConsumed)
}
as.mu.Lock()
defer as.mu.Unlock()
if index < 0 || index >= len(as.computation.ResultConsumers) {
return []byte{}, ErrUndeclaredConsumer
}
if !as.resultsConsumed && currentState == ConsumingResults {
as.resultsConsumed = true
defer as.sm.SendEvent(ResultsConsumed)
}
// Return the result file or an error
return as.result, as.runError
}
func (as *agentService) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error) {
switch attType {
case attestation.SNP:
rawQuote, err := as.provider.TeeAttestation(reportData[:])
if err != nil {
return []byte{}, errors.Wrap(ErrAttestationFailed, err)
}
return rawQuote, nil
case attestation.VTPM:
vTPMQuote, err := as.provider.VTpmAttestation(nonce[:])
if err != nil {
return []byte{}, errors.Wrap(ErrAttestationVTpmFailed, err)
}
return vTPMQuote, nil
case attestation.SNPvTPM:
vTPMQuote, err := as.provider.Attestation(reportData[:], nonce[:])
if err != nil {
return []byte{}, errors.Wrap(ErrAttestationVTpmFailed, err)
}
return vTPMQuote, nil
default:
return []byte{}, ErrAttestationType
func (as *agentService) Attestation(ctx context.Context, reportData [ReportDataSize]byte) ([]byte, error) {
provider, err := client.GetQuoteProvider()
if err != nil {
return []byte{}, err
}
rawQuote, err := provider.GetRawQuote(reportData)
if err != nil {
return []byte{}, err
}
return rawQuote, nil
}
func (as *agentService) AttestationResult(ctx context.Context, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error) {
switch attType {
case attestation.AzureToken:
token, err := as.provider.AzureAttestationToken(nonce[:])
if err != nil {
return []byte{}, err
}
return token, nil
default:
return []byte{}, ErrAttestationType
}
}
func (as *agentService) runComputation(state statemachine.State) {
as.publishEvent(Starting.String())(state)
as.logger.Debug("computation run started")
defer func() {
if as.runError != nil {
as.sm.SendEvent(RunFailed)
} else {
as.sm.SendEvent(RunComplete)
}
}()
if err := os.Mkdir(algorithm.ResultsDir, 0o755); err != nil {
as.runError = fmt.Errorf("error creating results directory: %s", err.Error())
as.logger.Warn(as.runError.Error())
as.publishEvent(Failed.String())(state)
return
}
defer func() {
if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
as.logger.Warn(fmt.Sprintf("error removing results directory and its contents: %s", err.Error()))
}
if err := os.RemoveAll(algorithm.DatasetsDir); err != nil {
as.logger.Warn(fmt.Sprintf("error removing datasets directory and its contents: %s", err.Error()))
}
}()
as.publishEvent(InProgress.String())(state)
if err := as.algorithm.Run(); err != nil {
as.runError = err
as.logger.Warn(fmt.Sprintf("failed to run computation: %s", err.Error()))
as.publishEvent(Failed.String())(state)
return
}
results, err := internal.ZipDirectoryToMemory(algorithm.ResultsDir)
func (as *agentService) runComputation() {
as.publishEvent("starting", json.RawMessage{})()
as.sm.logger.Debug("computation run started")
defer as.sm.SendEvent(runComplete)
as.publishEvent("in-progress", json.RawMessage{})()
result, err := run(as.algorithms[0], as.datasets[0])
if err != nil {
as.runError = err
as.logger.Warn(fmt.Sprintf("failed to zip results: %s", err.Error()))
as.publishEvent(Failed.String())(state)
as.publishEvent("failed", json.RawMessage{})()
return
}
as.publishEvent(Completed.String())(state)
as.result = results
as.publishEvent("complete", json.RawMessage{})()
as.result = result
}
func (as *agentService) publishEvent(status string) statemachine.Action {
return func(state statemachine.State) {
as.eventSvc.SendEvent(as.computation.ID, state.String(), status, json.RawMessage{})
func (as *agentService) publishEvent(status string, details json.RawMessage) func() {
return func() {
if err := as.eventSvc.SendEvent(as.sm.State.String(), status, details); err != nil {
as.sm.logger.Warn(err.Error())
}
}
}
func (as *agentService) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
data, err := os.ReadFile(ImaMeasurementsFilePath)
func run(algoContent, dataContent []byte) ([]byte, error) {
listener, err := socket.StartUnixSocketServer(socketPath)
if err != nil {
return nil, nil, fmt.Errorf("Error reading Linux IMA measurements file: %s", err.Error())
return nil, fmt.Errorf("error creating stdout pipe: %v", err)
}
defer listener.Close()
// Create channels for received data and errors
dataChannel := make(chan []byte)
errorChannel := make(chan error)
var result []byte
go socket.AcceptConnection(listener, dataChannel, errorChannel)
// Construct the Python script content with CSV data as a command-line argument
script := string(algoContent)
data := string(dataContent)
cmd := exec.Command(pyRuntime, "-c", script, data, socketPath)
if err := cmd.Start(); err != nil {
return nil, fmt.Errorf("error starting Python script: %v", err)
}
pcr10, err := vtpm.GetPCRSHA1Value(ImaPcrIndex)
if err != nil {
return nil, nil, fmt.Errorf("Error reading TPM PCR #10: %s", err.Error())
if err := cmd.Wait(); err != nil {
return nil, fmt.Errorf("python script execution error: %v", err)
}
return data, pcr10, nil
select {
case result = <-dataChannel:
return result, nil
case err = <-errorChannel:
return nil, fmt.Errorf("error receiving data: %v", err)
}
}
-483
View File
@@ -1,483 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package agent
import (
"context"
"crypto/rand"
"log"
"os"
"testing"
"time"
mglog "github.com/absmach/magistrala/logger"
"github.com/absmach/magistrala/pkg/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/ultravioletrs/cocos/agent/algorithm"
"github.com/ultravioletrs/cocos/agent/algorithm/python"
"github.com/ultravioletrs/cocos/agent/events/mocks"
"github.com/ultravioletrs/cocos/agent/statemachine"
smmocks "github.com/ultravioletrs/cocos/agent/statemachine/mocks"
"github.com/ultravioletrs/cocos/pkg/attestation"
mocks2 "github.com/ultravioletrs/cocos/pkg/attestation/mocks"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
"golang.org/x/crypto/sha3"
"google.golang.org/grpc/metadata"
)
var (
algoPath = "../test/manual/algo/lin_reg.py"
reqPath = "../test/manual/algo/requirements.txt"
dataPath = "../test/manual/data/iris.csv"
)
const datasetFile = "iris.csv"
func TestAlgo(t *testing.T) {
algo, err := os.ReadFile(algoPath)
require.NoError(t, err)
algoHash := sha3.Sum256(algo)
reqFile, err := os.ReadFile(reqPath)
require.NoError(t, err)
testCases := []struct {
name string
err error
algo Algorithm
algoType string
}{
{
name: "Test Algo successfully",
algo: Algorithm{
Algorithm: algo,
Hash: algoHash,
},
algoType: "python",
err: nil,
},
{
name: "Test Algo successfully with requirements file",
algo: Algorithm{
Algorithm: algo,
Hash: algoHash,
Requirements: reqFile,
},
algoType: "python",
err: nil,
},
{
name: "Test Algo type binary successfully",
algo: Algorithm{
Algorithm: algo,
Hash: algoHash,
},
algoType: "bin",
err: nil,
},
{
name: "Test Algo type wasm successfully",
algo: Algorithm{
Algorithm: algo,
Hash: algoHash,
},
algoType: "wasm",
err: nil,
},
{
name: "Test Algo type docker successfully",
algo: Algorithm{
Algorithm: algo,
Hash: algoHash,
},
algoType: "docker",
err: nil,
},
{
name: "Test algo hash mismatch",
algo: Algorithm{},
algoType: "python",
err: ErrHashMismatch,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
err = os.RemoveAll("datasets")
require.NoError(t, err)
ctx := metadata.NewIncomingContext(context.Background(),
metadata.Pairs(algorithm.AlgoTypeKey, tc.algoType, python.PyRuntimeKey, python.PyRuntime),
)
events := new(mocks.Service)
events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
ctx, cancel := context.WithCancel(ctx)
defer cancel()
svc := New(ctx, mglog.NewMock(), events, &attestation.EmptyProvider{}, 0)
err := svc.InitComputation(ctx, testComputation(t))
require.NoError(t, err)
time.Sleep(300 * time.Millisecond)
err = svc.Algo(ctx, tc.algo)
assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
t.Cleanup(func() {
err = os.RemoveAll("venv")
err = os.RemoveAll("algo")
err = os.RemoveAll("datasets")
})
})
}
}
func TestData(t *testing.T) {
algo, err := os.ReadFile(algoPath)
require.NoError(t, err)
algoHash := sha3.Sum256(algo)
alg := Algorithm{
Hash: algoHash,
Algorithm: algo,
}
data, err := os.ReadFile(dataPath)
require.NoError(t, err)
dataHash := sha3.Sum256(data)
cases := []struct {
name string
data Dataset
err error
}{
{
name: "Test data successfully",
data: Dataset{
Hash: dataHash,
Dataset: data,
Filename: datasetFile,
},
},
{
name: "Test State not ready",
data: Dataset{
Dataset: data,
Hash: dataHash,
Filename: datasetFile,
},
err: ErrStateNotReady,
},
{
name: "Test File name does not match manifest",
data: Dataset{
Dataset: data,
Hash: dataHash,
Filename: "invalid",
},
err: ErrFileNameMismatch,
},
{
name: "Test dataset not declared in manifest",
data: Dataset{
Filename: datasetFile,
},
err: ErrUndeclaredDataset,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
ctx := metadata.NewIncomingContext(context.Background(),
metadata.Pairs(
algorithm.AlgoTypeKey, "python",
python.PyRuntimeKey, python.PyRuntime),
)
events := new(mocks.Service)
events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
if tc.err != ErrUndeclaredDataset {
ctx = IndexToContext(ctx, 0)
}
ctx, cancel := context.WithCancel(ctx)
defer cancel()
svc := New(ctx, mglog.NewMock(), events, &attestation.EmptyProvider{}, 0)
err := svc.InitComputation(ctx, testComputation(t))
require.NoError(t, err)
time.Sleep(300 * time.Millisecond)
if tc.err != ErrStateNotReady {
err = svc.Algo(ctx, alg)
require.NoError(t, err)
time.Sleep(300 * time.Millisecond)
}
err = svc.Data(ctx, tc.data)
assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
t.Cleanup(func() {
_ = os.RemoveAll("datasets")
_ = os.RemoveAll("results")
err = os.RemoveAll("venv")
err = os.RemoveAll("algo")
})
})
}
}
func TestResult(t *testing.T) {
cases := []struct {
name string
err error
setup func(svc *agentService)
ctxSetup func(ctx context.Context) context.Context
state statemachine.State
}{
{
name: "Test results not ready",
err: ErrResultsNotReady,
setup: func(svc *agentService) {
},
state: Running,
},
{
name: "Test undeclared consumer",
err: ErrUndeclaredConsumer,
setup: func(svc *agentService) {
svc.computation.ResultConsumers = []ResultConsumer{{UserKey: []byte("user")}}
},
ctxSetup: func(ctx context.Context) context.Context {
return ctx
},
state: ConsumingResults,
},
{
name: "Test results consumed and event sent",
err: nil,
setup: func(svc *agentService) {
svc.computation.ResultConsumers = []ResultConsumer{{UserKey: []byte("key")}}
},
ctxSetup: func(ctx context.Context) context.Context {
return IndexToContext(ctx, 0)
},
state: ConsumingResults,
},
}
for _, tc := range cases {
events := new(mocks.Service)
events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
t.Run(tc.name, func(t *testing.T) {
ctx := metadata.NewIncomingContext(context.Background(),
metadata.Pairs(algorithm.AlgoTypeKey, "python", python.PyRuntimeKey, python.PyRuntime),
)
if tc.ctxSetup != nil {
ctx = tc.ctxSetup(ctx)
}
sm := new(smmocks.StateMachine)
sm.On("Start", ctx).Return(nil)
sm.On("GetState").Return(tc.state)
sm.On("SendEvent", mock.Anything).Return()
svc := &agentService{
sm: sm,
eventSvc: events,
provider: &attestation.EmptyProvider{},
computation: testComputation(t),
}
go func() {
if err := svc.sm.Start(ctx); err != nil {
t.Errorf("Error starting state machine: %v", err)
}
}()
tc.setup(svc)
_, err := svc.Result(ctx)
t.Cleanup(func() {
_ = os.RemoveAll("datasets")
_ = os.RemoveAll("results")
})
assert.ErrorIs(t, err, tc.err, "expected %v, got %v", tc.err, err)
})
}
}
func TestAttestation(t *testing.T) {
provider := new(mocks2.Provider)
cases := []struct {
name string
reportData [quoteprovider.Nonce]byte
nonce [vtpm.Nonce]byte
rawQuote []uint8
platform attestation.PlatformType
err error
}{
{
name: "Test SNP attestation successful",
reportData: generateReportData(),
nonce: [32]byte{},
rawQuote: make([]uint8, 0),
platform: attestation.SNP,
err: nil,
},
{
name: "Test SNP attestation failed",
reportData: generateReportData(),
nonce: [32]byte{},
rawQuote: nil,
platform: attestation.SNP,
err: ErrAttestationFailed,
},
{
name: "Test vTPM attestation successful",
reportData: generateReportData(),
nonce: [32]byte{},
rawQuote: make([]uint8, 0),
platform: attestation.VTPM,
err: nil,
},
{
name: "Test vTPM attestation failed",
reportData: generateReportData(),
nonce: [32]byte{},
rawQuote: nil,
platform: attestation.VTPM,
err: ErrAttestationVTpmFailed,
},
{
name: "Test SNP-vTPM attestation successful",
reportData: generateReportData(),
nonce: [32]byte{},
rawQuote: make([]uint8, 0),
platform: attestation.SNPvTPM,
err: nil,
},
{
name: "Test SNP-vTPM attestation failed",
reportData: generateReportData(),
nonce: [32]byte{},
rawQuote: nil,
platform: attestation.SNPvTPM,
err: ErrAttestationVTpmFailed,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
events := new(mocks.Service)
events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
ctx := metadata.NewIncomingContext(context.Background(),
metadata.Pairs(algorithm.AlgoTypeKey, "python", python.PyRuntimeKey, python.PyRuntime),
)
ctx, cancel := context.WithCancel(ctx)
defer cancel()
getQuote := provider.On("TeeAttestation", mock.Anything).Return(tc.rawQuote, tc.err)
if tc.err != ErrAttestationFailed && tc.err != ErrAttestationVTpmFailed {
getQuote = provider.On("TeeAttestation", mock.Anything).Return(tc.nonce, nil)
}
defer getQuote.Unset()
svc := New(ctx, mglog.NewMock(), events, provider, 0)
time.Sleep(300 * time.Millisecond)
_, err := svc.Attestation(ctx, tc.reportData, tc.nonce, 0)
assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
})
}
}
func TestAttestationResult(t *testing.T) {
provider := new(mocks2.Provider)
cases := []struct {
name string
nonce [vtpm.Nonce]byte
platform attestation.PlatformType
token []byte
err error
}{
{
name: "Azure token fetch successful",
nonce: [32]byte{1, 2, 3}, // any test nonce
platform: attestation.AzureToken,
token: []byte("mockToken"),
err: nil,
},
{
name: "Azure token fetch failed",
nonce: [32]byte{4, 5, 6},
platform: attestation.AzureToken,
token: []byte{},
err: ErrFetchAzureToken,
},
{
name: "Invalid attestation type",
nonce: [32]byte{7, 8, 9},
platform: attestation.SNP,
token: []byte{},
err: ErrAttestationType,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
events := new(mocks.Service)
events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
if tc.platform == attestation.AzureToken {
provider.On("AzureAttestationToken", tc.nonce[:]).Return(tc.token, tc.err)
}
ctx := context.Background()
svc := New(ctx, mglog.NewMock(), events, provider, 0)
result, err := svc.AttestationResult(ctx, tc.nonce, tc.platform)
assert.True(t, errors.Contains(err, tc.err), "expected error %v, got %v", tc.err, err)
assert.Equal(t, tc.token, result)
})
}
}
func generateReportData() [quoteprovider.Nonce]byte {
bytes := make([]byte, quoteprovider.Nonce)
_, err := rand.Read(bytes)
if err != nil {
log.Fatalf("Failed to generate random bytes: %v", err)
}
return [64]byte(bytes)
}
func testComputation(t *testing.T) Computation {
algo, err := os.ReadFile(algoPath)
require.NoError(t, err)
algoHash := sha3.Sum256(algo)
data, err := os.ReadFile(dataPath)
require.NoError(t, err)
dataHash := sha3.Sum256(data)
return Computation{
ID: "1",
Name: "sample computation",
Description: "sample description",
Datasets: []Dataset{{Hash: dataHash, UserKey: []byte("key"), Dataset: data, Filename: datasetFile}},
Algorithm: Algorithm{Hash: algoHash, UserKey: []byte("key"), Algorithm: algo},
ResultConsumers: []ResultConsumer{{UserKey: []byte("key")}},
}
}
+121
View File
@@ -0,0 +1,121 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package agent
import (
"context"
"fmt"
"log/slog"
"sync"
)
//go:generate stringer -type=state
type state int
const (
idle state = iota
receivingManifests
receivingAlgorithms
receivingData
running
resultsReady
complete
)
type event int
const (
start event = iota
manifestsReceived
algorithmsReceived
dataReceived
runComplete
resultsConsumed
)
// StateMachine represents the state machine.
type StateMachine struct {
mu sync.Mutex
State state
EventChan chan event
Transitions map[state]map[event]state
StateFunctions map[state]func()
logger *slog.Logger
wg *sync.WaitGroup
}
// NewStateMachine creates a new StateMachine.
func NewStateMachine(logger *slog.Logger) *StateMachine {
sm := &StateMachine{
State: idle,
EventChan: make(chan event),
Transitions: make(map[state]map[event]state),
StateFunctions: make(map[state]func()),
logger: logger,
wg: &sync.WaitGroup{},
}
sm.Transitions[idle] = make(map[event]state)
sm.Transitions[idle][start] = receivingManifests
sm.Transitions[receivingManifests] = make(map[event]state)
sm.Transitions[receivingManifests][manifestsReceived] = receivingAlgorithms
sm.Transitions[receivingAlgorithms] = make(map[event]state)
sm.Transitions[receivingAlgorithms][algorithmsReceived] = receivingData
sm.Transitions[receivingData] = make(map[event]state)
sm.Transitions[receivingData][dataReceived] = running
sm.Transitions[running] = make(map[event]state)
sm.Transitions[running][runComplete] = resultsReady
sm.Transitions[resultsReady] = make(map[event]state)
sm.Transitions[resultsReady][resultsConsumed] = complete
return sm
}
// Start the state machine.
func (sm *StateMachine) Start(ctx context.Context) {
sm.wg.Add(1)
defer sm.wg.Done()
for {
select {
case event := <-sm.EventChan:
nextState, valid := sm.Transitions[sm.GetState()][event]
if valid {
sm.mu.Lock()
sm.State = nextState
sm.mu.Unlock()
sm.logger.Debug(fmt.Sprintf("Transition: %v -> %v\n", sm.GetState(), nextState))
} else {
sm.logger.Error(fmt.Sprintf("Invalid transition: %v -> ???\n", sm.GetState()))
}
stateFunc, exists := sm.StateFunctions[sm.GetState()]
if exists {
go stateFunc()
}
case <-ctx.Done():
return
}
}
}
// SendEvent sends an event to the state machine.
func (sm *StateMachine) SendEvent(event event) {
sm.EventChan <- event
}
func (sm *StateMachine) GetState() state {
sm.mu.Lock()
state := sm.State
sm.mu.Unlock()
return state
}
func (sm *StateMachine) SetState(state state) {
sm.mu.Lock()
sm.State = state
sm.mu.Unlock()
}
+29
View File
@@ -0,0 +1,29 @@
// Code generated by "stringer -type=state"; DO NOT EDIT.
package agent
import "strconv"
func _() {
// An "invalid array index" compiler error signifies that the constant values have changed.
// Re-run the stringer command to generate them again.
var x [1]struct{}
_ = x[idle-0]
_ = x[receivingManifests-1]
_ = x[receivingAlgorithms-2]
_ = x[receivingData-3]
_ = x[running-4]
_ = x[resultsReady-5]
_ = x[complete-6]
}
const _state_name = "idlereceivingManifestsreceivingAlgorithmsreceivingDatarunningresultsReadycomplete"
var _state_index = [...]uint8{0, 4, 22, 41, 54, 61, 73, 81}
func (i state) String() string {
if i < 0 || i >= state(len(_state_index)-1) {
return "state(" + strconv.FormatInt(int64(i), 10) + ")"
}
return _state_name[_state_index[i]:_state_index[i+1]]
}
+39 -230
View File
@@ -4,249 +4,58 @@ package agent
import (
"context"
sync "sync"
"fmt"
"testing"
"time"
"github.com/ultravioletrs/cocos/agent/statemachine"
mglog "github.com/absmach/magistrala/logger"
)
type MockState int
type MockEvent int
func (s MockState) String() string {
return []string{"State1", "State2", "State3"}[s]
}
func (e MockEvent) String() string {
return []string{"Event1", "Event2", "Event3"}[e]
}
const (
State1 MockState = iota
State2
State3
)
const (
Event1 MockEvent = iota
Event2
Event3
)
func TestNewStateMachine(t *testing.T) {
sm := statemachine.NewStateMachine(State1)
if sm == nil {
t.Fatal("NewStateMachine returned nil")
}
if sm.GetState() != State1 {
t.Errorf("Initial state not set correctly, got %v, want %v", sm.GetState(), State1)
}
}
func TestAddTransition(t *testing.T) {
sm := statemachine.NewStateMachine(State1)
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
go func() {
if err := sm.Start(ctx); err != context.Canceled {
t.Errorf("Start returned error: %v", err)
}
}()
sm.SendEvent(Event1)
time.Sleep(50 * time.Millisecond)
if sm.GetState() != State2 {
t.Errorf("Transition not applied correctly, got state %v, want %v", sm.GetState(), State2)
}
}
func TestSetAction(t *testing.T) {
sm := statemachine.NewStateMachine(State1)
var wg sync.WaitGroup
wg.Add(1)
sm.SetAction(State2, func(s statemachine.State) {
defer wg.Done()
})
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
go func() {
if err := sm.Start(ctx); err != context.Canceled {
t.Errorf("Start returned error: %v", err)
}
}()
sm.SendEvent(Event1)
wg.Wait()
if ctx.Err() != nil {
t.Error("Action was not called within the expected time")
}
}
func TestInvalidTransition(t *testing.T) {
sm := statemachine.NewStateMachine(State1)
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
errChan := make(chan error)
go func() {
errChan <- sm.Start(ctx)
}()
sm.SendEvent(Event2)
select {
case err := <-errChan:
if err == nil {
t.Errorf("Expected invalid transition error, got: %v", err)
}
case <-time.After(150 * time.Millisecond):
t.Error("Timeout waiting for invalid transition error")
}
}
func TestMultipleTransitions(t *testing.T) {
sm := statemachine.NewStateMachine(State1)
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
sm.AddTransition(statemachine.Transition{From: State2, Event: Event2, To: State3})
sm.AddTransition(statemachine.Transition{From: State3, Event: Event3, To: State1})
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
defer cancel()
go func() {
if err := sm.Start(ctx); err != context.Canceled {
t.Errorf("Start returned error: %v", err)
}
}()
transitions := []struct {
event MockEvent
want MockState
func TestStateMachineTransitions(t *testing.T) {
testCases := []struct {
fromState state
event event
expected state
}{
{Event1, State2},
{Event2, State3},
{Event3, State1},
{idle, start, receivingManifests},
{receivingManifests, manifestsReceived, receivingAlgorithms},
{receivingAlgorithms, algorithmsReceived, receivingData},
{receivingData, dataReceived, running},
{running, runComplete, resultsReady},
{resultsReady, resultsConsumed, complete},
}
for _, tt := range transitions {
sm.SendEvent(tt.event)
time.Sleep(50 * time.Millisecond)
for _, testCase := range testCases {
t.Run(fmt.Sprintf("Transition from %v to %v", testCase.fromState, testCase.expected), func(t *testing.T) {
sm := NewStateMachine(mglog.NewMock())
ctx, cancel := context.WithCancel(context.Background())
go func() {
sm.Start(ctx)
}()
sm.wg.Wait()
sm.SetState(testCase.fromState)
if sm.GetState() != tt.want {
t.Errorf("After event %v, got state %v, want %v", tt.event, sm.GetState(), tt.want)
}
sm.SendEvent(testCase.event)
if sm.GetState() != testCase.expected {
t.Errorf("Expected state %v after the event, but got %v", testCase.expected, sm.GetState())
}
close(sm.EventChan)
cancel()
})
}
}
func TestConcurrency(t *testing.T) {
sm := statemachine.NewStateMachine(State1)
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
sm.AddTransition(statemachine.Transition{From: State2, Event: Event2, To: State1})
func TestStateMachineInvalidTransition(t *testing.T) {
sm := NewStateMachine(mglog.NewMock())
ctx, cancel := context.WithCancel(context.Background())
go sm.Start(ctx)
ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
defer cancel()
sm.SetState(idle)
go func() {
if err := sm.Start(ctx); err == nil {
t.Errorf("Expected context error, got nil")
}
}()
sm.SendEvent(dataReceived)
for i := 0; i < 100; i++ {
go func() {
sm.SendEvent(Event1)
sm.SendEvent(Event2)
}()
}
time.Sleep(400 * time.Millisecond)
finalState := sm.GetState()
if finalState != State1 && finalState != State2 {
t.Errorf("Unexpected final state: %v", finalState)
}
}
func TestAgentEventString(t *testing.T) {
tests := []struct {
event AgentEvent
want string
}{
{Start, "Start"},
{ManifestReceived, "ManifestReceived"},
{AlgorithmReceived, "AlgorithmReceived"},
{DataReceived, "DataReceived"},
{RunComplete, "RunComplete"},
{ResultsConsumed, "ResultsConsumed"},
{RunFailed, "RunFailed"},
{AgentEvent(-1), "AgentEvent(-1)"},
}
for _, tt := range tests {
if got := tt.event.String(); got != tt.want {
t.Errorf("AgentEvent.String() = %v, want %v", got, tt.want)
}
}
}
func TestAgentStateString(t *testing.T) {
tests := []struct {
state AgentState
want string
}{
{Idle, "Idle"},
{ReceivingManifest, "ReceivingManifest"},
{ReceivingAlgorithm, "ReceivingAlgorithm"},
{ReceivingData, "ReceivingData"},
{Running, "Running"},
{ConsumingResults, "ConsumingResults"},
{Complete, "Complete"},
{Failed, "Failed"},
{AgentState(-1), "AgentState(-1)"},
}
for _, tt := range tests {
if got := tt.state.String(); got != tt.want {
t.Errorf("AgentState.String() = %v, want %v", got, tt.want)
}
}
}
func TestStatusString(t *testing.T) {
tests := []struct {
status Status
want string
}{
{IdleState, "IdleState"},
{InProgress, "InProgress"},
{Ready, "Ready"},
{Completed, "Completed"},
{Terminated, "Terminated"},
{Warning, "Warning"},
{Starting, "Starting"},
{Status(uint8(8)), "Status(8)"},
}
for _, tt := range tests {
if got := tt.status.String(); got != tt.want {
t.Errorf("Status.String() = %v, want %v", got, tt.want)
}
if sm.State != idle {
t.Errorf("State should not change on an invalid event, but got %v", sm.State)
}
cancel()
}
-266
View File
@@ -1,266 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mocks
import (
context "context"
mock "github.com/stretchr/testify/mock"
statemachine "github.com/ultravioletrs/cocos/agent/statemachine"
)
// StateMachine is an autogenerated mock type for the StateMachine type
type StateMachine struct {
mock.Mock
}
type StateMachine_Expecter struct {
mock *mock.Mock
}
func (_m *StateMachine) EXPECT() *StateMachine_Expecter {
return &StateMachine_Expecter{mock: &_m.Mock}
}
// AddTransition provides a mock function with given fields: t
func (_m *StateMachine) AddTransition(t statemachine.Transition) {
_m.Called(t)
}
// StateMachine_AddTransition_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AddTransition'
type StateMachine_AddTransition_Call struct {
*mock.Call
}
// AddTransition is a helper method to define mock.On call
// - t statemachine.Transition
func (_e *StateMachine_Expecter) AddTransition(t interface{}) *StateMachine_AddTransition_Call {
return &StateMachine_AddTransition_Call{Call: _e.mock.On("AddTransition", t)}
}
func (_c *StateMachine_AddTransition_Call) Run(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(statemachine.Transition))
})
return _c
}
func (_c *StateMachine_AddTransition_Call) Return() *StateMachine_AddTransition_Call {
_c.Call.Return()
return _c
}
func (_c *StateMachine_AddTransition_Call) RunAndReturn(run func(statemachine.Transition)) *StateMachine_AddTransition_Call {
_c.Run(run)
return _c
}
// GetState provides a mock function with no fields
func (_m *StateMachine) GetState() statemachine.State {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetState")
}
var r0 statemachine.State
if rf, ok := ret.Get(0).(func() statemachine.State); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(statemachine.State)
}
}
return r0
}
// StateMachine_GetState_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetState'
type StateMachine_GetState_Call struct {
*mock.Call
}
// GetState is a helper method to define mock.On call
func (_e *StateMachine_Expecter) GetState() *StateMachine_GetState_Call {
return &StateMachine_GetState_Call{Call: _e.mock.On("GetState")}
}
func (_c *StateMachine_GetState_Call) Run(run func()) *StateMachine_GetState_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *StateMachine_GetState_Call) Return(_a0 statemachine.State) *StateMachine_GetState_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *StateMachine_GetState_Call) RunAndReturn(run func() statemachine.State) *StateMachine_GetState_Call {
_c.Call.Return(run)
return _c
}
// Reset provides a mock function with given fields: initialState
func (_m *StateMachine) Reset(initialState statemachine.State) {
_m.Called(initialState)
}
// StateMachine_Reset_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Reset'
type StateMachine_Reset_Call struct {
*mock.Call
}
// Reset is a helper method to define mock.On call
// - initialState statemachine.State
func (_e *StateMachine_Expecter) Reset(initialState interface{}) *StateMachine_Reset_Call {
return &StateMachine_Reset_Call{Call: _e.mock.On("Reset", initialState)}
}
func (_c *StateMachine_Reset_Call) Run(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(statemachine.State))
})
return _c
}
func (_c *StateMachine_Reset_Call) Return() *StateMachine_Reset_Call {
_c.Call.Return()
return _c
}
func (_c *StateMachine_Reset_Call) RunAndReturn(run func(statemachine.State)) *StateMachine_Reset_Call {
_c.Run(run)
return _c
}
// SendEvent provides a mock function with given fields: event
func (_m *StateMachine) SendEvent(event statemachine.Event) {
_m.Called(event)
}
// StateMachine_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
type StateMachine_SendEvent_Call struct {
*mock.Call
}
// SendEvent is a helper method to define mock.On call
// - event statemachine.Event
func (_e *StateMachine_Expecter) SendEvent(event interface{}) *StateMachine_SendEvent_Call {
return &StateMachine_SendEvent_Call{Call: _e.mock.On("SendEvent", event)}
}
func (_c *StateMachine_SendEvent_Call) Run(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(statemachine.Event))
})
return _c
}
func (_c *StateMachine_SendEvent_Call) Return() *StateMachine_SendEvent_Call {
_c.Call.Return()
return _c
}
func (_c *StateMachine_SendEvent_Call) RunAndReturn(run func(statemachine.Event)) *StateMachine_SendEvent_Call {
_c.Run(run)
return _c
}
// SetAction provides a mock function with given fields: state, action
func (_m *StateMachine) SetAction(state statemachine.State, action statemachine.Action) {
_m.Called(state, action)
}
// StateMachine_SetAction_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetAction'
type StateMachine_SetAction_Call struct {
*mock.Call
}
// SetAction is a helper method to define mock.On call
// - state statemachine.State
// - action statemachine.Action
func (_e *StateMachine_Expecter) SetAction(state interface{}, action interface{}) *StateMachine_SetAction_Call {
return &StateMachine_SetAction_Call{Call: _e.mock.On("SetAction", state, action)}
}
func (_c *StateMachine_SetAction_Call) Run(run func(state statemachine.State, action statemachine.Action)) *StateMachine_SetAction_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(statemachine.State), args[1].(statemachine.Action))
})
return _c
}
func (_c *StateMachine_SetAction_Call) Return() *StateMachine_SetAction_Call {
_c.Call.Return()
return _c
}
func (_c *StateMachine_SetAction_Call) RunAndReturn(run func(statemachine.State, statemachine.Action)) *StateMachine_SetAction_Call {
_c.Run(run)
return _c
}
// Start provides a mock function with given fields: ctx
func (_m *StateMachine) Start(ctx context.Context) error {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for Start")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context) error); ok {
r0 = rf(ctx)
} else {
r0 = ret.Error(0)
}
return r0
}
// StateMachine_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
type StateMachine_Start_Call struct {
*mock.Call
}
// Start is a helper method to define mock.On call
// - ctx context.Context
func (_e *StateMachine_Expecter) Start(ctx interface{}) *StateMachine_Start_Call {
return &StateMachine_Start_Call{Call: _e.mock.On("Start", ctx)}
}
func (_c *StateMachine_Start_Call) Run(run func(ctx context.Context)) *StateMachine_Start_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *StateMachine_Start_Call) Return(_a0 error) *StateMachine_Start_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *StateMachine_Start_Call) RunAndReturn(run func(context.Context) error) *StateMachine_Start_Call {
_c.Call.Return(run)
return _c
}
// NewStateMachine creates a new instance of StateMachine. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewStateMachine(t interface {
mock.TestingT
Cleanup(func())
}) *StateMachine {
mock := &StateMachine{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
-127
View File
@@ -1,127 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package statemachine
import (
"context"
"fmt"
"sync"
)
type State interface {
String() string
}
type Event interface {
String() string
}
type Action func(State)
type Transition struct {
From State
Event Event
To State
}
type StateMachine interface {
AddTransition(t Transition)
SetAction(state State, action Action)
GetState() State
SendEvent(event Event)
Start(ctx context.Context) error
Reset(initialState State)
}
type stateMachine struct {
mu sync.Mutex
currentState State
transitions map[State]map[Event]State
actions map[State]Action
eventChan chan Event
}
func NewStateMachine(initialState State) StateMachine {
return &stateMachine{
currentState: initialState,
transitions: make(map[State]map[Event]State),
actions: make(map[State]Action),
eventChan: make(chan Event),
}
}
func (sm *stateMachine) AddTransition(t Transition) {
sm.mu.Lock()
defer sm.mu.Unlock()
if _, ok := sm.transitions[t.From]; !ok {
sm.transitions[t.From] = make(map[Event]State)
}
sm.transitions[t.From][t.Event] = t.To
}
func (sm *stateMachine) SetAction(state State, action Action) {
sm.mu.Lock()
defer sm.mu.Unlock()
sm.actions[state] = action
}
func (sm *stateMachine) GetState() State {
sm.mu.Lock()
defer sm.mu.Unlock()
return sm.currentState
}
func (sm *stateMachine) SendEvent(event Event) {
sm.eventChan <- event
}
func (sm *stateMachine) Start(ctx context.Context) error {
for {
select {
case event := <-sm.eventChan:
if err := sm.handleEvent(event); err != nil {
return err
}
case <-ctx.Done():
return ctx.Err()
}
}
}
func (sm *stateMachine) Reset(initialState State) {
sm.mu.Lock()
defer sm.mu.Unlock()
// Reset current state to initial state
sm.currentState = initialState
// Close the existing event channel to stop processing events
close(sm.eventChan)
// Create a new event channel
sm.eventChan = make(chan Event)
}
func (sm *stateMachine) handleEvent(event Event) error {
sm.mu.Lock()
currentState := sm.currentState
nextState, valid := sm.transitions[currentState][event]
sm.mu.Unlock()
if !valid {
return fmt.Errorf("invalid transition: %v -> %v", currentState, event)
}
sm.mu.Lock()
sm.currentState = nextState
action := sm.actions[nextState]
sm.mu.Unlock()
if action != nil {
go action(nextState)
}
return nil
}
-29
View File
@@ -1,29 +0,0 @@
// Code generated by "stringer -type=Status"; DO NOT EDIT.
package agent
import "strconv"
func _() {
// An "invalid array index" compiler error signifies that the constant values have changed.
// Re-run the stringer command to generate them again.
var x [1]struct{}
_ = x[IdleState-0]
_ = x[InProgress-1]
_ = x[Ready-2]
_ = x[Completed-3]
_ = x[Terminated-4]
_ = x[Warning-5]
_ = x[Starting-6]
}
const _Status_name = "IdleStateInProgressReadyCompletedTerminatedWarningStarting"
var _Status_index = [...]uint8{0, 9, 19, 24, 33, 43, 50, 58}
func (i Status) String() string {
if i >= Status(len(_Status_index)-1) {
return "Status(" + strconv.FormatInt(int64(i), 10) + ")"
}
return _Status_name[_Status_index[i]:_Status_index[i+1]]
}
BIN
View File
Binary file not shown.
-9
View File
@@ -1,9 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cocosai
import _ "embed"
//go:embed attestation.bin
var EmbeddedAttestation []byte
+14 -45
View File
@@ -1,6 +1,6 @@
# Agent CLI
This repository contains the command-line interface (CLI) tool for interacting with the Agent and manager service. The CLI allows you to perform various tasks such as running computations, uploading algorithm and datasets, and retrieving results.
This repository contains the command-line interface (CLI) tool for interacting with the Agent and manager service. The CLI allows you to perform various tasks such as running computations, uploading algorithms and datasets, and retrieving results.
## Build
@@ -12,18 +12,26 @@ make cli
## Usage
#### Run Computation
To run a computation, use the following command:
```bash
./build/cocos-cli manager run --computation '{"name": "my-computation"}'
```
#### Get attestation
Retrieves attestation information from the SEV guest and saves it to a file.
To retrieve attestation from agent, use the following command:
```bash
./build/cocos-cli attestation get '<report_data>'
./build/cocos-cli agent attestation get '<report_data>'
```
#### Validate attestation
Validates the retrieved attestation information against a specified policy and checks its authenticity.
To validate and verify attestation from agent, use the following command:
```bash
./build/cocos-cli attestation validate '<attestation>' --report_data '<report_data>'
./build/cocos-cli agent attestation validate '<attestation>' --report_data '<report_data>'
```
##### Flags
- --config: Path to a JSON file containing the validation configuration. This can be used to override individual flags.
@@ -62,60 +70,21 @@ To validate and verify attestation from agent, use the following command:
To upload an algorithm, use the following command:
```bash
./build/cocos-cli algo /path/to/algorithm <private_key_file_path>
./build/cocos-cli agent algo /path/to/algorithm
```
##### Flags
- -a, --algorithm string Algorithm type to run (default "bin")
- --args stringArray Arguments to pass to the algorithm
- --python-runtime string Python runtime to use (default "python3")
- -r, --requirements string Python requirements file
#### Upload Dataset
To upload a dataset, use the following command:
```bash
./build/cocos-cli data /path/to/dataset.csv <private_key_file_path>
./build/cocos-cli agent data /path/to/dataset.csv
```
Users can also upload directories which will be compressed on transit. Once received by agent they will be stored as compressed files or decompressed if the user passed the decompression argument.
##### Flags
- -d, --decompress Decompress the dataset on agent
#### Retrieve result
To retrieve the computation result, use the following command:
```bash
./build/cocos-cli result <private_key_file_path>
```
#### Checksum
When defining the manifest dataset and algorithm checksums are required. This can be done as below:
```bash
./build/cocos-cli checksum <path_to_dataset_or_algorithm>
```
#### Measure IGVM file
We assume that our current working directory is the root of the cocos repository, both on the host machine and in the VM.
`igvmmeasure` calculates the launch measurement for an IGVM file and can generate a signed version. It ensures integrity by precomputing the expected launch digest, which can be verified against the attestation report. The tool parses IGVM directives, outputs the measurement as a hex string, or creates a signed file for verification at guest launch.
##### Example
We measure an IGVM file using our measure command, run:
```bash
./build/cocos-cli igvmmeasure /path/to/igvm/file
```
The tool will parse the directives in the IGVM file, calculate the launch measurement, and output the computed digest. If successful, it prints the measurement to standard output.
Here is a sample output
```
91c4929bec2d0ecf11a708e09f0a57d7d82208bcba2451564444a4b01c22d047995ca27f9053f86de4e8063e9f810548
./build/cocos-cli agent result
```
-180
View File
@@ -1,180 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bytes"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/pem"
"errors"
"os"
"testing"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
)
const algorithmFile = "test_algo_file.py"
func generateRSAPrivateKeyFile(fileName string) error {
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
return err
}
privKeyFile, err := os.Create(fileName)
if err != nil {
return err
}
defer privKeyFile.Close()
privateKeyPEM := &pem.Block{
Type: rsaKeyType,
Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
}
err = pem.Encode(privKeyFile, privateKeyPEM)
if err != nil {
return err
}
return nil
}
func TestAlgorithmCmd(t *testing.T) {
tests := []struct {
name string
setupMock func(*mocks.SDK)
setupFiles func() error
args []string
connectErr error
expectedOutput string
cleanup func()
}{
{
name: "successful upload",
setupMock: func(m *mocks.SDK) {
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
},
setupFiles: func() error {
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
return err
}
return generateRSAPrivateKeyFile(privateKeyFile)
},
args: []string{algorithmFile, privateKeyFile},
expectedOutput: "Successfully uploaded algorithm",
cleanup: func() {
os.Remove(privateKeyFile)
os.Remove(algorithmFile)
},
},
{
name: "missing algorithm file",
setupMock: func(m *mocks.SDK) {
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
},
args: []string{"non_existent_algo_file.py", privateKeyFile},
expectedOutput: "Error reading algorithm file",
},
{
name: "missing private key file",
setupMock: func(m *mocks.SDK) {
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
},
setupFiles: func() error {
return os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
},
args: []string{algorithmFile, "non_existent_private_key.pem"},
expectedOutput: "Error reading private key file",
cleanup: func() {
os.Remove(algorithmFile)
},
},
{
name: "upload failure",
setupMock: func(m *mocks.SDK) {
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
},
setupFiles: func() error {
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
return err
}
return generateRSAPrivateKeyFile(privateKeyFile)
},
args: []string{algorithmFile, privateKeyFile},
expectedOutput: "Failed to upload algorithm",
cleanup: func() {
os.Remove(privateKeyFile)
os.Remove(algorithmFile)
},
},
{
name: "invalid private key",
setupMock: func(m *mocks.SDK) {
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
},
setupFiles: func() error {
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
return err
}
privKeyFile, err := os.Create(privateKeyFile)
if err != nil {
return err
}
defer privKeyFile.Close()
_, err = privKeyFile.WriteString("invalid private key")
return err
},
args: []string{algorithmFile, privateKeyFile},
expectedOutput: "Error decoding private key",
cleanup: func() {
os.Remove(algorithmFile)
os.Remove(privateKeyFile)
},
},
{
name: "connection error",
setupMock: func(m *mocks.SDK) {
},
connectErr: errors.New("failed to connect to agent"),
args: []string{algorithmFile, privateKeyFile},
expectedOutput: "Failed to connect to agent",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mockSDK := new(mocks.SDK)
if tt.setupMock != nil {
tt.setupMock(mockSDK)
}
testCLI := CLI{
agentSDK: mockSDK,
connectErr: tt.connectErr,
}
if tt.setupFiles != nil {
err := tt.setupFiles()
require.NoError(t, err)
}
cmd := testCLI.NewAlgorithmCmd()
buf := new(bytes.Buffer)
cmd.SetOut(buf)
cmd.SetArgs(tt.args)
err := cmd.Execute()
require.NoError(t, err)
require.Contains(t, buf.String(), tt.expectedOutput)
if tt.cleanup != nil {
tt.cleanup()
}
})
}
}
+16 -70
View File
@@ -3,94 +3,40 @@
package cli
import (
"context"
"encoding/pem"
"log"
"os"
"github.com/fatih/color"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/agent/algorithm"
"github.com/ultravioletrs/cocos/agent/algorithm/python"
"google.golang.org/grpc/metadata"
"github.com/ultravioletrs/cocos/agent"
)
var (
pythonRuntime string
algoType string
requirementsFile string
algoArgs []string
)
func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
cmd := &cobra.Command{
func (cli *CLI) NewAlgorithmsCmd() *cobra.Command {
return &cobra.Command{
Use: "algo",
Short: "Upload an algorithm binary",
Example: "algo <algo_file> <private_key_file_path>",
Args: cobra.ExactArgs(2),
Example: "algo <algo_file> <id> <provider>",
Args: cobra.ExactArgs(3),
Run: func(cmd *cobra.Command, args []string) {
if cli.connectErr != nil {
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
return
}
algorithmFile := args[0]
cmd.Println("Uploading algorithm file:", algorithmFile)
log.Println("Uploading algorithm binary:", algorithmFile)
algorithm, err := os.Open(algorithmFile)
algorithm, err := os.ReadFile(algorithmFile)
if err != nil {
printError(cmd, "Error reading algorithm file: %v", err)
return
log.Fatalf("Error reading algorithm file: %v", err)
}
defer algorithm.Close()
var req *os.File
if requirementsFile != "" {
req, err = os.Open(requirementsFile)
if err != nil {
printError(cmd, "Error reading requirments file: %v ❌ ", err)
return
}
defer req.Close()
algoReq := agent.Algorithm{
Algorithm: algorithm,
ID: args[1],
Provider: args[2],
}
privKeyFile, err := os.ReadFile(args[1])
if err != nil {
printError(cmd, "Error reading private key file: %v ❌ ", err)
return
if err := cli.agentSDK.Algo(cmd.Context(), algoReq); err != nil {
log.Fatalf("Error uploading algorithm with ID %s and provider %s: %v", algoReq.ID, algoReq.Provider, err)
}
pemBlock, _ := pem.Decode(privKeyFile)
privKey, err := decodeKey(pemBlock)
if err != nil {
printError(cmd, "Error decoding private key: %v ❌ ", err)
return
}
ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
if err := cli.agentSDK.Algo(addAlgoMetadata(ctx), algorithm, req, privKey); err != nil {
printError(cmd, "Failed to upload algorithm due to error: %v ❌ ", err)
return
}
cmd.Println(color.New(color.FgGreen).Sprint("Successfully uploaded algorithm! ✔ "))
log.Println("Successfully uploaded algorithm")
},
}
cmd.Flags().StringVarP(&algoType, "algorithm", "a", string(algorithm.AlgoTypeBin), "Algorithm type to run")
cmd.Flags().StringVar(&pythonRuntime, "python-runtime", python.PyRuntime, "Python runtime to use")
cmd.Flags().StringVarP(&requirementsFile, "requirements", "r", "", "Python requirements file")
cmd.Flags().StringArrayVar(&algoArgs, "args", []string{}, "Arguments to pass to the algorithm")
return cmd
}
func addAlgoMetadata(ctx context.Context) context.Context {
ctx = algorithm.AlgorithmTypeToContext(ctx, algoType)
ctx = algorithm.AlgorithmArgsToContext(ctx, algoArgs)
ctx = python.PythonRunTimeToContext(ctx, pythonRuntime)
return ctx
}
+120 -616
View File
@@ -3,61 +3,44 @@
package cli
import (
"encoding/base64"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"log"
"os"
"strconv"
"strings"
"time"
"github.com/absmach/magistrala/pkg/errors"
"github.com/fatih/color"
"github.com/google/go-sev-guest/abi"
"github.com/google/go-sev-guest/proto/check"
"github.com/google/go-sev-guest/proto/sevsnp"
"github.com/google/go-sev-guest/tools/lib/report"
tpmAttest "github.com/google/go-tpm-tools/proto/attest"
"github.com/google/go-sev-guest/validate"
"github.com/google/go-sev-guest/verify"
"github.com/google/go-sev-guest/verify/trust"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/azure"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
"github.com/ultravioletrs/cocos/agent"
"google.golang.org/protobuf/encoding/protojson"
"google.golang.org/protobuf/encoding/prototext"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/wrapperspb"
)
const (
defaultMinimumTcb = 0
defaultMinimumLaunchTcb = 0
defaultMinimumGuestSvn = 0
defaultGuestPolicy = 0x0000000000030000
defaultMinimumBuild = 0
defaultCheckCrl = false
defaultTimeout = 2 * time.Minute
defaultMaxRetryDelay = 30 * time.Second
defaultRequireAuthor = false
defaultRequireIdBlock = false
defaultMinVersion = "0.0"
size16 = 16
size32 = 32
size48 = 48
size64 = 64
attestationFilePath = "attestation.bin"
azureAttestResultFilePath = "azure_attest_result.json"
azureAttestTokenFilePath = "azure_attest_token.jwt"
vtpmFilePath = "../quote.dat"
attestationReportJson = "attestation.json"
sevProductNameMilan = "Milan"
sevProductNameGenoa = "Genoa"
FormatBinaryPB = "binarypb"
FormatTextProto = "textproto"
exampleJSONConfig = `
defaultMinimumTcb = 0
defaultMinimumLaunchTcb = 0
defaultMinimumGuestSvn = 0
defaultGuestPolicy = 0x0000000000030000
defaultMinimumBuild = 0
defaultCheckCrl = false
defaultTimeout = 2 * time.Minute
defaultMaxRetryDelay = 30 * time.Second
defaultRequireAuthor = false
defaultRequireIdBlock = false
defaultMinVersion = "0.0"
size16 = 16
size32 = 32
size48 = 48
size64 = 64
attestationFilePath = "attestation.bin"
exampleJSONConfig = `
{
"rootOfTrust":{
"product":"test_product",
@@ -110,56 +93,35 @@ const (
}
}
`
SNP = "snp"
VTPM = "vtpm"
SNPvTPM = "snp-vtpm"
AzureToken = "azure-token"
CCNone = "none"
CCAzure = "azure"
CCGCP = "gcp"
)
var (
mode string
cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
cfgString string
timeout time.Duration
maxRetryDelay time.Duration
platformInfo string
stepping string
trustedAuthorKeys []string
trustedAuthorHashes []string
trustedIdKeys []string
trustedIdKeyHashes []string
attestationFile string
tpmAttestationFile string
attestationRaw []byte
empty16 = [size16]byte{}
empty32 = [size32]byte{}
empty64 = [size64]byte{}
defaultReportIdMa = []byte{255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255}
errReportSize = errors.New("attestation contents too small")
ErrBadAttestation = errors.New("attestation file is corrupted or in wrong format")
output string
nonce []byte
format string
teeNonce []byte
tokenNonce []byte
getTextProtoAttestationReport bool
getAzureTokenJWT bool
cloud string
cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
cfgString string
timeout time.Duration
maxRetryDelay time.Duration
platformInfo string
stepping string
trustedAuthorKeys []string
trustedAuthorHashes []string
trustedIdKeys []string
trustedIdKeyHashes []string
attestationFile string
attestation []byte
empty16 = [size16]byte{}
empty32 = [size32]byte{}
empty64 = [size64]byte{}
defaultReportIdMa = []byte{255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255}
)
var errEmptyFile = errors.New("input file is empty")
func (cli *CLI) NewAttestationCmd() *cobra.Command {
return &cobra.Command{
Use: "attestation [command]",
Short: "Get and validate attestations",
Run: func(cmd *cobra.Command, args []string) {
cmd.Printf("Get and validate attestations\n\n")
cmd.Printf("Usage:\n %s [command]\n\n", cmd.CommandPath())
cmd.Printf("Available Commands:\n")
fmt.Printf("Get and validate attestations\n\n")
fmt.Printf("Usage:\n %s [command]\n\n", cmd.CommandPath())
fmt.Printf("Available Commands:\n")
// Filter out "completion" command
availableCommands := make([]*cobra.Command, 0)
@@ -170,353 +132,89 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
}
for _, subCmd := range availableCommands {
cmd.Printf(" %-15s%s\n", subCmd.Name(), subCmd.Short)
fmt.Printf(" %-15s%s\n", subCmd.Name(), subCmd.Short)
}
cmd.Printf("\nFlags:\n")
fmt.Printf("\nFlags:\n")
cmd.Flags().VisitAll(func(flag *pflag.Flag) {
cmd.Printf(" -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
fmt.Printf(" -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
})
cmd.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
fmt.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
},
}
}
func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "get",
Short: "Retrieve attestation information from agent. The argument of the command must be the type of the report (snp or vtpm or snp-vtpm).",
ValidArgs: []cobra.Completion{SNP, VTPM, SNPvTPM, AzureToken},
Example: fmt.Sprintf(`Based on attestation report type:
get %s --tee <512 bit hex value>
get %s --vtpm <256 bit hex value>
get %s --tee <512 bit hex value> --vtpm <256 bit hex value>
get %s --token <256 bit hex value>`, SNP, VTPM, SNPvTPM, AzureToken),
Args: cobra.ExactArgs(1),
return &cobra.Command{
Use: "get",
Short: "Retrieve attestation information from agent. Report data expected in hex enoded string of length 64 bytes.",
Example: "report <report_data>",
Args: cobra.ExactArgs(1),
Run: func(cmd *cobra.Command, args []string) {
if cli.connectErr != nil {
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
return
}
log.Println("Getting attestation")
if err := cobra.OnlyValidArgs(cmd, args); err != nil {
printError(cmd, "Bad attestation type: %v ❌ ", err)
return
}
attestationType := args[0]
attType := attestation.SNP
switch attestationType {
case SNP:
cmd.Println("Fetching SEV-SNP attestation report")
case VTPM:
cmd.Println("Fetching vTPM report")
attType = attestation.VTPM
case SNPvTPM:
cmd.Println("Fetching SEV-SNP and vTPM report")
attType = attestation.SNPvTPM
case AzureToken:
cmd.Println("Fetching Azure token")
attType = attestation.AzureToken
}
if (attType == attestation.VTPM || attType == attestation.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 {
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 {
msg := color.New(color.FgRed).Sprint("Token nonce must be defined for Azure attestation ❌ ")
cmd.Println(msg)
return
}
var fixedReportData [quoteprovider.Nonce]byte
if attType == attestation.SNP || attType == attestation.SNPvTPM {
if len(teeNonce) > quoteprovider.Nonce {
msg := color.New(color.FgRed).Sprintf("nonce must be a hex encoded string of length lesser or equal %d bytes ❌ ", quoteprovider.Nonce)
cmd.Println(msg)
return
}
copy(fixedReportData[:], teeNonce)
}
var fixedVtpmNonceByte [vtpm.Nonce]byte
if attType != attestation.SNP {
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 {
copy(fixedVtpmNonceByte[:], tokenNonce)
} else {
copy(fixedVtpmNonceByte[:], nonce)
}
}
filename := attestationFilePath
if attType == attestation.AzureToken {
filename = azureAttestResultFilePath
}
if getTextProtoAttestationReport {
filename = attestationReportJson
} else if getAzureTokenJWT {
filename = azureAttestTokenFilePath
}
attestationFile, err := os.Create(filename)
reportData, err := hex.DecodeString(args[0])
if err != nil {
printError(cmd, "Error creating attestation file: %v ❌ ", err)
return
log.Fatalf("attestation validation and verification failed with error: %s", err)
}
if len(reportData) != agent.ReportDataSize {
log.Fatalf("report data must be a hex encoded string of length %d bytes", agent.ReportDataSize)
}
var returnJsonAzureToken bool
if attType == attestation.AzureToken {
err := cli.agentSDK.AttestationResult(cmd.Context(), fixedVtpmNonceByte, int(attType), attestationFile)
if err != nil {
printError(cmd, "Failed to get attestation result due to error: %v ❌", err)
return
}
returnJsonAzureToken = !getAzureTokenJWT
} else {
err := cli.agentSDK.Attestation(cmd.Context(), fixedReportData, fixedVtpmNonceByte, int(attType), attestationFile)
if err != nil {
printError(cmd, "Failed to get attestation due to error: %v ❌", err)
return
}
result, err := cli.agentSDK.Attestation(cmd.Context(), [agent.ReportDataSize]byte(reportData))
if err != nil {
log.Fatalf("Error retrieving attestation: %v", err)
}
if err := attestationFile.Close(); err != nil {
printError(cmd, "Error closing attestation file: %v", err)
return
if err = os.WriteFile(attestationFilePath, result, 0o644); err != nil {
log.Fatalf("Error saving attestation result: %v", err)
}
if getTextProtoAttestationReport || returnJsonAzureToken {
result, err := os.ReadFile(filename)
if err != nil {
printError(cmd, "Error reading attestation file: %v ❌ ", err)
return
}
switch attestationType {
case SNP:
result, err = attesationToJSON(result)
if err != nil {
printError(cmd, "Error converting SNP attestation to JSON: %v ❌", err)
return
}
case VTPM, SNPvTPM:
marshalOptions := prototext.MarshalOptions{
Multiline: true,
EmitASCII: true,
}
var attvTPM tpmAttest.Attestation
err = proto.Unmarshal(result, &attvTPM)
if err != nil {
printError(cmd, "Failed to unmarshal the attestation report: %v ❌", err)
return
}
result = []byte(marshalOptions.Format(&attvTPM))
case AzureToken:
result, err = decodeJWTToJSON(result)
if err != nil {
printError(cmd, "Error decoding Azure token: %v ❌", err)
return
}
}
if err := os.WriteFile(filename, result, 0o644); err != nil {
printError(cmd, "Error writing attestation file: %v ❌ ", err)
return
}
}
cmd.Println("Attestation result retrieved and saved successfully!")
log.Println("Attestation result retrieved and saved successfully!")
},
}
cmd.Flags().BoolVarP(&getAzureTokenJWT, "azurejwt", "t", false, "Get azure attestation token as jwt format")
cmd.Flags().BoolVarP(&getTextProtoAttestationReport, "reporttextproto", "r", false, "Get attestation report in textproto format")
cmd.Flags().BytesHexVar(&teeNonce, "tee", []byte{}, "Define the nonce for the SNP attestation report (must be used with attestation type snp and snp-vtpm)")
cmd.Flags().BytesHexVar(&nonce, "vtpm", []byte{}, "Define the nonce for the vTPM attestation report (must be used with attestation type vtpm and snp-vtpm)")
cmd.Flags().BytesHexVar(&tokenNonce, "token", []byte{}, "Define the nonce for the Azure attestation token (must be used with attestation type azure-token)")
return cmd
}
func attesationToJSON(report []byte) ([]byte, error) {
if len(report) < abi.ReportSize {
return nil, errors.Wrap(errReportSize, fmt.Errorf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes", len(report), abi.ReportSize))
}
attestationPB, err := abi.ReportCertsToProto(report[:abi.ReportSize])
if err != nil {
return nil, err
}
return json.MarshalIndent(attestationPB, "", " ")
}
func attesationFromJSON(reportFile []byte) ([]byte, error) {
var attestationPB sevsnp.Attestation
if err := json.Unmarshal(reportFile, &attestationPB); err != nil {
return nil, err
}
return report.Transform(&attestationPB, "bin")
}
func isFileJSON(filename string) bool {
return strings.HasSuffix(filename, ".json")
}
func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "validate",
Short: fmt.Sprintf("Validate and verify attestation information. You can define the confidential computing cloud provider (%s, %s, %s; %s is the default) and can choose from 3 modes: %s, %s and %s. Default mode is %s.", CCNone, CCAzure, CCGCP, CCNone, SNP, VTPM, SNPvTPM, SNP),
Example: `Based on mode:
validate <attestationreportfilepath> --report_data <reportdata> --product <product data> --platform <cc platform> //default
validate --mode snp <attestationreportfilepath> --report_data <reportdata> --product <product data>
validate --mode vtpm <attestationreportfilepath> --nonce <noncevalue> --format <formatvalue> --output <outputvalue>
validate --mode snp-vtpm <attestationreportfilepath> --report_data <reportdata> --product <product data> --nonce <noncevalue> --format <formatvalue> --output <outputvalue>
validate --cloud none --mode snp <attestationreportfilepath> --report_data <reportdata> --product <product data>
validate --cloud azure --mode vtpm <attestationreportfilepath> --nonce <noncevalue> --format <formatvalue> --output <outputvalue>
validate --cloud gcp --mode snp-vtpm <attestationreportfilepath> --report_data <reportdata> --product <product data> --nonce <noncevalue> --format <formatvalue> --output <outputvalue>`,
PreRunE: func(cmd *cobra.Command, args []string) error {
mode, _ := cmd.Flags().GetString("mode")
if len(args) != 1 {
return fmt.Errorf("please pass the attestation report file path")
Use: "validate",
Short: "Validate and verify attestation information. The report is provided as a file path.",
Example: "validate <attestation_report_file_path>",
Args: cobra.ExactArgs(1),
Run: func(cmd *cobra.Command, args []string) {
log.Println("Checking attestation")
attestationFile = string(args[0])
if err := parseConfig(); err != nil {
log.Fatalf("attestation validation and verification failed with error: %s", err)
}
if err := parseHashes(); err != nil {
log.Fatalf("attestation validation and verification failed with error: %s", err)
}
if err := parseFiles(); err != nil {
log.Fatalf("attestation validation and verification failed with error: %s", err)
}
// This format is the attestation report in AMD's specified ABI format, immediately
// followed by the certificate table bytes.
if len(attestation) < abi.ReportSize {
log.Fatalf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes", len(attestation), abi.ReportSize)
}
if err := parseUints(); err != nil {
log.Fatalf("attestation validation and verification failed with error: %s", err)
}
cfg.Policy.Vmpl = wrapperspb.UInt32(0)
if err := validateInput(); err != nil {
log.Fatalf("attestation validation and verification failed with error: %s", err)
}
// Validate flags based on the mode
switch mode {
case SNP:
if err := cmd.MarkFlagRequired("report_data"); err != nil {
return fmt.Errorf("failed to mark 'report_data' as required for SEV-%s mode: %v", SNP, err)
}
if err := cmd.MarkFlagRequired("product"); err != nil {
return fmt.Errorf("failed to mark flag as required: %v ❌ ", err)
}
case SNPvTPM:
if err := cmd.MarkFlagRequired("nonce"); err != nil {
return fmt.Errorf("failed to mark 'nonce' as required for %s mode: %v", VTPM, err)
}
if err := cmd.MarkFlagRequired("report_data"); err != nil {
return fmt.Errorf("failed to mark 'report_data' as required for SEV-%s mode: %v", SNP, err)
}
if err := cmd.MarkFlagRequired("product"); err != nil {
return fmt.Errorf("failed to mark flag as required: %v ❌ ", err)
}
if err := cmd.MarkFlagRequired("format"); err != nil {
return fmt.Errorf("failed to mark 'format' as required for %s mode: %v", VTPM, err)
}
if err := cmd.MarkFlagRequired("output"); err != nil {
return fmt.Errorf("failed to mark 'output' as required for %s mode: %v", VTPM, err)
}
case VTPM:
if err := cmd.MarkFlagRequired("nonce"); err != nil {
return fmt.Errorf("failed to mark 'nonce' as required for %s mode: %v", VTPM, err)
}
if err := cmd.MarkFlagRequired("format"); err != nil {
return fmt.Errorf("failed to mark 'format' as required for %s mode: %v", VTPM, err)
}
if err := cmd.MarkFlagRequired("output"); err != nil {
return fmt.Errorf("failed to mark 'output' as required for %s mode: %v", VTPM, err)
}
default:
return fmt.Errorf("unknown mode: %s", mode)
if err := verifyAndValidateAttestation(attestation); err != nil {
log.Fatalf("attestation validation and verification failed with error: %s", err)
}
return nil
log.Println("Attestation validation and verification is successful!")
},
RunE: func(cmd *cobra.Command, args []string) error {
mode, _ := cmd.Flags().GetString("mode")
cloud, _ := cmd.Flags().GetString("cloud")
output, err := createOutputFile()
if err != nil {
return fmt.Errorf("failed to create output file: %v ❌ ", err)
}
if closer, ok := output.(*os.File); ok {
defer closer.Close()
}
var provider attestation.Provider
switch cloud {
case CCNone:
provider = vtpm.New(nil, false, 0, output)
case CCAzure:
provider = azure.New(output)
case CCGCP:
provider = vtpm.New(nil, false, 0, output)
default:
provider = vtpm.New(nil, false, 0, output)
}
switch mode {
case SNP:
return sevsnpverify(cmd, provider, args)
case SNPvTPM:
return vtpmSevSnpverify(args, provider)
case VTPM:
return vtpmverify(args, provider)
default:
return fmt.Errorf("unknown mode: %s", mode)
}
},
SilenceUsage: true,
SilenceErrors: true,
}
cmd.Flags().StringVar(
&cloud,
"cloud",
"none", // default CC provider
"The confidential computing cloud provider. Example: azure",
)
cmd.Flags().StringVar(
&mode,
"mode",
"snp", // default mode
"The attestation validation mode. Example: snp",
)
// VTPM FLAGS
cmd.Flags().BytesHexVar(
&nonce,
"nonce",
[]byte{},
"hex encoded nonce for vTPM attestation, cannot be empty",
)
cmd.Flags().StringVar(
&format,
"format",
"binarypb", // default value
"type of output file where attestation report stored <binarypb|textproto>",
)
cmd.Flags().StringVar(
&output,
"output",
"",
"output file",
)
// SEV-SNP FLAGS
cmd.Flags().StringVar(
&cfgString,
"config",
@@ -668,7 +366,7 @@ func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
"Hex-encoded SHA-384 hash values of trusted identity keys in AMD public key format",
)
cmd.Flags().StringVar(
&cfg.RootOfTrust.ProductLine,
&cfg.RootOfTrust.Product,
"product",
"",
"The AMD product name for the chip that generated the attestation report.",
@@ -692,180 +390,41 @@ func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
"PEM format CA bundles for the AMD product. Combined with contents of cabundle_paths.",
)
if err := cmd.MarkFlagRequired("report_data"); err != nil {
log.Fatalf("Failed to mark flag as required: %s", err)
}
return cmd
}
func (cli *CLI) NewMeasureCmd(igvmBinaryPath string) *cobra.Command {
igvmmeasureCmd := &cobra.Command{
Use: "igvmmeasure <INPUT>",
Short: "Measure an IGVM file",
Long: `igvmmeasure measures an IGVM file and outputs the calculated measurement.
It ensures integrity verification for the IGVM file.`,
Args: cobra.MinimumNArgs(1),
RunE: func(cmd *cobra.Command, args []string) error {
if len(args) == 0 {
return fmt.Errorf("error: No input file provided")
}
inputFile := args[0]
measurement, err := cli.measurement.Run(inputFile)
if err != nil {
return err
}
outputString := string(measurement)
lines := strings.Split(strings.TrimSpace(outputString), "\n")
if len(lines) == 1 {
outputString = strings.ToLower(outputString)
} else {
return fmt.Errorf("error: %s", outputString)
}
cmd.Print(outputString)
return nil
},
}
return igvmmeasureCmd
}
func sevsnpverify(cmd *cobra.Command, provider attestation.Provider, args []string) error {
cmd.Println("Checking attestation")
attestationFile = string(args[0])
if err := parseAttestationFile(); err != nil {
return fmt.Errorf("error parsing config: %v ❌ ", err)
}
// This format is the attestation report in AMD's specified ABI format, immediately
// followed by the certificate table bytes.
if len(attestationRaw) < abi.ReportSize {
return fmt.Errorf("attestation too small: got 0x%x bytes, need at least 0x%x bytes", len(attestationRaw), abi.ReportSize)
}
if err := parseAttestationConfig(); err != nil {
return err
}
// Used for verification of SNP attestation report
attestation.AttestationPolicy.Config = &cfg
if err := provider.VerifTeeAttestation(attestationRaw, cfg.Policy.ReportData); err != nil {
return fmt.Errorf("attestation validation and verification failed with error: %v ❌ ", err)
}
cmd.Println("Attestation validation and verification is successful!")
return nil
}
func parseAttestationConfig() error {
if err := parseConfig(); err != nil {
return fmt.Errorf("error parsing config: %v ❌ ", err)
}
if err := parseHashes(); err != nil {
return fmt.Errorf("error parsing hashes: %v ❌ ", err)
}
if err := parseTrustedKeys(); err != nil {
return fmt.Errorf("error parsing files: %v ❌ ", err)
}
if err := parseUints(); err != nil {
return fmt.Errorf("error parsing uints: %v ❌ ", err)
}
if err := validateInput(); err != nil {
return fmt.Errorf("error validating input: %v ❌ ", err)
}
return nil
}
func vtpmSevSnpverify(args []string, provider attestation.Provider) error {
attest, err := returnvTPMAttestation(args)
func verifyAndValidateAttestation(attestation []byte) error {
sopts, err := verify.RootOfTrustToOptions(cfg.RootOfTrust)
if err != nil {
return err
}
if err := parseAttestationConfig(); err != nil {
return err
sopts.Product = cfg.Policy.Product
sopts.Getter = &trust.RetryHTTPSGetter{
Timeout: timeout,
MaxRetryDelay: maxRetryDelay,
Getter: &trust.SimpleHTTPSGetter{},
}
// Used for verification of SNP attestation report
attestation.AttestationPolicy.Config = &cfg
if err := provider.VerifyAttestation(attest, cfg.Policy.ReportData, nonce); err != nil {
return fmt.Errorf("attestation validation and verification failed with error: %v ❌ ", err)
}
return nil
}
func vtpmverify(args []string, provider attestation.Provider) error {
attestation, err := returnvTPMAttestation(args)
attestationPB, err := abi.ReportCertsToProto(attestation)
if err != nil {
return err
}
if err := provider.VerifVTpmAttestation(attestation, nonce); err != nil {
return fmt.Errorf("attestation validation and verification failed with error: %v ❌ ", err)
if err = verify.SnpAttestation(attestationPB, sopts); err != nil {
return err
}
opts, err := validate.PolicyToOptions(cfg.Policy)
if err != nil {
return err
}
if err = validate.SnpAttestation(attestationPB, opts); err != nil {
return err
}
return nil
}
func returnvTPMAttestation(args []string) ([]byte, error) {
tpmAttestationFile = string(args[0])
input, err := openInputFile()
if err != nil {
return nil, err
}
if closer, ok := input.(*os.File); ok {
defer closer.Close()
}
attestationBytes, err := io.ReadAll(input)
if err != nil {
return nil, err
}
attestation := &tpmAttest.Attestation{}
if format == FormatBinaryPB {
return attestationBytes, nil
} else if format == FormatTextProto {
unmarshalOptions := prototext.UnmarshalOptions{}
err = unmarshalOptions.Unmarshal(attestationBytes, attestation)
} else {
return nil, fmt.Errorf("format should be either binarypb or textproto")
}
if err != nil {
return nil, fmt.Errorf("fail to unmarshal attestation report: %v", err)
}
attestationBytes, err = proto.Marshal(attestation)
if err != nil {
return nil, fmt.Errorf("fail to marshal vTPM attestation report: %v", err)
}
return attestationBytes, nil
}
func openInputFile() (io.Reader, error) {
if tpmAttestationFile == "" {
return nil, errEmptyFile
}
return os.Open(tpmAttestationFile)
}
func createOutputFile() (io.Writer, error) {
if output == "" {
return os.Stdout, nil
}
return os.Create(output)
}
// parseConfig decodes config passed as json for check.Config struct.
// example
/* {
@@ -954,23 +513,12 @@ func parseHashes() error {
return nil
}
func parseAttestationFile() error {
func parseFiles() error {
file, err := os.ReadFile(attestationFile)
if err != nil {
return err
}
attestationRaw = file
if isFileJSON(attestationFile) {
attestationRaw, err = attesationFromJSON(attestationRaw)
if err != nil {
return err
}
}
return nil
}
func parseTrustedKeys() error {
attestation = file
for _, path := range trustedAuthorKeys {
file, err := os.ReadFile(path)
if err != nil {
@@ -1036,7 +584,7 @@ func getBase(val string) int {
}
func validateInput() error {
if len(cfg.RootOfTrust.CabundlePaths) != 0 || len(cfg.RootOfTrust.Cabundles) != 0 && cfg.RootOfTrust.ProductLine == "" {
if len(cfg.RootOfTrust.CabundlePaths) != 0 || len(cfg.RootOfTrust.Cabundles) != 0 && cfg.RootOfTrust.Product == "" {
return fmt.Errorf("product name must be set if CA bundles are provided")
}
@@ -1083,47 +631,3 @@ func validateFieldLength(fieldName string, field []byte, expectedLength int) err
}
return nil
}
func decodeJWTToJSON(tokenBytes []byte) ([]byte, error) {
token := string(tokenBytes) // convert to string
parts := strings.Split(token, ".")
if len(parts) < 2 {
return nil, fmt.Errorf("invalid JWT: must have at least 2 parts")
}
decode := func(seg string) (map[string]interface{}, error) {
// Add padding if missing
if m := len(seg) % 4; m != 0 {
seg += strings.Repeat("=", 4-m)
}
data, err := base64.URLEncoding.DecodeString(seg)
if err != nil {
return nil, err
}
var result map[string]interface{}
if err := json.Unmarshal(data, &result); err != nil {
return nil, err
}
return result, nil
}
header, err := decode(parts[0])
if err != nil {
return nil, fmt.Errorf("failed to decode header: %v", err)
}
payload, err := decode(parts[1])
if err != nil {
return nil, fmt.Errorf("failed to decode payload: %v", err)
}
combined := map[string]interface{}{
"header": header,
"payload": payload,
}
return json.MarshalIndent(combined, "", " ")
}
-319
View File
@@ -1,319 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bytes"
"crypto/sha512"
"encoding/base64"
"encoding/json"
"fmt"
"os"
"strconv"
"github.com/absmach/magistrala/pkg/errors"
"github.com/google/go-sev-guest/proto/check"
"github.com/google/go-tpm-tools/proto/attest"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/azure"
"github.com/ultravioletrs/cocos/pkg/attestation/gcp"
"google.golang.org/protobuf/proto"
)
type fieldType int
const (
measurementField fieldType = iota
hostDataField
)
const (
// 0o744 file permission gives RWX permission to the user and only the R permission to others.
filePermission = 0o744
// Length of the expected host data and measurement field in bytes.
hostDataLength = 32
measurementLength = 48
)
var (
errDecode = errors.New("base64 string could not be decoded")
errDataLength = errors.New("data does not have an adequate length")
errReadingAttestationPolicyFile = errors.New("error while reading the attestation policy file")
errUnmarshalJSON = errors.New("failed to unmarshal json")
errMarshalJSON = errors.New("failed to marshal json")
errWriteFile = errors.New("failed to write to file")
errAttestationPolicyField = errors.New("the specified field type does not exist in the attestation policy")
policy uint64 = 196639
)
func (cli *CLI) NewAttestationPolicyCmd() *cobra.Command {
return &cobra.Command{
Use: "policy [command]",
Short: "Change attestation policy",
Run: func(cmd *cobra.Command, args []string) {
fmt.Printf("Change attestation policy\n\n")
fmt.Printf("Usage:\n %s [command]\n\n", cmd.CommandPath())
fmt.Printf("Available Commands:\n")
// Filter out "completion" command
availableCommands := make([]*cobra.Command, 0)
for _, subCmd := range cmd.Commands() {
if subCmd.Name() != "completion" {
availableCommands = append(availableCommands, subCmd)
}
}
for _, subCmd := range availableCommands {
fmt.Printf(" %-15s%s\n", subCmd.Name(), subCmd.Short)
}
fmt.Printf("\nFlags:\n")
cmd.Flags().VisitAll(func(flag *pflag.Flag) {
fmt.Printf(" -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
})
fmt.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
},
}
}
func (cli *CLI) NewAddMeasurementCmd() *cobra.Command {
return &cobra.Command{
Use: "measurement",
Short: "Add measurement to the attestation policy file. The value should be in base64. The second parameter is attestation_policy.json file",
Example: "measurement <measurement> <attestation_policy.json>",
Args: cobra.ExactArgs(2),
Run: func(cmd *cobra.Command, args []string) {
if err := changeAttestationConfiguration(args[1], args[0], measurementLength, measurementField); err != nil {
printError(cmd, "Error could not change measurement data: %v ❌ ", err)
return
}
},
}
}
func (cli *CLI) NewAddHostDataCmd() *cobra.Command {
return &cobra.Command{
Use: "hostdata",
Short: "Add host data to the attestation policy file. The value should be in base64. The second parameter is attestation_policy.json file",
Example: "hostdata <host-data> <attestation_policy.json>",
Args: cobra.ExactArgs(2),
Run: func(cmd *cobra.Command, args []string) {
if err := changeAttestationConfiguration(args[1], args[0], hostDataLength, hostDataField); err != nil {
printError(cmd, "Error could not change host data: %v ❌ ", err)
return
}
},
}
}
func (cli *CLI) NewGCPAttestationPolicy() *cobra.Command {
return &cobra.Command{
Use: "gcp",
Short: "Get attestation policy for GCP CVM",
Example: `gcp <bin_vtmp_attestation_report_file> <vcpu_count>`,
Args: cobra.ExactArgs(2),
Run: func(cmd *cobra.Command, args []string) {
attestationBin, err := os.ReadFile(args[0])
if err != nil {
printError(cmd, "Error reading attestation report file: %v ❌ ", err)
return
}
vcpuCount, err := strconv.Atoi(args[1])
if err != nil {
printError(cmd, "Error converting vCPU count to integer: %v ❌ ", err)
return
}
attestation := &attest.Attestation{}
if err := proto.Unmarshal(attestationBin, attestation); err != nil {
printError(cmd, "Error unmarshaling attestation report: %v ❌ ", err)
return
}
attestationPB := attestation.GetSevSnpAttestation()
measurement, err := gcp.Extract384BitMeasurement(attestationPB)
if err != nil {
printError(cmd, "Error extracting 384-bit measurement: %v ❌ ", err)
return
}
launchEndorsement, err := gcp.GetLaunchEndorsement(cmd.Context(), measurement)
if err != nil {
printError(cmd, "Error getting launch endorsement: %v ❌ ", err)
return
}
attestationPolicy, err := gcp.GenerateAttestationPolicy(launchEndorsement, uint32(vcpuCount))
if err != nil {
printError(cmd, "Error generating attestation policy: %v ❌ ", err)
return
}
attestationPolicyJson, err := json.MarshalIndent(attestationPolicy, "", " ")
if err != nil {
printError(cmd, "Error marshaling attestation policy: %v ❌ ", err)
return
}
if err := os.WriteFile("attestation_policy.json", attestationPolicyJson, filePermission); err != nil {
printError(cmd, "Error writing attestation policy file: %v ❌ ", err)
return
}
cmd.Println("Attestation policy file generated successfully ✅")
},
}
}
func (cli *CLI) NewDownloadGCPOvmfFile() *cobra.Command {
return &cobra.Command{
Use: "download",
Short: "Download GCP OVMF file",
Example: `download <bin_vtmp_attestation_report_file>`,
Args: cobra.ExactArgs(1),
Run: func(cmd *cobra.Command, args []string) {
attestationBin, err := os.ReadFile(args[0])
if err != nil {
printError(cmd, "Error reading attestation report file: %v ❌ ", err)
return
}
attestation := &attest.Attestation{}
if err := proto.Unmarshal(attestationBin, attestation); err != nil {
printError(cmd, "Error unmarshaling attestation report: %v ❌ ", err)
return
}
attestationPB := attestation.GetSevSnpAttestation()
measurement, err := gcp.Extract384BitMeasurement(attestationPB)
if err != nil {
printError(cmd, "Error extracting 384-bit measurement: %v ❌ ", err)
return
}
launchEndorsement, err := gcp.GetLaunchEndorsement(cmd.Context(), measurement)
if err != nil {
printError(cmd, "Error getting launch endorsement: %v ❌ ", err)
return
}
ovmf, err := gcp.DownloadOvmfFile(cmd.Context(), fmt.Sprintf("%x", launchEndorsement.Digest))
if err != nil {
printError(cmd, "Error downloading OVMF file: %v ❌ ", err)
return
}
sum384 := sha512.Sum384(ovmf)
if !bytes.Equal(sum384[:], launchEndorsement.Digest) {
printError(cmd, "Error OVMF file does not match the measurement: %v ❌ ", fmt.Errorf("digest mismatch"))
} else {
cmd.Println("OVMF firmware in vm is unmodified ✅")
}
if err := os.WriteFile("ovmf.fd", ovmf, filePermission); err != nil {
printError(cmd, "Error writing OVMF file: %v ❌ ", err)
return
}
cmd.Println("OVMF file downloaded successfully ✅")
},
}
}
func (cli *CLI) NewAzureAttestationPolicy() *cobra.Command {
cmd := &cobra.Command{
Use: "azure",
Short: "Get attestation policy for Azure CVM",
Example: `azure <azure_maa_token_file> <product_name>`,
Args: cobra.ExactArgs(2),
Run: func(cmd *cobra.Command, args []string) {
token, err := os.ReadFile(args[0])
if err != nil {
printError(cmd, "Error reading attestation report file: %v ❌ ", err)
return
}
product := args[1]
config, err := azure.GenerateAttestationPolicy(string(token), product, policy)
if err != nil {
printError(cmd, "Error generating attestation policy: %v ❌ ", err)
return
}
attestationPolicyJson, err := json.MarshalIndent(&config, "", " ")
if err != nil {
printError(cmd, "Error marshaling attestation policy: %v ❌ ", err)
return
}
if err := os.WriteFile("attestation_policy.json", attestationPolicyJson, filePermission); err != nil {
printError(cmd, "Error writing attestation policy file: %v ❌ ", err)
return
}
cmd.Println("Attestation policy file generated successfully ✅")
},
}
cmd.Flags().Uint64Var(
&policy,
"policy",
policy,
"Policy of the guest CVM",
)
return cmd
}
func changeAttestationConfiguration(fileName, base64Data string, expectedLength int, field fieldType) error {
data, err := base64.StdEncoding.DecodeString(base64Data)
if err != nil {
return errDecode
}
if len(data) != expectedLength {
return errDataLength
}
ac := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
f, err := os.ReadFile(fileName)
if err != nil {
return errors.Wrap(errReadingAttestationPolicyFile, err)
}
if err = attestation.ReadAttestationPolicyFromByte(f, &ac); err != nil {
return errors.Wrap(errUnmarshalJSON, err)
}
if ac.Config.Policy == nil {
ac.Config.Policy = &check.Policy{}
}
switch field {
case measurementField:
ac.Config.Policy.Measurement = data
case hostDataField:
ac.Config.Policy.HostData = data
default:
return errAttestationPolicyField
}
fileJson, err := json.MarshalIndent(&ac, "", " ")
if err != nil {
return errors.Wrap(errMarshalJSON, err)
}
if err = os.WriteFile(fileName, fileJson, filePermission); err != nil {
return errors.Wrap(errWriteFile, err)
}
return nil
}
-132
View File
@@ -1,132 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"encoding/base64"
"encoding/json"
"os"
"testing"
"github.com/google/go-sev-guest/proto/check"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/ultravioletrs/cocos/pkg/attestation"
)
func TestChangeAttestationConfiguration(t *testing.T) {
tmpfile, err := os.CreateTemp("", "attestation_policy.json")
require.NoError(t, err)
defer os.Remove(tmpfile.Name())
initialConfig := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
initialJSON, err := json.Marshal(initialConfig)
require.NoError(t, err)
err = os.WriteFile(tmpfile.Name(), initialJSON, 0o644)
require.NoError(t, err)
tests := []struct {
name string
base64Data string
expectedLength int
field fieldType
expectError bool
errorType error
}{
{
name: "Valid Measurement",
base64Data: base64.StdEncoding.EncodeToString(make([]byte, measurementLength)),
expectedLength: measurementLength,
field: measurementField,
expectError: false,
},
{
name: "Valid Host Data",
base64Data: base64.StdEncoding.EncodeToString(make([]byte, hostDataLength)),
expectedLength: hostDataLength,
field: hostDataField,
expectError: false,
},
{
name: "Invalid Base64",
base64Data: "Invalid Base64",
expectedLength: measurementLength,
field: measurementField,
expectError: true,
errorType: errDecode,
},
{
name: "Invalid Data Length",
base64Data: base64.StdEncoding.EncodeToString(make([]byte, measurementLength-1)),
expectedLength: measurementLength,
field: measurementField,
expectError: true,
errorType: errDataLength,
},
{
name: "Invalid Field Type",
base64Data: base64.StdEncoding.EncodeToString(make([]byte, measurementLength)),
expectedLength: measurementLength,
field: fieldType(999),
expectError: true,
errorType: errAttestationPolicyField,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := changeAttestationConfiguration(tmpfile.Name(), tt.base64Data, tt.expectedLength, tt.field)
if tt.expectError {
assert.Error(t, err)
assert.ErrorIs(t, err, tt.errorType)
} else {
assert.NoError(t, err)
content, err := os.ReadFile(tmpfile.Name())
require.NoError(t, err)
ap := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
err = attestation.ReadAttestationPolicyFromByte(content, &ap)
require.NoError(t, err)
decodedData, _ := base64.StdEncoding.DecodeString(tt.base64Data)
if tt.field == measurementField {
assert.Equal(t, decodedData, ap.Config.Policy.Measurement)
} else if tt.field == hostDataField {
assert.Equal(t, decodedData, ap.Config.Policy.HostData)
}
}
})
}
}
func TestNewAttestationPolicyCmd(t *testing.T) {
cli := &CLI{}
cmd := cli.NewAttestationPolicyCmd()
assert.Equal(t, "policy [command]", cmd.Use)
assert.Equal(t, "Change attestation policy", cmd.Short)
assert.NotNil(t, cmd.Run)
}
func TestNewAddMeasurementCmd(t *testing.T) {
cli := &CLI{}
cmd := cli.NewAddMeasurementCmd()
assert.Equal(t, "measurement", cmd.Use)
assert.Equal(t, "Add measurement to the attestation policy file. The value should be in base64. The second parameter is attestation_policy.json file", cmd.Short)
assert.Equal(t, "measurement <measurement> <attestation_policy.json>", cmd.Example)
assert.NotNil(t, cmd.Run)
}
func TestNewAddHostDataCmd(t *testing.T) {
cli := &CLI{}
cmd := cli.NewAddHostDataCmd()
assert.Equal(t, "hostdata", cmd.Use)
assert.Equal(t, "Add host data to the attestation policy file. The value should be in base64. The second parameter is attestation_policy.json file", cmd.Short)
assert.Equal(t, "hostdata <host-data> <attestation_policy.json>", cmd.Example)
assert.NotNil(t, cmd.Run)
}
-709
View File
@@ -1,709 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bytes"
"encoding/hex"
"encoding/json"
"fmt"
"os"
"testing"
"github.com/absmach/magistrala/pkg/errors"
"github.com/google/go-sev-guest/abi"
"github.com/google/go-sev-guest/proto/check"
"github.com/google/go-sev-guest/proto/sevsnp"
"github.com/spf13/cobra"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
)
func TestNewAttestationCmd(t *testing.T) {
mockSDK := new(mocks.SDK)
cli := &CLI{agentSDK: mockSDK}
cmd := cli.NewAttestationCmd()
assert.Equal(t, "attestation [command]", cmd.Use)
assert.Equal(t, "Get and validate attestations", cmd.Short)
var buf bytes.Buffer
cmd.SetOut(&buf)
reportData := bytes.Repeat([]byte{0x01}, quoteprovider.Nonce)
mockSDK.On("Attestation", mock.Anything, [quoteprovider.Nonce]byte(reportData), mock.Anything).Return(nil)
cmd.SetArgs([]string{hex.EncodeToString(reportData)})
err := cmd.Execute()
assert.NoError(t, err)
assert.Contains(t, buf.String(), "Get and validate attestations")
}
func TestNewGetAttestationCmd(t *testing.T) {
validattestation, err := os.ReadFile("../attestation.bin")
require.NoError(t, err)
teeNonce := hex.EncodeToString(bytes.Repeat([]byte{0x00}, quoteprovider.Nonce))
vtpmNonce := hex.EncodeToString(bytes.Repeat([]byte{0x00}, vtpm.Nonce))
tokenNonce := hex.EncodeToString(bytes.Repeat([]byte{0x00}, vtpm.Nonce))
testCases := []struct {
name string
args []string
mockResponse []byte
mockError error
expectedErr string
expectedOut string
}{
{
name: "successful SNP attestation retrieval",
args: []string{"snp", "--tee", teeNonce},
mockResponse: []byte("mock attestation"),
mockError: nil,
expectedOut: "Attestation result 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!",
},
{
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!",
},
{
name: "missing vTPM nonce",
args: []string{"snp-vtpm", "--tee", teeNonce},
mockResponse: []byte("mock attestation"),
mockError: nil,
expectedOut: "vTPM nonce must be defined for vTPM attestation",
},
{
name: "missing TEE nonce",
args: []string{"snp-vtpm", "--vtpm", vtpmNonce},
mockResponse: []byte("mock attestation"),
mockError: nil,
expectedOut: "TEE nonce must be defined for SEV-SNP attestation",
},
{
name: "invalid report data size",
args: []string{"snp", "--tee", hex.EncodeToString(bytes.Repeat([]byte{0x00}, 65))},
mockResponse: nil,
mockError: errors.New("error"),
expectedErr: "nonce must be a hex encoded string of length lesser or equal 64 bytes",
},
{
name: "invalid vTPM data size",
args: []string{"vtpm", "--vtpm", hex.EncodeToString(bytes.Repeat([]byte{0x00}, 33))},
mockResponse: nil,
mockError: errors.New("error"),
expectedErr: "vTPM nonce must be a hex encoded string of length lesser or equal 32 bytes",
},
{
name: "invalid arguments",
args: []string{"invalid"},
mockResponse: nil,
mockError: errors.New("error"),
expectedErr: "Bad attestation type: invalid argument ",
},
{
name: "failed to get attestation",
args: []string{"snp", "--tee", teeNonce},
mockResponse: nil,
mockError: errors.New("error"),
expectedErr: "Failed to get attestation due to error",
},
{
name: "Textproto report error",
args: []string{"snp", "--tee", teeNonce, "--reporttextproto"},
mockResponse: []byte("mock attestation"),
mockError: nil,
expectedErr: "Fetching SEV-SNP attestation report\nError converting SNP attestation to JSON: attestation contents too small : attestation contents too small (0x10 bytes). Want at least 0x4a0 bytes ❌\n",
},
{
name: "successful Textproto report",
args: []string{"snp", "--tee", teeNonce, "--reporttextproto"},
mockResponse: validattestation,
mockError: nil,
expectedOut: "Attestation result retrieved and saved successfully!",
},
{
name: "connection error",
args: []string{"snp", "--tee", teeNonce},
mockResponse: nil,
mockError: errors.New("failed to connect to agent"),
expectedErr: "Failed to connect to agent",
},
{
name: "successful Azure token retrieval",
args: []string{"azure-token", "--token", tokenNonce},
mockResponse: []byte("eyJhbGciOiAiUlMyNTYifQ.eyJzdWIiOiAidGVzdC11c2VyIn0.signature"),
mockError: nil,
expectedOut: "Fetching Azure token\nAttestation result 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",
},
{
name: "invalid token nonce size",
args: []string{"azure-token", "--token", hex.EncodeToString(bytes.Repeat([]byte{0x00}, 33))},
mockResponse: nil,
mockError: errors.New("error"),
expectedErr: "Fetching Azure token\nvTPM nonce must be a hex encoded string of length lesser or equal 32 bytes ❌ \n",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
t.Cleanup(func() {
os.Remove(attestationFilePath)
os.Remove(attestationReportJson)
})
mockSDK := new(mocks.SDK)
cli := &CLI{agentSDK: mockSDK}
if tc.name == "connection error" {
cli.connectErr = errors.New("failed to connect to agent")
}
cmd := cli.NewGetAttestationCmd()
var buf bytes.Buffer
cmd.SetOut(&buf)
mockSDK.On("Attestation", mock.Anything, [quoteprovider.Nonce]byte(bytes.Repeat([]byte{0x00}, quoteprovider.Nonce)), [vtpm.Nonce]byte(bytes.Repeat([]byte{0x00}, vtpm.Nonce)), mock.Anything, mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
_, err := args.Get(4).(*os.File).Write(tc.mockResponse)
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) {
_, err := args.Get(3).(*os.File).Write(tc.mockResponse)
require.NoError(t, err)
})
cmd.SetArgs(tc.args)
err := cmd.Execute()
if tc.expectedErr != "" {
assert.Contains(t, buf.String(), tc.expectedErr)
} else {
assert.NoError(t, err)
assert.Contains(t, buf.String(), tc.expectedOut)
}
})
}
}
func TestNewValidateAttestationValidationCmdDefaults(t *testing.T) {
cli := &CLI{}
cmd := cli.NewValidateAttestationValidationCmd()
assert.Equal(t, "validate", cmd.Use)
expectedMessage := fmt.Sprintf("Validate and verify attestation information. You can define the confidential computing cloud provider (%s, %s, %s; %s is the default) and can choose from 3 modes: %s, %s and %s. Default mode is %s.", CCNone, CCAzure, CCGCP, CCNone, SNP, VTPM, SNPvTPM, SNP)
assert.Equal(t, expectedMessage, cmd.Short)
assert.Equal(t, fmt.Sprint(defaultMinimumTcb), cmd.Flag("minimum_tcb").Value.String())
assert.Equal(t, fmt.Sprint(defaultMinimumLaunchTcb), cmd.Flag("minimum_lauch_tcb").Value.String())
assert.Equal(t, fmt.Sprint(defaultGuestPolicy), cmd.Flag("guest_policy").Value.String())
assert.Equal(t, fmt.Sprint(defaultMinimumGuestSvn), cmd.Flag("minimum_guest_svn").Value.String())
assert.Equal(t, fmt.Sprint(defaultMinimumBuild), cmd.Flag("minimum_build").Value.String())
assert.Equal(t, defaultCheckCrl, cmd.Flag("check_crl").Value.String() == "true")
assert.Equal(t, fmt.Sprint(defaultTimeout), cmd.Flag("timeout").Value.String())
assert.Equal(t, fmt.Sprint(defaultMaxRetryDelay), cmd.Flag("max_retry_delay").Value.String())
}
func TestNewValidateAttestationValidationCmd(t *testing.T) {
cli := &CLI{}
cmd := cli.NewValidateAttestationValidationCmd()
t.Run("missing attestation report file path", func(t *testing.T) {
err := cmd.Execute()
assert.Error(t, err)
assert.Equal(t, "please pass the attestation report file path", err.Error())
})
t.Run("unknown mode", func(t *testing.T) {
cmd.SetArgs([]string{attestationFilePath, "--mode=invalid"})
err := cmd.Execute()
assert.Error(t, err)
assert.Contains(t, err.Error(), "unknown mode")
})
t.Run("snp mode with missing flags", func(t *testing.T) {
cmd.SetArgs([]string{attestationFilePath, "--mode=snp"})
err := cmd.Execute()
assert.Error(t, err)
assert.Contains(t, err.Error(), "required flag(s) \"product\", \"report_data\" not set")
})
t.Run("vtpm mode with missing flags", func(t *testing.T) {
cmd.SetArgs([]string{vtpmFilePath, "--mode=vtpm"})
err := cmd.Execute()
assert.Error(t, err)
assert.Contains(t, err.Error(), "required flag(s) \"format\", \"nonce\", \"output\", \"product\", \"report_data\" not set")
})
t.Run("snp-vtpm mode with missing flags", func(t *testing.T) {
cmd.SetArgs([]string{vtpmFilePath, "--mode=snp-vtpm"})
err := cmd.Execute()
assert.Error(t, err)
assert.Contains(t, err.Error(), "required flag(s) \"format\", \"nonce\", \"output\", \"product\", \"report_data\" not set")
})
t.Run("valid snp mode execution", func(t *testing.T) {
cli := CLI{}
cmd := cli.NewValidateAttestationValidationCmd()
cmd.RunE = func(_ *cobra.Command, _ []string) error {
t.Log("Mock RunE executed instead of sevsnpverify")
return nil
}
cmd.SetArgs([]string{
"../attestation.bin",
"--mode=snp",
"--report_data=" +
"11223344556677889900aabbccddeeff11223344556677889900aabbccddeeff" +
"11223344556677889900aabbccddeeff11223344556677889900aabbccddeeff",
"--product=Milan",
})
err := cmd.PreRunE(cmd, []string{"../attestation.bin"})
assert.NoError(t, err)
})
t.Run("valid vtpm mode execution", func(t *testing.T) {
cli := CLI{}
cmd := cli.NewValidateAttestationValidationCmd()
cmd.RunE = func(_ *cobra.Command, _ []string) error {
t.Log("Mock RunE executed instead of vtpmverify")
return nil
}
cmd.SetArgs([]string{vtpmFilePath, "--mode=vtpm", "--nonce=123abc", "--format=binarypb", "--output=some_output"})
err := cmd.PreRunE(cmd, []string{"../quote.dat"})
assert.NoError(t, err)
})
t.Run("valid snp-vtpm mode execution", func(t *testing.T) {
cli := CLI{}
cmd := cli.NewValidateAttestationValidationCmd()
cmd.RunE = func(_ *cobra.Command, _ []string) error {
t.Log("Mock RunE executed instead of vtpmSevSnpverify")
return nil
}
cmd.SetArgs([]string{vtpmFilePath, "--mode=snp-vtpm", "--nonce=123abc", "--format=textproto", "--output=some_output"})
err := cmd.PreRunE(cmd, []string{"../quote.dat"})
assert.NoError(t, err)
})
}
type MockMeasurement struct {
mock.Mock
}
func (m *MockMeasurement) Run(igvmBinaryPath string) ([]byte, error) {
args := m.Called(igvmBinaryPath)
return nil, args.Error(0)
}
func (m *MockMeasurement) Stop() error {
args := m.Called()
return args.Error(0)
}
func TestNewMeasureCmd_RunSuccess(t *testing.T) {
cliInstance := &CLI{}
mockMeasurement := new(MockMeasurement)
cliInstance.measurement = mockMeasurement
mockMeasurement.On("Run", "testfile.igvm").Return(nil)
cmd := cliInstance.NewMeasureCmd("fake_binary_path")
buf := new(bytes.Buffer)
cmd.SetOut(buf)
cmd.SetErr(buf)
cmd.SetArgs([]string{"testfile.igvm"})
err := cmd.Execute()
assert.NoError(t, err)
mockMeasurement.AssertExpectations(t)
}
func TestNewMeasureCmd_RunError(t *testing.T) {
cliInstance := &CLI{}
mockMeasurement := new(MockMeasurement)
cliInstance.measurement = mockMeasurement
expectedError := errors.New("mocked measurement error")
mockMeasurement.On("Run", "testfile.igvm").Return(expectedError)
cmd := cliInstance.NewMeasureCmd("fake_binary_path")
buf := new(bytes.Buffer)
cmd.SetOut(buf)
cmd.SetErr(buf)
cmd.SetArgs([]string{"testfile.igvm"})
err := cmd.Execute()
assert.Error(t, err)
assert.Equal(t, expectedError.Error(), err.Error())
mockMeasurement.AssertExpectations(t)
}
func TestParseConfig(t *testing.T) {
cfgString = ""
err := parseConfig()
assert.NoError(t, err)
assert.NotNil(t, cfg.RootOfTrust)
assert.NotNil(t, cfg.Policy)
cfgString = `{"rootOfTrust":{"product":"test_product"},"policy":{"minimumGuestSvn":1}}`
err = parseConfig()
assert.NoError(t, err)
assert.Equal(t, "test_product", cfg.RootOfTrust.Product)
assert.Equal(t, uint32(1), cfg.Policy.MinimumGuestSvn)
cfgString = `{"invalid_json"`
err = parseConfig()
assert.Error(t, err)
}
func TestParseHashes(t *testing.T) {
trustedAuthorHashes = []string{"0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"}
trustedIdKeyHashes = []string{"fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210"}
cfg = check.Config{}
if cfg.Policy == nil {
cfg.Policy = &check.Policy{}
}
err := parseHashes()
assert.NoError(t, err)
assert.Len(t, cfg.Policy.TrustedAuthorKeyHashes, 1)
assert.Len(t, cfg.Policy.TrustedIdKeyHashes, 1)
trustedAuthorHashes = []string{"invalid_hash"}
err = parseHashes()
assert.Error(t, err)
}
func TestParseFiles(t *testing.T) {
attestationFile = "test_attestation.bin"
authorKeyFile := "test_author_key.pem"
idKeyFile := "test_id_key.pem"
err := os.WriteFile(attestationFile, []byte("test attestation"), 0o644)
assert.NoError(t, err)
err = os.WriteFile(authorKeyFile, []byte("test author key"), 0o644)
assert.NoError(t, err)
err = os.WriteFile(idKeyFile, []byte("test id key"), 0o644)
assert.NoError(t, err)
trustedAuthorKeys = []string{authorKeyFile}
trustedIdKeys = []string{idKeyFile}
err = parseAttestationFile()
assert.NoError(t, err)
err = parseTrustedKeys()
assert.NoError(t, err)
assert.Equal(t, []byte("test attestation"), attestationRaw)
assert.Len(t, cfg.Policy.TrustedAuthorKeys, 1)
assert.Len(t, cfg.Policy.TrustedIdKeys, 1)
os.Remove(attestationFile)
os.Remove(authorKeyFile)
os.Remove(idKeyFile)
attestationFile = "non_existent_file.bin"
err = parseAttestationFile()
assert.Error(t, err)
}
func TestParseUints(t *testing.T) {
stepping = "10"
platformInfo = "0xFF"
cfg = check.Config{}
if cfg.Policy == nil {
cfg.Policy = &check.Policy{
Product: &sevsnp.SevProduct{},
}
}
err := parseUints()
assert.NoError(t, err)
assert.Equal(t, uint32(10), cfg.Policy.Product.MachineStepping.Value)
assert.Equal(t, uint64(255), cfg.Policy.PlatformInfo.Value)
stepping = "invalid"
err = parseUints()
assert.Error(t, err)
stepping = "10"
platformInfo = "invalid"
err = parseUints()
assert.Error(t, err)
}
func TestValidateInput(t *testing.T) {
cfg = check.Config{}
if cfg.Policy == nil {
cfg.Policy = &check.Policy{}
}
if cfg.RootOfTrust == nil {
cfg.RootOfTrust = &check.RootOfTrust{}
}
cfg.Policy.ReportData = make([]byte, 64)
cfg.Policy.HostData = make([]byte, 32)
cfg.Policy.FamilyId = make([]byte, 16)
cfg.Policy.ImageId = make([]byte, 16)
cfg.Policy.ReportId = make([]byte, 32)
cfg.Policy.ReportIdMa = make([]byte, 32)
cfg.Policy.Measurement = make([]byte, 48)
cfg.Policy.ChipId = make([]byte, 64)
err := validateInput()
assert.NoError(t, err)
cfg.Policy.ReportData = make([]byte, 32)
err = validateInput()
assert.Error(t, err)
}
func TestGetBase(t *testing.T) {
assert.Equal(t, 16, getBase("0xFF"))
assert.Equal(t, 8, getBase("0o77"))
assert.Equal(t, 2, getBase("0b1010"))
assert.Equal(t, 10, getBase("123"))
}
func TestAttestationToJSON(t *testing.T) {
validReport, err := os.ReadFile("../attestation.bin")
require.NoError(t, err)
tests := []struct {
name string
input []byte
err error
}{
{
name: "Valid report",
input: validReport,
err: nil,
},
{
name: "Invalid report size",
input: make([]byte, abi.ReportSize-1),
err: errReportSize,
},
{
name: "Nil input",
input: nil,
err: errReportSize,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := attesationToJSON(tt.input)
assert.True(t, errors.Contains(err, tt.err))
if tt.err != nil {
assert.Nil(t, got)
return
}
require.NotNil(t, got)
var js map[string]interface{}
err = json.Unmarshal(got, &js)
assert.NoError(t, err)
})
}
}
func TestAttestationFromJSON(t *testing.T) {
tests := []struct {
name string
input []byte
err error
validate func(t *testing.T, output []byte)
}{
{
name: "Valid JSON",
input: func() []byte {
att := &sevsnp.Attestation{
Report: &sevsnp.Report{
CurrentTcb: 1,
FamilyId: make([]byte, 16),
ImageId: make([]byte, 16),
ReportData: make([]byte, 64),
Measurement: make([]byte, 48),
HostData: make([]byte, 32),
IdKeyDigest: make([]byte, 48),
AuthorKeyDigest: make([]byte, 48),
ReportId: make([]byte, 32),
ReportIdMa: make([]byte, 32),
ChipId: make([]byte, 64),
Signature: make([]byte, 512),
},
}
data, err := json.Marshal(att)
require.NoError(t, err)
return data
}(),
err: nil,
validate: func(t *testing.T, output []byte) {
assert.NotEmpty(t, output)
},
},
{
name: "Invalid JSON",
input: []byte(`{"invalid": json`),
err: errors.New("invalid character 'j' looking for beginning of value"),
validate: func(t *testing.T, output []byte) {
assert.Nil(t, output)
},
},
{
name: "Empty input",
input: []byte{},
err: errors.New("unexpected end of JSON input"),
validate: func(t *testing.T, output []byte) {
assert.Nil(t, output)
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := attesationFromJSON(tt.input)
assert.True(t, errors.Contains(err, tt.err))
tt.validate(t, got)
})
}
}
func TestIsFileJSON(t *testing.T) {
tests := []struct {
name string
filename string
want bool
}{
{
name: "Valid JSON extension",
filename: "test.json",
want: true,
},
{
name: "Valid JSON extension with path",
filename: "/path/to/test.json",
want: true,
},
{
name: "Invalid extension",
filename: "test.txt",
want: false,
},
{
name: "No extension",
filename: "test",
want: false,
},
{
name: "JSON in filename",
filename: "json.txt",
want: false,
},
{
name: "Empty string",
filename: "",
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := isFileJSON(tt.filename)
assert.Equal(t, tt.want, got)
})
}
}
func TestRoundTrip(t *testing.T) {
originalReport, err := os.ReadFile("../attestation.bin")
require.NoError(t, err)
jsonData, err := attesationToJSON(originalReport)
require.NoError(t, err)
require.NotNil(t, jsonData)
roundTripReport, err := attesationFromJSON(jsonData)
require.NoError(t, err)
require.NotNil(t, roundTripReport)
}
func TestDecodeJWTToJSON(t *testing.T) {
tests := []struct {
name string
input []byte
err error
validate func(t *testing.T, output []byte)
}{
{
name: "Valid JWT",
input: []byte("eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9." +
"eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ." +
"SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"),
err: nil,
validate: func(t *testing.T, output []byte) {
assert.NotEmpty(t, output)
assert.Contains(t, string(output), `"header"`)
assert.Contains(t, string(output), `"payload"`)
},
},
{
name: "Invalid JWT - one part",
input: []byte("justonepart"),
err: fmt.Errorf("invalid JWT: must have at least 2 parts"),
validate: func(t *testing.T, output []byte) {
assert.Nil(t, output)
},
},
{
name: "Invalid Base64",
input: []byte("bad@@@.header"),
err: errors.New("illegal base64 data"),
validate: func(t *testing.T, output []byte) {
assert.Nil(t, output)
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := decodeJWTToJSON(tt.input)
if tt.err != nil {
assert.Error(t, err)
assert.Contains(t, err.Error(), tt.err.Error())
} else {
assert.NoError(t, err)
}
tt.validate(t, got)
})
}
}
-78
View File
@@ -1,78 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"fmt"
"os"
"path"
"github.com/google/go-sev-guest/abi"
"github.com/google/go-sev-guest/kds"
"github.com/google/go-sev-guest/proto/check"
"github.com/google/go-sev-guest/verify/trust"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/pkg/attestation"
)
const (
caBundleName = "ask_ark.pem"
filePermisionKeys = 0o766
)
func (cli *CLI) NewCABundleCmd(fileSavePath string) *cobra.Command {
return &cobra.Command{
Use: "ca-bundle",
Short: "Fetch AMD SEV-SNPs CA Bundle (ASK and ARK)",
Example: "ca-bundle <path_to_platform_info_json>",
Args: cobra.ExactArgs(1),
Run: func(cmd *cobra.Command, args []string) {
attestationConfiguration := attestation.Config{Config: &check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}, PcrConfig: &attestation.PcrConfig{}}
err := attestation.ReadAttestationPolicy(args[0], &attestationConfiguration)
if err != nil {
printError(cmd, "Error while reading manifest: %v ❌ ", err)
return
}
product := attestationConfiguration.Config.RootOfTrust.ProductLine
getter := trust.DefaultHTTPSGetter()
caURL := kds.ProductCertChainURL(abi.VcekReportSigner, product)
bundle, err := getter.Get(caURL)
if err != nil {
message := fmt.Sprintf("Error fetching ARK and ASK from AMD KDS for product: %s", product)
message += ", error: %v ❌ "
printError(cmd, message, err)
return
}
err = os.MkdirAll(path.Join(fileSavePath, product), filePermisionKeys)
if err != nil {
message := fmt.Sprintf("Error while creating directory for product name %s", product)
message += ", error: %v ❌ "
printError(cmd, message, err)
return
}
bundlePath := path.Join(fileSavePath, product, caBundleName)
if err = saveToFile(bundlePath, bundle); err != nil {
printError(cmd, "Error while saving ARK-ASK to file: %v ❌ ", err)
return
}
},
}
}
func saveToFile(fileSavePath string, content []byte) error {
file, err := os.OpenFile(fileSavePath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, filePermisionKeys)
if err != nil {
return err
}
if _, err := file.Write(content); err != nil {
return err
}
return nil
}
-59
View File
@@ -1,59 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bytes"
"os"
"path"
"testing"
"github.com/stretchr/testify/assert"
)
func TestNewCABundleCmd(t *testing.T) {
cli := &CLI{}
tempDir, err := os.MkdirTemp("", "ca-bundle-test")
assert.NoError(t, err)
defer os.RemoveAll(tempDir)
manifestContent := []byte(`{"root_of_trust": {"product_line": "Milan"}}`)
manifestPath := path.Join(tempDir, "manifest.json")
err = os.WriteFile(manifestPath, manifestContent, 0o644)
assert.NoError(t, err)
cmd := cli.NewCABundleCmd(tempDir)
cmd.SetArgs([]string{manifestPath})
output := &bytes.Buffer{}
cmd.SetOutput(output)
err = cmd.Execute()
assert.NoError(t, err)
expectedFilePath := path.Join(tempDir, "Milan", caBundleName)
_, err = os.Stat(expectedFilePath)
assert.NoError(t, err)
content, err := os.ReadFile(expectedFilePath)
assert.NoError(t, err)
assert.NotNil(t, content)
}
func TestSaveToFile(t *testing.T) {
tempDir, err := os.MkdirTemp("", "save-to-file-test")
assert.NoError(t, err)
defer os.RemoveAll(tempDir)
filePath := path.Join(tempDir, "test-file.txt")
content := []byte("test content")
err = saveToFile(filePath, content)
assert.NoError(t, err)
savedContent, err := os.ReadFile(filePath)
assert.NoError(t, err)
assert.Equal(t, content, savedContent)
_, err = os.Stat(filePath)
assert.NoError(t, err)
}
-95
View File
@@ -1,95 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"encoding/base64"
"encoding/hex"
"encoding/json"
"os"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/internal"
"golang.org/x/crypto/sha3"
)
var (
ismanifest bool
toBase64 bool
)
func (cli *CLI) NewFileHashCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "checksum",
Short: "Compute the sha3-256 hash of a file",
Example: "checksum <file>",
Args: cobra.ExactArgs(1),
Run: func(cmd *cobra.Command, args []string) {
path := args[0]
if ismanifest {
hash, err := manifestChecksum(path)
if err != nil {
printError(cmd, "Error computing hash: %v ❌ ", err)
return
}
cmd.Println("Hash of manifest file:", hashOut(hash))
return
}
hash, err := internal.ChecksumHex(path)
if err != nil {
printError(cmd, "Error computing hash: %v ❌ ", err)
return
}
cmd.Println("Hash of file:", hashOut(hash))
},
}
cmd.Flags().BoolVarP(&ismanifest, "manifest", "m", false, "Compute the hash of the manifest file")
cmd.Flags().BoolVarP(&toBase64, "base64", "b", false, "Output the hash in base64")
return cmd
}
func manifestChecksum(path string) (string, error) {
file, err := os.ReadFile(path)
if err != nil {
return "", err
}
var cmp agent.Computation
if err := json.Unmarshal(file, &cmp); err != nil {
return "", err
}
jsonBytes, err := json.Marshal(cmp)
if err != nil {
return "", err
}
sum := sha3.Sum256(jsonBytes)
return hex.EncodeToString(sum[:]), nil
}
func hashOut(hashHex string) string {
if toBase64 {
return hexToBase64(hashHex)
}
return hashHex
}
func hexToBase64(hexStr string) string {
decoded, err := hex.DecodeString(hexStr)
if err != nil {
return ""
}
return base64.StdEncoding.EncodeToString(decoded)
}
-230
View File
@@ -1,230 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bytes"
"fmt"
"os"
"strings"
"testing"
"github.com/stretchr/testify/assert"
)
func TestNewFileHashCmd(t *testing.T) {
cli := &CLI{}
cmd := cli.NewFileHashCmd()
if cmd.Use != "checksum" {
t.Errorf("Expected Use to be 'checksum', got %s", cmd.Use)
}
if cmd.Short != "Compute the sha3-256 hash of a file" {
t.Errorf("Expected Short to be 'Compute the sha3-256 hash of a file', got %s", cmd.Short)
}
if cmd.Example != "checksum <file>" {
t.Errorf("Expected Example to be 'checksum <file>', got %s", cmd.Example)
}
}
func TestNewFileHashCmdRun(t *testing.T) {
testCases := []struct {
name string
isManifest bool
toBase64 bool
expectedOut string
expectedErr string
}{
{
name: "Valid file",
isManifest: false,
toBase64: false,
expectedOut: "Hash of file:",
expectedErr: "",
},
{
name: "Valid manifest file",
isManifest: true,
toBase64: false,
expectedOut: "Hash of manifest file:",
expectedErr: "",
},
{
name: "Valid file with base64 output",
isManifest: false,
toBase64: true,
expectedOut: "Hash of file:",
expectedErr: "",
},
{
name: "Non-existent file",
isManifest: false,
toBase64: false,
expectedOut: "Error computing hash:",
expectedErr: "",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
cli := &CLI{}
cmd := cli.NewFileHashCmd()
var output bytes.Buffer
cmd.SetOut(&output)
cmd.SetErr(&output)
err := cmd.Flags().Set("manifest", fmt.Sprint(tc.isManifest))
assert.Nil(t, err)
err = cmd.Flags().Set("base64", fmt.Sprint(tc.toBase64))
assert.Nil(t, err)
if tc.name == "Non-existent file" {
cmd.SetArgs([]string{"non_existent_file.txt"})
} else {
content := []byte("{}")
tmpfile, err := os.CreateTemp("", "example")
if err != nil {
t.Fatal(err)
}
defer os.Remove(tmpfile.Name())
if _, err := tmpfile.Write(content); err != nil {
t.Fatal(err)
}
if err := tmpfile.Close(); err != nil {
t.Fatal(err)
}
cmd.SetArgs([]string{tmpfile.Name()})
}
err = cmd.Execute()
if err != nil {
t.Fatalf("Error executing command: %v", err)
}
out := output.String()
if !strings.Contains(out, tc.expectedOut) {
t.Errorf("Expected output to contain '%s', got '%s'", tc.expectedOut, out)
}
if tc.expectedErr != "" && !strings.Contains(out, tc.expectedErr) {
t.Errorf("Expected output to contain '%s', got '%s'", tc.expectedErr, out)
}
})
}
}
func TestManifestChecksum(t *testing.T) {
testCases := []struct {
name string
jsonContent string
expectedSum string
}{
{
name: "Valid manifest file",
jsonContent: `{
"id": "1234",
"name": "Example Computation",
"description": "This is an example computation"
}`,
expectedSum: "a99683e4d22ba54cefa51aa49fb2e97a92b828c088395992ddff16a6236f3299",
},
{
name: "Invalid JSON",
jsonContent: `{`,
expectedSum: "",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
f, err := os.CreateTemp("", "test")
assert.Nil(t, err)
t.Cleanup(func() {
os.Remove(f.Name())
})
_, err = f.WriteString(tc.jsonContent)
assert.NoError(t, err)
err = f.Close()
assert.Nil(t, err)
hash, err := manifestChecksum(f.Name())
if tc.expectedSum == "" && err == nil {
t.Errorf("Expected error, got nil")
}
if tc.expectedSum != "" && err != nil {
t.Errorf("Unexpected error: %v", err)
}
if hash != tc.expectedSum {
t.Errorf("Expected hash %s, got %s", tc.expectedSum, hash)
}
})
}
}
func TestHexToBase64(t *testing.T) {
testCases := []struct {
name string
hexInput string
expectedOut string
}{
{
name: "Valid hex input",
hexInput: "48656c6c6f",
expectedOut: "SGVsbG8=",
},
{
name: "Invalid hex input",
hexInput: "invalid-hex",
expectedOut: "",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
out := hexToBase64(tc.hexInput)
if out != tc.expectedOut {
t.Errorf("Expected %s, got %s", tc.expectedOut, out)
}
})
}
}
func TestHashOut(t *testing.T) {
testCases := []struct {
name string
hashHex string
toBase64 bool
expectedOut string
}{
{
name: "Hex output",
hashHex: "48656c6c6f",
toBase64: false,
expectedOut: "48656c6c6f",
},
{
name: "Base64 output",
hashHex: "48656c6c6f",
toBase64: true,
expectedOut: "SGVsbG8=",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
toBase64 = tc.toBase64
out := hashOut(tc.hashHex)
if out != tc.expectedOut {
t.Errorf("Expected %s, got %s", tc.expectedOut, out)
}
})
}
}
+16 -92
View File
@@ -3,116 +3,40 @@
package cli
import (
"context"
"crypto/x509"
"encoding/pem"
"log"
"os"
"path"
"github.com/absmach/magistrala/pkg/errors"
"github.com/fatih/color"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/internal"
"google.golang.org/grpc/metadata"
)
var decompressDataset bool
func (cli *CLI) NewDatasetsCmd() *cobra.Command {
cmd := &cobra.Command{
return &cobra.Command{
Use: "data",
Short: "Upload a dataset",
Example: "data <dataset_path> <private_key_file_path>",
Args: cobra.ExactArgs(2),
Short: "Upload a dataset CSV file",
Example: "data <dataset.csv> <id> <provider>",
Args: cobra.ExactArgs(3),
Run: func(cmd *cobra.Command, args []string) {
if cli.connectErr != nil {
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
return
}
datasetFile := args[0]
datasetPath := args[0]
log.Println("Uploading dataset CSV:", datasetFile)
cmd.Println("Uploading dataset:", datasetPath)
f, err := os.Stat(datasetPath)
dataset, err := os.ReadFile(datasetFile)
if err != nil {
printError(cmd, "Error reading dataset file: %v", err)
return
log.Fatalf("Error reading dataset file: %v", err)
}
var dataset *os.File
if f.IsDir() {
dataset, err = internal.ZipDirectoryToTempFile(datasetPath)
if err != nil {
printError(cmd, "Error zipping dataset directory: %v ❌ ", err)
return
}
defer dataset.Close()
defer os.Remove(dataset.Name())
} else {
dataset, err = os.Open(datasetPath)
if err != nil {
printError(cmd, "Error reading dataset file: %v ❌ ", err)
return
}
defer dataset.Close()
dataReq := agent.Dataset{
Dataset: dataset,
ID: args[1],
Provider: args[2],
}
privKeyFile, err := os.ReadFile(args[1])
if err != nil {
printError(cmd, "Error reading private key file: %v ❌ ", err)
return
if err := cli.agentSDK.Data(cmd.Context(), dataReq); err != nil {
log.Fatalf("Error uploading dataset: %v", err)
}
pemBlock, _ := pem.Decode(privKeyFile)
privKey, err := decodeKey(pemBlock)
if err != nil {
printError(cmd, "Error decoding private key: %v ❌ ", err)
return
}
ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
if err := cli.agentSDK.Data(addDatasetMetadata(ctx), dataset, path.Base(datasetPath), privKey); err != nil {
printError(cmd, "Failed to upload dataset due to error: %v ❌ ", err)
return
}
cmd.Println(color.New(color.FgGreen).Sprint("Successfully uploaded dataset! ✔ "))
log.Println("Successfully uploaded dataset")
},
}
cmd.Flags().BoolVarP(&decompressDataset, "decompress", "d", false, "Decompress the dataset on agent")
return cmd
}
func decodeKey(b *pem.Block) (interface{}, error) {
if b == nil {
return nil, errors.New("error decoding key")
}
switch b.Type {
case rsaKeyType:
privKey, err := x509.ParsePKCS8PrivateKey(b.Bytes)
if err != nil {
privKey, err = x509.ParsePKCS1PrivateKey(b.Bytes)
if err != nil {
return nil, err
}
}
return privKey, nil
case ecdsaKeyType:
privKey, err := x509.ParseECPrivateKey(b.Bytes)
if err != nil {
return nil, err
}
return privKey, nil
default:
return nil, errors.New("error decoding key")
}
}
func addDatasetMetadata(ctx context.Context) context.Context {
return agent.DecompressToContext(ctx, decompressDataset)
}
-158
View File
@@ -1,158 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bytes"
"errors"
"os"
"testing"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
)
func createTempDatasetFile(content string) (string, error) {
tmpFile, err := os.CreateTemp("", "dataset-*.txt")
if err != nil {
return "", err
}
defer tmpFile.Close()
_, err = tmpFile.WriteString(content)
if err != nil {
return "", err
}
return tmpFile.Name(), nil
}
func TestDatasetsCmd(t *testing.T) {
tests := []struct {
name string
setupMock func(*mocks.SDK)
setupFiles func() (string, error)
connectErr error
expectedOutput string
cleanup func(string, string)
}{
{
name: "successful upload",
setupMock: func(m *mocks.SDK) {
m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
},
setupFiles: func() (string, error) {
datasetFile, err := createTempDatasetFile("test dataset content")
if err != nil {
return "", err
}
err = generateRSAPrivateKeyFile(privateKeyFile)
return datasetFile, err
},
expectedOutput: "Successfully uploaded dataset",
cleanup: func(datasetFile, privateKeyFile string) {
os.Remove(datasetFile)
os.Remove(privateKeyFile)
},
},
{
name: "missing dataset file",
setupMock: func(m *mocks.SDK) {
m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
},
setupFiles: func() (string, error) {
return "", nil
},
expectedOutput: "Error reading dataset file",
},
{
name: "missing private key file",
setupMock: func(m *mocks.SDK) {
m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
},
setupFiles: func() (string, error) {
return createTempDatasetFile("test dataset content")
},
expectedOutput: "Error reading private key file",
cleanup: func(datasetFile, _ string) {
os.Remove(datasetFile)
},
},
{
name: "upload failure",
setupMock: func(m *mocks.SDK) {
m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
},
setupFiles: func() (string, error) {
datasetFile, err := createTempDatasetFile("test dataset content")
if err != nil {
return "", err
}
err = generateRSAPrivateKeyFile(privateKeyFile)
return datasetFile, err
},
expectedOutput: "Failed to upload dataset due to error",
cleanup: func(datasetFile, privateKeyFile string) {
os.Remove(datasetFile)
os.Remove(privateKeyFile)
},
},
{
name: "invalid private key",
setupMock: func(m *mocks.SDK) {
m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
},
setupFiles: func() (string, error) {
datasetFile, err := createTempDatasetFile("test dataset content")
if err != nil {
return "", err
}
err = os.WriteFile(privateKeyFile, []byte("invalid private key"), 0o644)
return datasetFile, err
},
expectedOutput: "Error decoding private key",
cleanup: func(datasetFile, privateKeyFile string) {
os.Remove(datasetFile)
os.Remove(privateKeyFile)
},
},
{
name: "connection error",
setupMock: func(m *mocks.SDK) {
},
setupFiles: func() (string, error) { return "", nil },
connectErr: errors.New("failed to connect to agent"),
expectedOutput: "Failed to connect to agent",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mockSDK := new(mocks.SDK)
if tt.setupMock != nil {
tt.setupMock(mockSDK)
}
testCLI := CLI{
agentSDK: mockSDK,
connectErr: tt.connectErr,
}
datasetFile, err := tt.setupFiles()
require.NoError(t, err)
cmd := testCLI.NewDatasetsCmd()
buf := new(bytes.Buffer)
cmd.SetOut(buf)
cmd.SetArgs([]string{datasetFile, privateKeyFile})
err = cmd.Execute()
require.NoError(t, err)
require.Contains(t, buf.String(), tt.expectedOutput)
if tt.cleanup != nil {
tt.cleanup(datasetFile, privateKeyFile)
}
})
}
}
-49
View File
@@ -1,49 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"github.com/absmach/magistrala/pkg/errors"
"github.com/fatih/color"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/agent/auth"
"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
var (
errAgentUnavailable = errors.New("agent is unavailable on the current address")
errDigitalSignatureVerificationFailed = errors.New("digital signature verification failed, check the provided public key")
)
func decodeErros(err error) error {
statusErr, ok := status.FromError(err)
if ok {
switch statusErr.Code() {
case codes.PermissionDenied:
return errDigitalSignatureVerificationFailed
case codes.Unavailable:
return errAgentUnavailable
case codes.Unknown:
return err
}
}
switch {
case errors.Contains(err, auth.ErrSignatureVerificationFailed):
return auth.ErrSignatureVerificationFailed
case errors.Contains(err, agent.ErrAgentServiceUnavailable):
return agent.ErrAgentServiceUnavailable
default:
return err
}
}
func printError(cmd *cobra.Command, message string, err error) {
if !Verbose {
err = decodeErros(err)
}
msg := color.New(color.FgRed).Sprintf(message, err)
cmd.Println(msg)
}
-110
View File
@@ -1,110 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bytes"
"errors"
"testing"
mgerrors "github.com/absmach/magistrala/pkg/errors"
"github.com/fatih/color"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/agent/auth"
"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
func TestDecodeErros(t *testing.T) {
tests := []struct {
name string
input error
expected error
}{
{
name: "Permission Denied",
input: status.Error(codes.PermissionDenied, "permission denied"),
expected: errDigitalSignatureVerificationFailed,
},
{
name: "Unavailable",
input: status.Error(codes.Unavailable, "service unavailable"),
expected: errAgentUnavailable,
},
{
name: "Unknown",
input: status.Error(codes.Unknown, "unknown error"),
expected: status.Error(codes.Unknown, "unknown error"),
},
{
name: "Signature Verification Failed",
input: mgerrors.Wrap(auth.ErrSignatureVerificationFailed, errors.New("wrapped error")),
expected: auth.ErrSignatureVerificationFailed,
},
{
name: "Other Error",
input: errors.New("other error"),
expected: errors.New("other error"),
},
{
name: "Agent Service Unavailable",
input: mgerrors.Wrap(agent.ErrAgentServiceUnavailable, errors.New("wrapped error")),
expected: agent.ErrAgentServiceUnavailable,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := decodeErros(tt.input)
if result.Error() != tt.expected.Error() {
t.Errorf("decodeErros(%v) = %v, want %v", tt.input, result, tt.expected)
}
})
}
}
func TestPrintError(t *testing.T) {
// Save the original color.NoColor value and restore it after the test
origNoColor := color.NoColor
color.NoColor = true
defer func() { color.NoColor = origNoColor }()
tests := []struct {
name string
message string
err error
verbose bool
expected string
}{
{
name: "Non-verbose mode",
message: "Error: %s",
err: status.Error(codes.PermissionDenied, "permission denied"),
verbose: false,
expected: "Error: digital signature verification failed, check the provided public key\n",
},
{
name: "Verbose mode",
message: "Error: %s",
err: status.Error(codes.PermissionDenied, "permission denied"),
verbose: true,
expected: "Error: rpc error: code = PermissionDenied desc = permission denied\n",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
Verbose = tt.verbose
cmd := &cobra.Command{}
buf := new(bytes.Buffer)
cmd.SetOut(buf)
printError(cmd, tt.message, tt.err)
if got := buf.String(); got != tt.expected {
t.Errorf("printError() output = %q, want %q", got, tt.expected)
}
})
}
}
-96
View File
@@ -1,96 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bufio"
"crypto/sha1"
"encoding/hex"
"os"
"strings"
"github.com/fatih/color"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
)
const (
imaMeasurementsFilename = "ima_measurements"
)
func (cli *CLI) NewIMAMeasurementsCmd() *cobra.Command {
return &cobra.Command{
Use: "ima-measurements",
Short: "Retrieve Linux IMA measurements file",
Example: "ima-measurements <optional_file_name>",
Run: func(cmd *cobra.Command, args []string) {
if cli.connectErr != nil {
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
return
}
cmd.Println("⏳ Retrieving computation Linux IMA measurements file")
filename := imaMeasurementsFilename
if len(args) >= 1 {
filename = args[0]
}
imaMeasurementsFile, err := os.Create(filename)
if err != nil {
printError(cmd, "Error creating imaMeasurements file: %v ❌ ", err)
return
}
defer imaMeasurementsFile.Close()
pcr10, err := cli.agentSDK.IMAMeasurements(cmd.Context(), imaMeasurementsFile)
if err != nil {
printError(cmd, "Error retrieving Linux IMA measurements file: %v ❌ ", err)
return
}
cmd.Println(color.New(color.FgGreen).Sprintf("Linux IMA measurements file retrieved and saved successfully as %s! PCR10 = %s ✔ ", filename, hex.EncodeToString(pcr10)))
calculatedPCR10 := make([]byte, vtpm.Hash1)
file, err := os.Open(filename)
if err != nil {
printError(cmd, "Failed to open file: %v ❌ ", err)
}
defer file.Close()
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := scanner.Text()
parts := strings.Fields(line)
if parts[0] != "10" {
continue
}
digestHex := parts[1]
if digestHex == strings.Repeat("0", 40) {
digestHex = strings.Repeat("f", 40)
}
digest, err := hex.DecodeString(digestHex)
if err != nil {
printError(cmd, "Failed to decode digest: %v ❌ ", err)
continue
}
hasher := sha1.New()
hasher.Write(calculatedPCR10)
hasher.Write(digest)
calculatedPCR10 = hasher.Sum(nil)
}
if hex.EncodeToString(pcr10) != hex.EncodeToString(calculatedPCR10) {
printError(cmd, "Measurements file not verified ❌ ", err)
} else {
cmd.Println(color.New(color.FgGreen).Sprintf("Measurements file verified!"))
}
},
}
}
-140
View File
@@ -1,140 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"crypto/ecdsa"
"crypto/ed25519"
"crypto/elliptic"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/pem"
"os"
"github.com/spf13/cobra"
)
const (
keyBitSize = 4096
rsaKeyType = "PRIVATE KEY"
ecdsaKeyType = "EC PRIVATE KEY"
ed25519KeyType = "PRIVATE KEY"
publicKeyType = "PUBLIC KEY"
publicKeyFile = "public.pem"
privateKeyFile = "private.pem"
ECDSA = "ecdsa"
ED25519 = "ed25519"
)
var KeyType string
func (cli *CLI) NewKeysCmd() *cobra.Command {
return &cobra.Command{
Use: "keys",
Short: "Generate a new public/private key pair",
Long: "Generates a new public/private key pair using an algorithm of the users choice.\n" +
"Supported algorithms are RSA, ecdsa, and ed25519.",
Example: "./build/cocos-cli keys -k rsa",
Args: cobra.ExactArgs(0),
Run: func(cmd *cobra.Command, args []string) {
switch KeyType {
case ECDSA:
privEcdsaKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
printError(cmd, "Error generating keys: %v ❌ ", err)
return
}
pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privEcdsaKey.PublicKey)
if err != nil {
printError(cmd, "Error marshalling public key: %v ❌ ", err)
return
}
if err := generateAndWriteKeys(privEcdsaKey, pubKeyBytes, ecdsaKeyType); err != nil {
printError(cmd, "Error generating and writing keys: %v ❌ ", err)
return
}
case ED25519:
pubEd25519Key, privEd25519Key, err := ed25519.GenerateKey(rand.Reader)
if err != nil {
printError(cmd, "Error generating keys: %v ❌ ", err)
return
}
pubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
if err != nil {
printError(cmd, "Error marshalling public key: %v ❌ ", err)
return
}
if err := generateAndWriteKeys(privEd25519Key, pubKey, ed25519KeyType); err != nil {
printError(cmd, "Error generating and writing keys: %v ❌ ", err)
return
}
default:
privKey, err := rsa.GenerateKey(rand.Reader, keyBitSize)
if err != nil {
printError(cmd, "Error generating keys: %v ❌ ", err)
return
}
pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
if err != nil {
printError(cmd, "Error marshalling public key: %v ❌ ", err)
return
}
if err := generateAndWriteKeys(privKey, pubKeyBytes, rsaKeyType); err != nil {
printError(cmd, "Error generating and writing keys: %v ❌ ", err)
return
}
}
cmd.Printf("Successfully generated public/private key pair of type: %s", KeyType)
},
}
}
func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType string) error {
privFile, err := os.Create(privateKeyFile)
if err != nil {
return err
}
defer privFile.Close()
var b []byte
switch privKey := privKey.(type) {
case *rsa.PrivateKey:
b = x509.MarshalPKCS1PrivateKey(privKey)
case *ecdsa.PrivateKey:
b, err = x509.MarshalECPrivateKey(privKey)
case ed25519.PrivateKey:
b, err = x509.MarshalPKCS8PrivateKey(privKey)
}
if err != nil {
return err
}
if err := pem.Encode(privFile, &pem.Block{
Type: keyType,
Bytes: b,
}); err != nil {
return err
}
pubFile, err := os.Create(publicKeyFile)
if err != nil {
return err
}
defer pubFile.Close()
if err := pem.Encode(pubFile, &pem.Block{
Type: publicKeyType,
Bytes: pubKeyBytes,
}); err != nil {
return err
}
return nil
}
-92
View File
@@ -1,92 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"crypto/ecdsa"
"crypto/ed25519"
"crypto/rsa"
"crypto/x509"
"encoding/pem"
"os"
"testing"
)
func TestNewKeysCmd(t *testing.T) {
cli := &CLI{}
cmd := cli.NewKeysCmd()
if cmd.Use != "keys" {
t.Errorf("Expected Use to be 'keys', got %s", cmd.Use)
}
if cmd.Short != "Generate a new public/private key pair" {
t.Errorf("Unexpected Short description: %s", cmd.Short)
}
}
func TestGenerateAndWriteKeys(t *testing.T) {
tests := []struct {
name string
keyType string
}{
{"RSA", "rsa"},
{"ECDSA", "ecdsa"},
{"ED25519", "ed25519"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
KeyType = tt.keyType
cmd := (&CLI{}).NewKeysCmd()
cmd.Run(cmd, []string{})
if _, err := os.Stat(privateKeyFile); os.IsNotExist(err) {
t.Errorf("Private key file was not created")
}
if _, err := os.Stat(publicKeyFile); os.IsNotExist(err) {
t.Errorf("Public key file was not created")
}
privKeyData, err := os.ReadFile(privateKeyFile)
if err != nil {
t.Fatalf("Failed to read private key file: %v", err)
}
privPem, _ := pem.Decode(privKeyData)
if privPem == nil {
t.Fatalf("Failed to decode private key PEM")
}
var privKey interface{}
switch tt.keyType {
case "rsa":
privKey, err = x509.ParsePKCS1PrivateKey(privPem.Bytes)
case "ecdsa":
privKey, err = x509.ParseECPrivateKey(privPem.Bytes)
case "ed25519":
privKey, err = x509.ParsePKCS8PrivateKey(privPem.Bytes)
}
if err != nil {
t.Fatalf("Failed to parse private key: %v", err)
}
switch tt.keyType {
case "rsa":
if _, ok := privKey.(*rsa.PrivateKey); !ok {
t.Errorf("Expected RSA private key, got %T", privKey)
}
case "ecdsa":
if _, ok := privKey.(*ecdsa.PrivateKey); !ok {
t.Errorf("Expected ECDSA private key, got %T", privKey)
}
case "ed25519":
if _, ok := privKey.(ed25519.PrivateKey); !ok {
t.Errorf("Expected ED25519 private key, got %T", privKey)
}
}
os.Remove(privateKeyFile)
os.Remove(publicKeyFile)
})
}
}
-146
View File
@@ -1,146 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"os"
"time"
"github.com/fatih/color"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/manager"
)
const (
serverURL = "server-url"
serverCA = "server-ca"
clientKey = "client-key"
clientCrt = "client-crt"
caUrl = "ca-url"
logLevel = "log-level"
ttlFlag = "ttl"
)
var (
agentCVMServerUrl string
agentCVMServerCA string
agentCVMClientKey string
agentCVMClientCrt string
agentCVMCaUrl string
agentLogLevel string
ttl time.Duration
)
func (c *CLI) NewCreateVMCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "create-vm",
Short: "Create a new virtual machine",
Example: `create-vm`,
Args: cobra.ExactArgs(0),
Run: func(cmd *cobra.Command, args []string) {
if err := c.InitializeManagerClient(cmd); err != nil {
printError(cmd, "Failed to connect to manager: %v ❌ ", c.connectErr)
return
}
defer c.Close()
createReq, err := loadCerts()
if err != nil {
printError(cmd, "Error loading certs: %v ❌ ", err)
return
}
createReq.AgentCvmServerUrl = agentCVMServerUrl
createReq.AgentLogLevel = agentLogLevel
createReq.AgentCvmCaUrl = agentCVMCaUrl
if ttl > 0 {
createReq.Ttl = ttl.String()
}
cmd.Println("🔗 Creating a new virtual machine")
res, err := c.managerClient.CreateVm(cmd.Context(), createReq)
if err != nil {
printError(cmd, "Error creating virtual machine: %v ❌ ", err)
return
}
cmd.Println(color.New(color.FgGreen).Sprintf("✅ Virtual machine created successfully with id %s and port %s", res.CvmId, res.ForwardedPort))
},
}
cmd.Flags().StringVar(&agentCVMServerUrl, serverURL, "", "CVM server URL")
cmd.Flags().StringVar(&agentCVMServerCA, serverCA, "", "CVM server CA")
cmd.Flags().StringVar(&agentCVMClientKey, clientKey, "", "CVM client key")
cmd.Flags().StringVar(&agentCVMClientCrt, clientCrt, "", "CVM client crt")
cmd.Flags().StringVar(&agentCVMCaUrl, agentCVMCaUrl, "", "CVM CA service URL")
cmd.Flags().StringVar(&agentLogLevel, logLevel, "", "Agent Log level")
cmd.Flags().DurationVar(&ttl, ttlFlag, 0, "TTL for the VM")
if err := cmd.MarkFlagRequired(serverURL); err != nil {
printError(cmd, "Error marking flag as required: %v ❌ ", err)
return cmd
}
return cmd
}
func (c *CLI) NewRemoveVMCmd() *cobra.Command {
return &cobra.Command{
Use: "remove-vm",
Short: "Remove a virtual machine",
Example: `remove-vm <cvm_id>`,
Args: cobra.ExactArgs(1),
Run: func(cmd *cobra.Command, args []string) {
if err := c.InitializeManagerClient(cmd); err == nil {
defer c.Close()
}
if c.connectErr != nil {
printError(cmd, "Failed to connect to manager: %v ❌ ", c.connectErr)
return
}
cmd.Println("🔗 Removing virtual machine")
_, err := c.managerClient.RemoveVm(cmd.Context(), &manager.RemoveReq{CvmId: args[0]})
if err != nil {
printError(cmd, "Error removing virtual machine: %v ❌ ", err)
return
}
cmd.Println(color.New(color.FgGreen).Sprintf("✅ Virtual machine removed successfully"))
},
}
}
func fileReader(path string) ([]byte, error) {
if path == "" {
return nil, nil
}
return os.ReadFile(path)
}
func loadCerts() (*manager.CreateReq, error) {
clientKey, err := fileReader(agentCVMClientKey)
if err != nil {
return nil, err
}
clientCrt, err := fileReader(agentCVMClientCrt)
if err != nil {
return nil, err
}
serverCA, err := fileReader(agentCVMServerCA)
if err != nil {
return nil, err
}
return &manager.CreateReq{
AgentCvmServerCaCert: serverCA,
AgentCvmClientKey: clientKey,
AgentCvmClientCert: clientCrt,
}, nil
}
+11 -44
View File
@@ -3,65 +3,32 @@
package cli
import (
"encoding/pem"
"log"
"os"
"github.com/fatih/color"
"github.com/spf13/cobra"
)
const (
resultFilePrefix = "results"
resultFileExt = ".zip"
resultfilename = "results.zip"
)
const resultFilePath = "result.bin"
func (cli *CLI) NewResultsCmd() *cobra.Command {
return &cobra.Command{
Use: "result",
Short: "Retrieve computation result file",
Example: "result <private_key_file_path> <optional_file_name.zip>",
Args: cobra.MinimumNArgs(1),
Use: "result",
Short: "Retrieve computation result file",
Args: cobra.ExactArgs(1),
Run: func(cmd *cobra.Command, args []string) {
if cli.connectErr != nil {
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
return
}
log.Println("Retrieving computation result file")
cmd.Println("⏳ Retrieving computation result file")
privKeyFile, err := os.ReadFile(args[0])
result, err := cli.agentSDK.Result(cmd.Context(), args[0])
if err != nil {
printError(cmd, "Error reading private key file: %v", err)
return
log.Fatalf("Error retrieving computation result: %v", err)
}
filename := resultfilename
if len(args) > 1 {
filename = args[1]
if err := os.WriteFile(resultFilePath, result, 0o644); err != nil {
log.Fatalf("Error saving computation result to %s: %v", resultFilePath, err)
}
pemBlock, _ := pem.Decode(privKeyFile)
privKey, err := decodeKey(pemBlock)
if err != nil {
printError(cmd, "Error decoding private key: %v ❌ ", err)
return
}
resultFile, err := os.Create(filename)
if err != nil {
printError(cmd, "Error creating result file: %v ❌ ", err)
return
}
defer resultFile.Close()
if err = cli.agentSDK.Result(cmd.Context(), privKey, resultFile); err != nil {
printError(cmd, "Error retrieving computation result: %v ❌ ", err)
return
}
cmd.Println(color.New(color.FgGreen).Sprintf("Computation result retrieved and saved successfully as %s! ✔ ", filename))
log.Println("Computation result retrieved and saved successfully!")
},
}
}
-177
View File
@@ -1,177 +0,0 @@
// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package cli
import (
"bytes"
"errors"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
)
const compResult = "Test computation result"
func TestResultsCmd_MultipleExecutions(t *testing.T) {
mockSDK := new(mocks.SDK)
mockSDK.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
_, err := args.Get(2).(*os.File).WriteString(compResult)
require.NoError(t, err)
})
testCLI := CLI{agentSDK: mockSDK}
err := generateRSAPrivateKeyFile(privateKeyFile)
require.NoError(t, err)
cmd := testCLI.NewResultsCmd()
buf := new(bytes.Buffer)
cmd.SetOut(buf)
cmd.SetArgs([]string{privateKeyFile})
for i := 0; i < 3; i++ {
err = cmd.Execute()
require.NoError(t, err)
require.Contains(t, buf.String(), "Computation result retrieved and saved successfully")
buf.Reset()
}
files, err := filepath.Glob("results*.zip")
require.NoError(t, err)
t.Cleanup(func() {
for _, file := range files {
os.Remove(file)
}
os.Remove(privateKeyFile)
})
}
func TestResultsCmd_InvalidPrivateKey(t *testing.T) {
mockSDK := new(mocks.SDK)
mockSDK.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
_, err := args.Get(2).(*os.File).WriteString(compResult)
require.NoError(t, err)
})
testCLI := CLI{agentSDK: mockSDK}
invalidPrivateKey, err := os.CreateTemp("", "invalid_private_key.pem")
require.NoError(t, err)
err = invalidPrivateKey.Close()
require.NoError(t, err)
t.Cleanup(func() {
err := os.Remove(invalidPrivateKey.Name())
require.NoError(t, err)
})
cmd := testCLI.NewResultsCmd()
buf := new(bytes.Buffer)
cmd.SetOut(buf)
cmd.SetArgs([]string{invalidPrivateKey.Name()})
err = cmd.Execute()
require.NoError(t, err)
require.Contains(t, buf.String(), "Error decoding private key")
mockSDK.AssertNotCalled(t, "Result", mock.Anything, mock.Anything, mock.Anything)
}
func TestResultsCmd(t *testing.T) {
tests := []struct {
name string
setupMock func(*mocks.SDK)
setupFiles func() (string, error)
connectErr error
expectedOutput string
cleanup func()
}{
{
name: "successful result retrieval",
setupMock: func(m *mocks.SDK) {
m.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
_, err := args.Get(2).(*os.File).WriteString(compResult)
require.NoError(t, err)
})
},
setupFiles: func() (string, error) {
return privateKeyFile, generateRSAPrivateKeyFile(privateKeyFile)
},
expectedOutput: "Computation result retrieved and saved successfully",
cleanup: func() {
files, _ := filepath.Glob("results*.zip")
for _, file := range files {
os.Remove(file)
}
os.Remove(privateKeyFile)
},
},
{
name: "missing private key file",
setupMock: func(m *mocks.SDK) {
m.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
_, err := args.Get(2).(*os.File).WriteString(compResult)
require.NoError(t, err)
})
},
setupFiles: func() (string, error) {
return "non_existent_private_key.pem", nil
},
expectedOutput: "Error reading private key file",
},
{
name: "result retrieval failure",
setupMock: func(m *mocks.SDK) {
m.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error retrieving computation result"))
},
setupFiles: func() (string, error) {
return privateKeyFile, generateRSAPrivateKeyFile(privateKeyFile)
},
expectedOutput: "error retrieving computation result",
cleanup: func() {
os.Remove(privateKeyFile)
},
},
{
name: "connection error",
setupMock: func(m *mocks.SDK) {
},
setupFiles: func() (string, error) { return "", nil },
connectErr: errors.New("failed to connect to agent"),
expectedOutput: "Failed to connect to agent",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mockSDK := new(mocks.SDK)
if tt.setupMock != nil {
tt.setupMock(mockSDK)
}
testCLI := CLI{
agentSDK: mockSDK,
connectErr: tt.connectErr,
}
file, err := tt.setupFiles()
require.NoError(t, err)
cmd := testCLI.NewResultsCmd()
buf := new(bytes.Buffer)
cmd.SetOut(buf)
cmd.SetArgs([]string{file})
err = cmd.Execute()
require.NoError(t, err)
require.Contains(t, buf.String(), tt.expectedOutput)
if tt.cleanup != nil {
tt.cleanup()
}
})
}
}
+4 -53
View File
@@ -3,64 +3,15 @@
package cli
import (
"context"
"github.com/spf13/cobra"
"github.com/ultravioletrs/cocos/manager"
"github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig"
"github.com/ultravioletrs/cocos/pkg/clients/grpc"
"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
managergrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/manager"
"github.com/ultravioletrs/cocos/pkg/sdk"
"github.com/ultravioletrs/cocos/agent"
)
var Verbose bool
type CLI struct {
agentSDK sdk.SDK
agentConfig grpc.AgentClientConfig
managerConfig grpc.ManagerClientConfig
client grpc.Client
managerClient manager.ManagerServiceClient
connectErr error
measurement cmdconfig.MeasurementProvider
agentSDK agent.Service
}
func New(agentConfig grpc.AgentClientConfig, managerConfig grpc.ManagerClientConfig, measurement cmdconfig.MeasurementProvider) *CLI {
func New(agentSDK agent.Service) *CLI {
return &CLI{
agentConfig: agentConfig,
managerConfig: managerConfig,
measurement: measurement,
agentSDK: agentSDK,
}
}
func (c *CLI) InitializeAgentSDK(cmd *cobra.Command) error {
agentGRPCClient, agentClient, err := agent.NewAgentClient(context.Background(), c.agentConfig)
if err != nil {
c.connectErr = err
return err
}
cmd.Println("🔗 Connected to agent ", agentGRPCClient.Secure())
c.client = agentGRPCClient
c.agentSDK = sdk.NewAgentSDK(agentClient)
return nil
}
func (c *CLI) InitializeManagerClient(cmd *cobra.Command) error {
managerGRPCClient, managerClient, err := managergrpc.NewManagerClient(c.managerConfig)
if err != nil {
c.connectErr = err
return err
}
cmd.Println("🔗 Connected to manager using ", managerGRPCClient.Secure())
c.client = managerGRPCClient
c.managerClient = managerClient
return nil
}
func (c *CLI) Close() {
c.client.Close()
}
+64 -220
View File
@@ -4,282 +4,126 @@ package main
import (
"context"
"crypto/sha256"
"crypto/sha512"
"crypto/x509"
"encoding/pem"
"encoding/json"
"fmt"
"log"
"log/slog"
"os"
"os/signal"
"syscall"
"time"
mglog "github.com/absmach/magistrala/logger"
"github.com/absmach/magistrala/pkg/prometheus"
"github.com/caarlos0/env/v11"
"github.com/mdlayher/vsock"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/agent/api"
"github.com/ultravioletrs/cocos/agent/cvms"
cvmsapi "github.com/ultravioletrs/cocos/agent/cvms/api/grpc"
"github.com/ultravioletrs/cocos/agent/cvms/server"
agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
"github.com/ultravioletrs/cocos/agent/events"
"github.com/ultravioletrs/cocos/internal"
agentlogger "github.com/ultravioletrs/cocos/internal/logger"
"github.com/ultravioletrs/cocos/pkg/attestation"
"github.com/ultravioletrs/cocos/pkg/attestation/azure"
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
cvmsgrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/cvm"
"github.com/ultravioletrs/cocos/internal/server"
grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
"github.com/ultravioletrs/cocos/manager"
"golang.org/x/sync/errgroup"
"google.golang.org/grpc"
"google.golang.org/grpc/reflection"
)
const (
svcName = "agent"
defSvcGRPCPort = "7002"
retryInterval = 5 * time.Second
envPrefixCVMGRPC = "AGENT_CVM_GRPC_"
storageDir = "/var/lib/cocos/agent"
svcName = "agent"
defSvcGRPCPort = "7002"
)
type config struct {
LogLevel string `env:"AGENT_LOG_LEVEL" envDefault:"debug"`
Vmpl int `env:"AGENT_VMPL" envDefault:"2"`
AgentGrpcHost string `env:"AGENT_GRPC_HOST" envDefault:"0.0.0.0"`
CAUrl string `env:"AGENT_CVM_CA_URL" envDefault:""`
CVMId string `env:"AGENT_CVM_ID" envDefault:""`
AgentMaaURL string `env:"AGENT_MAA_URL" envDefault:"https://sharedeus2.eus2.attest.azure.net"`
AgentOSBuild string `env:"AGENT_OS_BUILD" envDefault:"UVC"`
AgentOSDistro string `env:"AGENT_OS_DISTRO" envDefault:"UVC"`
AgentOSType string `env:"AGENT_OS_TYPE" envDefault:"UVC"`
}
func main() {
ctx, cancel := context.WithCancel(context.Background())
g, ctx := errgroup.WithContext(ctx)
var cfg config
if err := env.Parse(&cfg); err != nil {
log.Fatalf("failed to load %s configuration : %s", svcName, err)
cfg, err := readConfig()
if err != nil {
log.Fatalf("failed to read agent configuration from vsock %s", err.Error())
}
var exitCode int
defer mglog.ExitWithError(&exitCode)
conn, err := vsock.Dial(vsock.Host, manager.ManagerVsockPort, nil)
if err != nil {
log.Fatal(err)
}
defer conn.Close()
var level slog.Level
if err := level.UnmarshalText([]byte(cfg.LogLevel)); err != nil {
if err := level.UnmarshalText([]byte(cfg.AgentConfig.LogLevel)); err != nil {
log.Println(err)
exitCode = 1
return
}
eventsLogsQueue := make(chan *cvms.ClientStreamMessage, 1000)
handler := agentlogger.NewProtoHandler(os.Stdout, &slog.HandlerOptions{Level: level}, eventsLogsQueue)
handler := agentlogger.NewProtoHandler(conn, &slog.HandlerOptions{Level: level})
logger := slog.New(handler)
eventSvc, err := events.New(svcName, eventsLogsQueue)
eventSvc, err := events.New(svcName, cfg.ID, manager.ManagerVsockPort)
if err != nil {
logger.Error(fmt.Sprintf("failed to create events service %s", err.Error()))
exitCode = 1
log.Printf("failed to create events service %s", err.Error())
return
}
defer eventSvc.Close()
svc := newService(ctx, logger, eventSvc, cfg)
var provider attestation.Provider
ccPlatform := attestation.CCPlatform()
azureConfig := azure.NewEnvConfigFromAgent(
cfg.AgentOSBuild,
cfg.AgentOSType,
cfg.AgentOSDistro,
cfg.AgentMaaURL,
)
azure.InitializeDefaultMAAVars(azureConfig)
switch ccPlatform {
case attestation.SNP:
provider = vtpm.New(nil, false, uint(cfg.Vmpl), nil)
case attestation.SNPvTPM:
provider = vtpm.New(nil, true, uint(cfg.Vmpl), nil)
case attestation.Azure:
provider = azure.New(nil)
case attestation.NoCC:
logger.Info("TEE device not found")
provider = &attestation.EmptyProvider{}
grpcServerConfig := server.Config{
Port: cfg.AgentConfig.Port,
Host: cfg.AgentConfig.Host,
CertFile: cfg.AgentConfig.CertFile,
KeyFile: cfg.AgentConfig.KeyFile,
ServerCAFile: cfg.AgentConfig.ServerCAFile,
ClientCAFile: cfg.AgentConfig.ClientCAFile,
}
cvmGrpcConfig := pkggrpc.CVMClientConfig{}
if err := env.ParseWithOptions(&cvmGrpcConfig, env.Options{Prefix: envPrefixCVMGRPC}); err != nil {
logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
exitCode = 1
return
}
cvmGRPCClient, cvmsClient, err := cvmsgrpc.NewCVMClient(cvmGrpcConfig)
if err != nil {
logger.Error(err.Error())
exitCode = 1
return
}
defer cvmGRPCClient.Close()
reconnectFn := func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) {
grpcClient, newClient, err := cvmsgrpc.NewCVMClient(cvmGrpcConfig)
if err != nil {
return nil, nil, err
}
// Don't defer close here as we want to keep the connection open
pc, err := newClient.Process(ctx)
if err != nil {
grpcClient.Close()
return nil, nil, err
}
return grpcClient, pc, nil
}
pc, err := cvmsClient.Process(ctx)
if err != nil {
logger.Error(err.Error())
exitCode = 1
return
}
if cfg.Vmpl < 0 || cfg.Vmpl > 3 {
logger.Error("vmpl level must be in a range [0, 3]")
exitCode = 1
return
}
svc := newService(ctx, logger, eventSvc, provider, cfg.Vmpl)
if err := os.MkdirAll(storageDir, 0o755); err != nil {
logger.Error(fmt.Sprintf("failed to create storage directory: %s", err))
exitCode = 1
return
}
mc, err := cvmsapi.NewClient(pc, svc, eventsLogsQueue, logger, server.NewServer(logger, svc, cfg.AgentGrpcHost, cfg.CAUrl, cfg.CVMId), storageDir, reconnectFn, cvmGRPCClient)
if err != nil {
logger.Error(err.Error())
exitCode = 1
return
registerAgentServiceServer := func(srv *grpc.Server) {
reflection.Register(srv)
agent.RegisterAgentServiceServer(srv, agentgrpc.NewServer(svc))
}
gs := grpcserver.New(ctx, cancel, svcName, grpcServerConfig, registerAgentServiceServer, logger)
g.Go(func() error {
ch := make(chan os.Signal, 1)
signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
defer signal.Stop(ch)
select {
case <-ch:
logger.Info("Received signal, shutting down...")
cancel()
return nil
case <-ctx.Done():
return ctx.Err()
}
return gs.Start()
})
g.Go(func() error {
return mc.Process(ctx, cancel)
return server.StopHandler(ctx, cancel, logger, svcName, gs)
})
attest, certSerialNumber, err := attestationFromCert(ctx, cvmGrpcConfig.ClientCert, svc)
if err != nil {
logger.Error(fmt.Sprintf("failed to get attestation: %s", err))
exitCode = 1
return
}
if ccPlatform == attestation.Azure {
azureAttestationResult, azureCertSerialNumber, err := azureAttestationFromCert(ctx, cvmGrpcConfig.ClientCert, svc)
if err != nil {
logger.Error(fmt.Sprintf("failed to get attestation: %s", err))
exitCode = 1
return
}
eventsLogsQueue <- &cvms.ClientStreamMessage{
Message: &cvms.ClientStreamMessage_AzureAttestationResult{
AzureAttestationResult: &cvms.AzureAttestationResponse{
File: azureAttestationResult,
CertSerialNumber: azureCertSerialNumber,
},
},
}
}
eventsLogsQueue <- &cvms.ClientStreamMessage{
Message: &cvms.ClientStreamMessage_VTPMattestationReport{
VTPMattestationReport: &cvms.AttestationResponse{
File: attest,
CertSerialNumber: certSerialNumber,
},
},
}
if err := g.Wait(); err != nil {
logger.Error(fmt.Sprintf("%s service terminated: %s", svcName, err))
}
}
func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Service, provider attestation.Provider, vmpl int) agent.Service {
svc := agent.New(ctx, logger, eventSvc, provider, vmpl)
func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp agent.Computation) agent.Service {
svc := agent.New(ctx, logger, eventSvc, cmp)
svc = api.LoggingMiddleware(svc, logger)
counter, latency := prometheus.MakeMetrics(svcName, "api")
counter, latency := internal.MakeMetrics(svcName, "api")
svc = api.MetricsMiddleware(svc, counter, latency)
return svc
}
func attestationFromCert(ctx context.Context, certFilePath string, svc agent.Service) ([]byte, string, error) {
if certFilePath == "" {
return nil, "", nil
}
certFile, err := os.ReadFile(certFilePath)
func readConfig() (agent.Computation, error) {
l, err := vsock.Listen(manager.VsockConfigPort, nil)
if err != nil {
return nil, "", err
return agent.Computation{}, err
}
certPem, _ := pem.Decode(certFile)
certx509, err := x509.ParseCertificate(certPem.Bytes)
defer l.Close()
conn, err := l.Accept()
if err != nil {
return nil, "", err
return agent.Computation{}, err
}
nonceSNP := sha512.Sum512(certFile)
nonceVTPM := sha256.Sum256(certFile)
attest, err := svc.Attestation(ctx, nonceSNP, nonceVTPM, attestation.SNPvTPM)
defer conn.Close()
b := make([]byte, 1024)
n, err := conn.Read(b)
if err != nil {
return nil, "", err
return agent.Computation{}, err
}
return attest, certx509.SerialNumber.String(), nil
}
func azureAttestationFromCert(ctx context.Context, certFilePath string, svc agent.Service) ([]byte, string, error) {
if certFilePath == "" {
return nil, "", nil
}
certFile, err := os.ReadFile(certFilePath)
if err != nil {
return nil, "", err
}
certPem, _ := pem.Decode(certFile)
certx509, err := x509.ParseCertificate(certPem.Bytes)
if err != nil {
return nil, "", err
}
nonceAzure := sha256.Sum256(certFile)
attestation, err := svc.AttestationResult(ctx, nonceAzure, attestation.AzureToken)
if err != nil {
return nil, "", err
}
return attestation, certx509.SerialNumber.String(), nil
ac := agent.Computation{
AgentConfig: agent.AgentConfig{},
}
if err := json.Unmarshal(b[:n], &ac); err != nil {
return agent.Computation{}, err
}
if ac.AgentConfig.LogLevel == "" {
ac.AgentConfig.LogLevel = "info"
}
if ac.AgentConfig.Port == "" {
ac.AgentConfig.Port = defSvcGRPCPort
}
return ac, nil
}
+42 -120
View File
@@ -4,36 +4,57 @@ package main
import (
"fmt"
"log"
"os"
"os/signal"
"path"
"syscall"
"github.com/caarlos0/env/v11"
"github.com/fatih/color"
mglog "github.com/absmach/magistrala/logger"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"github.com/ultravioletrs/cocos/cli"
"github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig"
"github.com/ultravioletrs/cocos/internal/env"
"github.com/ultravioletrs/cocos/pkg/clients/grpc"
cmd "github.com/virtee/sev-snp-measure-go/sevsnpmeasure/cmd"
"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
"github.com/ultravioletrs/cocos/pkg/sdk"
)
const (
svcName = "cli"
envPrefixAgentGRPC = "AGENT_GRPC_"
envPrefixManagerGRPC = "MANAGER_GRPC_"
completion = "completion"
filePermision = 0o755
cocosDirectory = ".cocos"
svcName = "cli"
envPrefixAgentGRPC = "AGENT_GRPC_"
completion = "completion"
)
type config struct {
LogLevel string `env:"AGENT_LOG_LEVEL" envDefault:"info"`
IgvmBinaryPath string `env:"IGVM_BINARY_PATH" envDefault:"./build/igvmmeasure"`
LogLevel string `env:"AGENT_LOG_LEVEL" envDefault:"info"`
}
func main() {
var cfg config
if err := env.Parse(&cfg); err != nil {
log.Fatalf("failed to load %s configuration : %s", svcName, err)
}
logger, err := mglog.New(os.Stdout, cfg.LogLevel)
if err != nil {
log.Fatalf("Error creating logger: %s", err)
}
agentGRPCConfig := grpc.Config{}
if err := env.Parse(&agentGRPCConfig, env.Options{Prefix: envPrefixAgentGRPC}); err != nil {
logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
return
}
agentGRPCClient, agentClient, err := agent.NewAgentClient(agentGRPCConfig)
if err != nil {
logger.Error(err.Error())
return
}
defer agentGRPCClient.Close()
agentSDK := sdk.NewAgentSDK(logger, agentClient)
cliSVC := cli.New(agentSDK)
rootCmd := &cobra.Command{
Use: "cocos-cli [command]",
Short: "CLI application for CoCos Service API",
@@ -62,118 +83,19 @@ func main() {
},
}
signalChan := make(chan os.Signal, 1)
signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
go func() {
<-signalChan
fmt.Println()
rootCmd.Println(color.New(color.FgRed).Sprint("Operation aborted by user!"))
os.Exit(2)
}()
var cfg config
if err := env.Parse(&cfg); err != nil {
message := color.New(color.FgRed).Sprintf("failed to load %s configuration : %s", svcName, err)
rootCmd.Println(message)
return
}
homePath, err := os.UserHomeDir()
if err != nil {
message := color.New(color.FgRed).Sprintf("failed to fetch user home directory: %s", err)
rootCmd.Println(message)
return
}
directoryCachePath := path.Join(homePath, cocosDirectory)
if err := os.MkdirAll(directoryCachePath, filePermision); err != nil {
message := color.New(color.FgRed).Sprintf("failed to create directory %s : %s", directoryCachePath, err)
rootCmd.Println(message)
return
}
agentGRPCConfig := grpc.AgentClientConfig{}
if err := env.ParseWithOptions(&agentGRPCConfig, env.Options{Prefix: envPrefixAgentGRPC}); err != nil {
message := color.New(color.FgRed).Sprintf("failed to load %s gRPC client configuration : %s", svcName, err)
rootCmd.Println(message)
return
}
managerGRPCConfig := grpc.ManagerClientConfig{}
if err := env.ParseWithOptions(&managerGRPCConfig, env.Options{Prefix: envPrefixManagerGRPC}); err != nil {
message := color.New(color.FgRed).Sprintf("failed to load %s gRPC client configuration : %s", svcName, err)
rootCmd.Println(message)
return
}
options := cmdconfig.IgvmMeasureOptions
measurement, err := cmdconfig.NewCmdConfig(cfg.IgvmBinaryPath, options, os.Stderr)
if err != nil {
message := color.New(color.FgRed).Sprintf("failed to initialize measurement: %s", err) // Use %s instead of %w
rootCmd.Println(message)
return
}
cliSVC := cli.New(agentGRPCConfig, managerGRPCConfig, measurement)
if err := cliSVC.InitializeAgentSDK(rootCmd); err == nil {
defer cliSVC.Close()
}
rootCmd.PersistentFlags().BoolVarP(&cli.Verbose, "verbose", "v", false, "Enable verbose output")
keysCmd := cliSVC.NewKeysCmd()
attestationCmd := cliSVC.NewAttestationCmd()
attestationPolicyCmd := cliSVC.NewAttestationPolicyCmd()
// Agent Commands
rootCmd.AddCommand(cliSVC.NewAlgorithmCmd())
rootCmd.AddCommand(cliSVC.NewAlgorithmsCmd())
rootCmd.AddCommand(cliSVC.NewDatasetsCmd())
rootCmd.AddCommand(cliSVC.NewResultsCmd())
rootCmd.AddCommand(attestationCmd)
rootCmd.AddCommand(cliSVC.NewFileHashCmd())
rootCmd.AddCommand(attestationPolicyCmd)
rootCmd.AddCommand(keysCmd)
rootCmd.AddCommand(cliSVC.NewCABundleCmd(directoryCachePath))
rootCmd.AddCommand(cliSVC.NewCreateVMCmd())
rootCmd.AddCommand(cliSVC.NewRemoveVMCmd())
rootCmd.AddCommand(cliSVC.NewIMAMeasurementsCmd())
attestaionCmd := cliSVC.NewAttestationCmd()
rootCmd.AddCommand(attestaionCmd)
// Attestation commands
attestationCmd.AddCommand(cliSVC.NewGetAttestationCmd())
attestationCmd.AddCommand(cliSVC.NewValidateAttestationValidationCmd())
// measure.
rootCmd.AddCommand(cmd.NewRootCmd())
rootCmd.AddCommand(cliSVC.NewMeasureCmd(cfg.IgvmBinaryPath))
// Flags
keysCmd.PersistentFlags().StringVarP(
&cli.KeyType,
"key-type",
"k",
"rsa",
"User Key type",
)
// Attestation Policy commands
attestationPolicyCmd.AddCommand(cliSVC.NewAddMeasurementCmd())
attestationPolicyCmd.AddCommand(cliSVC.NewAddHostDataCmd())
attestationPolicyCmd.AddCommand(cliSVC.NewGCPAttestationPolicy())
attestationPolicyCmd.AddCommand(cliSVC.NewDownloadGCPOvmfFile())
attestationPolicyCmd.AddCommand(cliSVC.NewAzureAttestationPolicy())
attestaionCmd.AddCommand(cliSVC.NewGetAttestationCmd())
attestaionCmd.AddCommand(cliSVC.NewValidateAttestationValidationCmd())
if err := rootCmd.Execute(); err != nil {
logErrorCmd(*rootCmd, err)
logger.Error(fmt.Sprintf("Command execution failed: %s", err))
return
}
}
func logErrorCmd(cmd cobra.Command, err error) {
boldRed := color.New(color.FgRed, color.Bold)
boldRed.Fprintf(cmd.ErrOrStderr(), "\nerror: ")
fmt.Fprintf(cmd.ErrOrStderr(), "%s\n\n", color.RedString(err.Error()))
}
+50 -72
View File
@@ -8,48 +8,40 @@ import (
"fmt"
"log"
"log/slog"
"net/url"
"os"
"strings"
mglog "github.com/absmach/magistrala/logger"
"github.com/absmach/magistrala/pkg/jaeger"
"github.com/absmach/magistrala/pkg/prometheus"
"github.com/absmach/magistrala/pkg/uuid"
"github.com/caarlos0/env/v11"
"github.com/ultravioletrs/cocos/internal/server"
grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
"github.com/ultravioletrs/cocos/internal"
"github.com/ultravioletrs/cocos/internal/env"
jaegerclient "github.com/ultravioletrs/cocos/internal/jaeger"
"github.com/ultravioletrs/cocos/manager"
"github.com/ultravioletrs/cocos/manager/api"
managergrpc "github.com/ultravioletrs/cocos/manager/api/grpc"
managerapi "github.com/ultravioletrs/cocos/manager/api/grpc"
"github.com/ultravioletrs/cocos/manager/qemu"
"github.com/ultravioletrs/cocos/manager/tracing"
"github.com/ultravioletrs/cocos/pkg/clients/grpc"
managergrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/manager"
pkgmanager "github.com/ultravioletrs/cocos/pkg/manager"
"go.opentelemetry.io/otel/trace"
"golang.org/x/sync/errgroup"
"google.golang.org/grpc"
"google.golang.org/grpc/reflection"
)
const (
svcName = "manager"
envPrefixGRPC = "MANAGER_GRPC_"
envPrefixQemu = "MANAGER_QEMU_"
clientBufferSize = 100
svcName = "manager"
envPrefixGRPC = "MANAGER_GRPC_"
envPrefixQemu = "MANAGER_QEMU_"
)
type config struct {
LogLevel string `env:"MANAGER_LOG_LEVEL" envDefault:"info"`
JaegerURL url.URL `env:"COCOS_JAEGER_URL" envDefault:"http://localhost:4318"`
TraceRatio float64 `env:"COCOS_JAEGER_TRACE_RATIO" envDefault:"1.0"`
InstanceID string `env:"MANAGER_INSTANCE_ID" envDefault:""`
AttestationPolicyBinary string `env:"MANAGER_ATTESTATION_POLICY_BINARY" envDefault:"../../build/attestation_policy"`
IgvmMeasureBinary string `env:"MANAGER_IGVMMEASURE_BINARY" envDefault:"../../build/igvmmeasure"`
PcrValues string `env:"MANAGER_PCR_VALUES" envDefault:""`
EosVersion string `env:"MANAGER_EOS_VERSION" envDefault:""`
LogLevel string `env:"MANAGER_LOG_LEVEL" envDefault:"info"`
JaegerURL string `env:"COCOS_JAEGER_URL" envDefault:"http://localhost:14268/api/traces"`
InstanceID string `env:"MANAGER_INSTANCE_ID" envDefault:""`
}
func main() {
ctx, cancel := context.WithCancel(context.Background())
ctx := context.Background()
g, ctx := errgroup.WithContext(ctx)
var cfg config
@@ -59,21 +51,18 @@ func main() {
logger, err := mglog.New(os.Stdout, cfg.LogLevel)
if err != nil {
log.Fatal(err.Error())
log.Fatalf(err.Error())
}
var exitCode int
defer mglog.ExitWithError(&exitCode)
if cfg.InstanceID == "" {
if cfg.InstanceID, err = uuid.New().ID(); err != nil {
cfg.InstanceID, err = uuid.New().ID()
if err != nil {
logger.Error(fmt.Sprintf("Failed to generate instance ID: %s", err))
exitCode = 1
return
}
}
tp, err := jaeger.NewProvider(ctx, svcName, cfg.JaegerURL, cfg.InstanceID, cfg.TraceRatio)
tp, err := jaegerclient.NewProvider(ctx, svcName, cfg.JaegerURL, cfg.InstanceID)
if err != nil {
logger.Error(fmt.Sprintf("Failed to init Jaeger: %s", err))
}
@@ -84,74 +73,63 @@ func main() {
}()
tracer := tp.Tracer(svcName)
qemuCfg, err := qemu.NewConfig()
if err != nil {
logger.Error(fmt.Sprintf("failed to create config: %v", err))
exitCode = 1
return
}
if qemuCfg.EnableTDX {
logger.Info("Manager started with TDX enabled")
} else if qemuCfg.EnableSEVSNP {
logger.Info("Manager started with SEV-SNP enabled")
} else if qemuCfg.EnableSEV {
logger.Info("Manager started with SEV enabled")
} else {
logger.Info("Manager started without confidential computing support")
}
if err := env.ParseWithOptions(&qemuCfg, env.Options{Prefix: envPrefixQemu}); err != nil {
qemuCfg := qemu.Config{}
if err := env.Parse(&qemuCfg, env.Options{Prefix: envPrefixQemu}); err != nil {
logger.Error(fmt.Sprintf("failed to load QEMU configuration: %s", err))
exitCode = 1
return
}
args := qemuCfg.ConstructQemuArgs()
logger.Info(strings.Join(args, " "))
exe, args, err := qemu.ExecutableAndArgs(qemuCfg)
if err != nil {
logger.Error(fmt.Sprintf("failed to parse QEMU configuration: %s", err))
return
}
logger.Info(fmt.Sprintf("%s %s", exe, strings.Join(args, " ")))
managerGRPCConfig := server.ServerConfig{}
if err := env.ParseWithOptions(&managerGRPCConfig, env.Options{Prefix: envPrefixGRPC}); err != nil {
managerGRPCConfig := grpc.Config{}
if err := env.Parse(&managerGRPCConfig, env.Options{Prefix: envPrefixGRPC}); err != nil {
logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
exitCode = 1
return
}
svc, err := newService(ctx, logger, tracer, *qemuCfg, cfg.AttestationPolicyBinary, cfg.IgvmMeasureBinary, cfg.PcrValues, cfg.EosVersion)
managerGRPCClient, managerClient, err := managergrpc.NewManagerClient(managerGRPCConfig)
if err != nil {
logger.Error(err.Error())
return
}
defer managerGRPCClient.Close()
pc, err := managerClient.Process(ctx)
if err != nil {
logger.Error(err.Error())
exitCode = 1
return
}
registerManagerServiceServer := func(srv *grpc.Server) {
reflection.Register(srv)
manager.RegisterManagerServiceServer(srv, managergrpc.NewServer(svc))
}
eventsChan := make(chan *pkgmanager.ClientStreamMessage)
svc := newService(logger, tracer, qemuCfg, eventsChan)
gs := grpcserver.New(ctx, cancel, svcName, managerGRPCConfig, registerManagerServiceServer, logger, nil, "", "")
mc := managerapi.NewClient(pc, svc, eventsChan)
g.Go(func() error {
return gs.Start()
})
g.Go(func() error {
return server.StopHandler(ctx, cancel, logger, svcName, gs)
return mc.Process(ctx)
})
if err := g.Wait(); err != nil {
logger.Error(fmt.Sprintf("%s service terminated: %s", svcName, err))
}
err = internal.DeleteFilesInDir(qemuCfg.TmpFileLoc)
if err != nil {
logger.Error(err.Error())
}
}
func newService(ctx context.Context, logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, attestationPolicyPath string, igvmMeasurementBinaryPath string, pcrValuesFilePath string, eosVersion string) (manager.Service, error) {
svc, err := manager.New(ctx, qemuCfg, attestationPolicyPath, igvmMeasurementBinaryPath, pcrValuesFilePath, logger, qemu.NewVM, eosVersion)
if err != nil {
return nil, err
}
func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, eventsChan chan *pkgmanager.ClientStreamMessage) manager.Service {
svc := manager.New(qemuCfg, logger, eventsChan)
go svc.RetrieveAgentEventsLogs()
svc = api.LoggingMiddleware(svc, logger)
counter, latency := prometheus.MakeMetrics(svcName, "api")
counter, latency := internal.MakeMetrics(svcName, "api")
svc = api.MetricsMiddleware(svc, counter, latency)
svc = tracing.New(svc, tracer)
return svc, nil
return svc
}
+270
View File
@@ -0,0 +1,270 @@
#!/bin/bash
#
# user changeable parameters
#
HDA_FILE="cmd/manager/img/focal-server-cloudimg-amd64.qcow2"
GUEST_SIZE_IN_MB="4096"
SEV_GUEST="1"
SMP_NCPUS="4"
CONSOLE="serial"
VNC_PORT=""
USE_VIRTIO="1"
UEFI_BIOS_CODE="/usr/share/OVMF/OVMF_CODE.fd"
UEFI_BIOS_VARS_ORIG="/usr/share/OVMF/OVMF_VARS.fd"
UEFI_BIOS_VARS_COPY="cmd/manager/img/OVMF_VARS.fd"
CBITPOS=51
HOST_HTTP_PORT=9301
GUEST_HTTP_PORT=9031
HOST_GRPC_PORT=7020
GUEST_GRPC_PORT=7002
ENABLE_FILE_LOG="0"
EXEC_QEMU_CMDLINE="0"
usage() {
echo "$0 [options]"
echo "Available <commands>:"
echo " -hda hard disk ($HDA_FILE)"
echo " -nosev disable sev support"
echo " -mem guest memory"
echo " -smp number of cpus"
echo " -console display console to use (serial or gxl)"
echo " -vnc VNC port to use"
echo " -bios bios to use (default $UEFI_BIOS_CODE)"
echo " -kernel kernel to use"
echo " -initrd initrd to use"
echo " -cdrom CDROM image"
echo " -virtio use virtio devices"
echo " -cbitpos location of the C-bit"
echo " -hosthttp host http port"
echo " -guesthttp guest http port"
echo " -hostgrpc host grpc port"
echo " -guestgrpc guest grpc port"
echo " -origuefivars UEFI BIOS vars original file (default $UEFI_BIOS_VARS_ORIG)"
echo " -copyuefivars UEFI BIOS vars copy file (default $UEFI_BIOS_VARS_COPY)"
echo " -exec execute the QEMU command (default $EXEC_QEMU_CMDLINE)"
echo " -filelog enable/disable QEMU cmd line file log (default: $ENABLE_FILE_LOG)"
exit 1
}
while [[ $1 != "" ]]; do
case "$1" in
-hda)
HDA_FILE=${2}
shift
;;
-nosev)
SEV_GUEST="0"
;;
-mem)
GUEST_SIZE_IN_MB=${2}
shift
;;
-console)
CONSOLE=${2}
shift
;;
-smp)
SMP_NCPUS=$2
shift
;;
-vnc)
VNC_PORT=$2
shift
;;
-bios)
UEFI_BIOS_CODE=$2
shift
;;
-initrd)
INITRD_FILE=$2
shift
;;
-kernel)
KERNEL_FILE=$2
shift
;;
-cdrom)
CDROM_FILE=$2
shift
;;
-virtio)
USE_VIRTIO="1"
;;
-cbitpos)
CBITPOS=$2
shift
;;
-hosthttp)
HOST_HTTP_PORT=$2
shift
;;
-guesthttp)
GUEST_HTTP_PORT=$2
shift
;;
-guestgrpc)
GUEST_GRPC_PORT=$2
shift
;;
-hostgrpc)
HOST_GRPC_PORT=$2
shift
;;
-origuefivars)
UEFI_BIOS_VARS_ORIG=$2
shift
;;
-copyuefivars)
UEFI_BIOS_VARS_COPY=$2
shift
;;
-exec)
EXEC_QEMU_CMDLINE="1"
;;
-filelog)
ENABLE_FILE_LOG="1"
;;
*)
usage;;
esac
shift
done
#
# func definitions
#
add_opts() {
echo -n "$* " >> ${QEMU_CMDLINE}
}
run_cmd() {
if ! "$@"; then
echo "Command '$*' failed"
exit 1
fi
}
# copy BIOS variables to new dest for VM use without modifying the original ones
cp "$UEFI_BIOS_VARS_ORIG" "$UEFI_BIOS_VARS_COPY"
#
# Qemu cmd line construction
#
# we add all the qemu command line options into a file
QEMU_CMDLINE=/tmp/cmdline.$$
rm -rf ${QEMU_CMDLINE}
add_opts "$(which qemu-system-x86_64)"
# Basic virtual machine property
add_opts "-enable-kvm -cpu EPYC -machine q35"
# add number of VCPUs
[ -n "$SMP_NCPUS" ] && add_opts "-smp ${SMP_NCPUS},maxcpus=64"
# define guest memory
add_opts "-m ${GUEST_SIZE_IN_MB}M,slots=5,maxmem=30G"
# The OVMF binary, including the non-volatile variable store, appears as a
# "normal" qemu drive on the host side, and it is exposed to the guest as a
# persistent flash device.
add_opts "-drive if=pflash,format=raw,unit=0,file=${UEFI_BIOS_CODE},readonly=on"
add_opts "-drive if=pflash,format=raw,unit=1,file=${UEFI_BIOS_VARS_COPY}"
# add CDROM if specified
[ -n "$CDROM_FILE" ] && add_opts "-drive file=${CDROM_FILE},media=cdrom -boot d"
add_opts "-netdev user,id=vmnic,hostfwd=tcp::2222-:22,hostfwd=tcp::$HOST_HTTP_PORT-:$GUEST_HTTP_PORT,hostfwd=tcp::$HOST_GRPC_PORT-:$GUEST_GRPC_PORT"
add_opts "-device virtio-net-pci,disable-legacy=on,iommu_platform=true,netdev=vmnic,romfile="
# If harddisk file is specified then add the HDD drive
if [ -n "$HDA_FILE" ]; then
if [ "$USE_VIRTIO" = "1" ]; then
if [[ ${HDA_FILE} = *"qcow2" ]]; then
add_opts "-drive file=${HDA_FILE},if=none,id=disk0,format=qcow2"
else
add_opts "-drive file=${HDA_FILE},if=none,id=disk0,format=raw"
fi
add_opts "-device virtio-scsi-pci,id=scsi,disable-legacy=on,iommu_platform=true"
add_opts "-device scsi-hd,drive=disk0"
else
if [[ ${HDA_FILE} = *"qcow2" ]]; then
add_opts "-drive file=${HDA_FILE},format=qcow2"
else
add_opts "-drive file=${HDA_FILE},format=raw"
fi
fi
fi
# If this is SEV guest then add the encryption device objects to enable support
if [ ${SEV_GUEST} = "1" ]; then
add_opts "-object sev-guest,id=sev0,cbitpos=${CBITPOS},reduced-phys-bits=1"
add_opts "-machine memory-encryption=sev0"
fi
# if console is serial then disable graphical interface
if [ "${CONSOLE}" = "serial" ]; then
add_opts "-nographic"
else
add_opts "-vga ${CONSOLE}"
fi
# if -kernel arg is specified then use the kernel provided in command line for boot
if [ "${KERNEL_FILE}" != "" ]; then
add_opts "-kernel $KERNEL_FILE"
add_opts "-append \"console=ttyS0 earlyprintk=serial root=/dev/sda2\""
[ -n "$INITRD_FILE" ] && add_opts "-initrd ${INITRD_FILE}"
fi
# start vnc server
[ -n "$VNC_PORT" ] && add_opts "-vnc :${VNC_PORT}" && echo "Starting VNC on port ${VNC_PORT}"
# start monitor on pty
add_opts "-monitor pty"
#
# Qemu cmd line log
#
# Set the log file path if ENABLE_FILE_LOG is 1
if [ "$ENABLE_FILE_LOG" = "1" ]; then
LOG_FILE=$(pwd)/stdout.log
# Save the command line args into log file
cat "$QEMU_CMDLINE" > "$LOG_FILE"
echo >> "$LOG_FILE"
fi
# Log the command line to the console
cat "$QEMU_CMDLINE"
#
# Qemu cmd line execution
#
if [[ "${EXEC_QEMU_CMDLINE}" = "0" ]]; then
exit 0
fi
# map CTRL-C to CTRL ]
echo "Mapping CTRL-C to CTRL-]"
stty intr ^]
echo "Launching VM ..."
if [ "$ENABLE_FILE_LOG" = "1" ]; then
bash ${QEMU_CMDLINE} 2>&1 | tee -a "${LOG_FILE}"
else
bash ${QEMU_CMDLINE} 2>&1
fi
# restore the mapping
stty intr ^c
rm -rf ${QEMU_CMDLINE}

Some files were not shown because too many files have changed in this diff Show More