mirror of
https://github.com/ultravioletrs/cocos.git
synced 2026-08-07 15:26:26 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 dorcaslitunya
|
0072e3cbcc |
@@ -16,3 +16,7 @@ updates:
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
day: "monday"
|
||||
groups:
|
||||
go-dependency:
|
||||
patterns:
|
||||
- "*"
|
||||
|
||||
@@ -33,8 +33,8 @@ jobs:
|
||||
|
||||
- name: Set up protoc
|
||||
run: |
|
||||
PROTOC_VERSION=29.3
|
||||
PROTOC_GEN_VERSION=v1.36.4
|
||||
PROTOC_VERSION=28.1
|
||||
PROTOC_GEN_VERSION=v1.34.2
|
||||
PROTOC_GRPC_VERSION=v1.5.1
|
||||
|
||||
# Download and install protoc
|
||||
@@ -55,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
|
||||
|
||||
@@ -67,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
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
name: Build and Release Hal
|
||||
name: Build and Release
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- "*"
|
||||
- '*'
|
||||
|
||||
jobs:
|
||||
build:
|
||||
@@ -32,8 +32,8 @@ jobs:
|
||||
with:
|
||||
root-reserve-mb: 35000
|
||||
swap-size-mb: 1024
|
||||
remove-dotnet: "true"
|
||||
remove-android: "true"
|
||||
remove-dotnet: 'true'
|
||||
remove-android: 'true'
|
||||
- name: Check free space
|
||||
run: |
|
||||
echo "Free space:"
|
||||
@@ -48,33 +48,26 @@ jobs:
|
||||
- name: Checkout cocos
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: "ultravioletrs/cocos"
|
||||
repository: 'ultravioletrs/cocos'
|
||||
path: cocos
|
||||
|
||||
- name: Checkout buildroot
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: "buildroot/buildroot"
|
||||
repository: 'buildroot/buildroot'
|
||||
path: buildroot
|
||||
ref: 2024.11-rc2
|
||||
|
||||
- name: Build hal
|
||||
- name: Build
|
||||
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
|
||||
|
||||
|
||||
@@ -35,7 +35,7 @@ jobs:
|
||||
mkdir coverage
|
||||
|
||||
- name: Run Agent tests
|
||||
run: go test --tags embed -v --race -covermode=atomic -coverprofile coverage/agent.out ./agent/...
|
||||
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/...
|
||||
|
||||
+1
-5
@@ -4,10 +4,6 @@ run:
|
||||
issues:
|
||||
max-issues-per-linter: 10
|
||||
max-same-issues: 10
|
||||
exclude-rules:
|
||||
- linters:
|
||||
- makezero
|
||||
text: "with non-zero initialized length"
|
||||
|
||||
linters-settings:
|
||||
importas:
|
||||
@@ -71,6 +67,7 @@ linters:
|
||||
- dogsled
|
||||
- errchkjson
|
||||
- errname
|
||||
- execinquery
|
||||
- copyloopvar
|
||||
- ginkgolinter
|
||||
- gocheckcompilerdirectives
|
||||
@@ -81,4 +78,3 @@ linters:
|
||||
- mirror
|
||||
- nakedret
|
||||
- dupword
|
||||
|
||||
|
||||
@@ -12,7 +12,6 @@ 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
|
||||
|
||||
define compile_service
|
||||
CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) \
|
||||
@@ -28,9 +27,8 @@ endef
|
||||
|
||||
all: $(SERVICES)
|
||||
|
||||
$(SERVICES):
|
||||
$(SERVICES):
|
||||
$(call compile_service,$@)
|
||||
@if [ "$@" = "cli" ] || [ "$@" = "manager" ]; then $(MAKE) build-igvm; fi
|
||||
|
||||
$(ATTESTATION_POLICY):
|
||||
$(MAKE) -C ./scripts/attestation_policy
|
||||
@@ -39,7 +37,6 @@ 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
|
||||
@@ -63,7 +60,3 @@ stop:
|
||||
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)
|
||||
|
||||
@@ -1,78 +1,65 @@
|
||||
<div align="center">
|
||||
|
||||
# Cocos AI 🥥
|
||||
|
||||
**Confidential Computing System for AI**
|
||||
|
||||
**Made with ❤️ by [Ultraviolet](https://ultraviolet.rs/)**
|
||||
# Cocos AI
|
||||
|
||||
[](https://codecov.io/gh/ultravioletrs/cocos)
|
||||
[](https://goreportcard.com/report/github.com/ultravioletrs/cocos)
|
||||
[](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 (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 algorithm.
|
||||
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)
|
||||
[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
|
||||
|
||||
+9
-9
@@ -6,15 +6,15 @@ 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 | "" |
|
||||
| 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 | "" |
|
||||
|
||||
|
||||
## Deployment
|
||||
|
||||
+208
-108
@@ -3,8 +3,8 @@
|
||||
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.4
|
||||
// protoc v5.29.3
|
||||
// protoc-gen-go v1.34.2
|
||||
// protoc v5.28.1
|
||||
// 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,21 @@ 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"`
|
||||
Requirements []byte `protobuf:"bytes,2,opt,name=requirements,proto3" json:"requirements,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 +49,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)
|
||||
@@ -77,16 +79,18 @@ func (x *AlgoRequest) GetRequirements() []byte {
|
||||
}
|
||||
|
||||
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 +101,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 +117,21 @@ 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"`
|
||||
Filename string `protobuf:"bytes,2,opt,name=filename,proto3" json:"filename,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 +142,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)
|
||||
@@ -165,16 +172,18 @@ func (x *DataRequest) GetFilename() string {
|
||||
}
|
||||
|
||||
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 +194,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 +210,18 @@ 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
|
||||
}
|
||||
|
||||
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 +232,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)
|
||||
@@ -237,17 +248,20 @@ func (*ResultRequest) Descriptor() ([]byte, []int) {
|
||||
}
|
||||
|
||||
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 +272,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 +295,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 +319,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 +334,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 +366,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)
|
||||
@@ -386,7 +390,7 @@ func (x *AttestationResponse) GetFile() []byte {
|
||||
|
||||
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,
|
||||
0x67, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1c, 0x0a, 0x09, 0x61, 0x6c, 0x67,
|
||||
@@ -403,43 +407,40 @@ var file_agent_agent_proto_rawDesc = string([]byte{
|
||||
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, 0x32, 0xfd, 0x01, 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, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61, 0x67, 0x65, 0x6e,
|
||||
0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
|
||||
})
|
||||
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, 0xfd, 0x01, 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, 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
|
||||
}
|
||||
@@ -476,11 +477,109 @@ 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 any, i int) any {
|
||||
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 any, i int) any {
|
||||
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 any, i int) any {
|
||||
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 any, i int) any {
|
||||
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 any, i int) any {
|
||||
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 any, i int) any {
|
||||
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 any, i int) any {
|
||||
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 any, i int) any {
|
||||
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: 8,
|
||||
NumExtensions: 0,
|
||||
@@ -491,6 +590,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
|
||||
}
|
||||
|
||||
+1
-3
@@ -36,9 +36,7 @@ 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 {
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
|
||||
// versions:
|
||||
// - protoc-gen-go-grpc v1.5.1
|
||||
// - protoc v5.29.3
|
||||
// - protoc v5.28.1
|
||||
// source: agent/agent.proto
|
||||
|
||||
package agent
|
||||
|
||||
@@ -46,7 +46,4 @@ func AlgorithmArgsFromContext(ctx context.Context) []string {
|
||||
type Algorithm interface {
|
||||
// Run executes the algorithm and returns the result.
|
||||
Run() error
|
||||
|
||||
// Stop stops the algorithm.
|
||||
Stop() error
|
||||
}
|
||||
|
||||
@@ -20,46 +20,29 @@ type binary struct {
|
||||
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 {
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string) algorithm.Algorithm {
|
||||
return &binary{
|
||||
algoFile: algoFile,
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
|
||||
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
|
||||
cmd := exec.Command(b.algoFile, b.args...)
|
||||
cmd.Stderr = b.stderr
|
||||
cmd.Stdout = b.stdout
|
||||
|
||||
if err := b.cmd.Start(); err != nil {
|
||||
if err := cmd.Start(); err != nil {
|
||||
return fmt.Errorf("error starting algorithm: %v", err)
|
||||
}
|
||||
|
||||
if err := b.cmd.Wait(); err != nil {
|
||||
if err := 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 err := b.cmd.Process.Kill(); err != nil {
|
||||
return fmt.Errorf("error stopping algorithm: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -18,7 +18,7 @@ func TestNewAlgorithm(t *testing.T) {
|
||||
algoFile := "/path/to/algo"
|
||||
args := []string{"arg1", "arg2"}
|
||||
|
||||
algo := NewAlgorithm(logger, eventsSvc, algoFile, args, "")
|
||||
algo := NewAlgorithm(logger, eventsSvc, algoFile, args)
|
||||
|
||||
b, ok := algo.(*binary)
|
||||
if !ok {
|
||||
@@ -74,7 +74,7 @@ func TestBinaryRun(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventsSvc := new(mocks.Service)
|
||||
|
||||
b := NewAlgorithm(logger, eventsSvc, tt.algoFile, tt.args, "").(*binary)
|
||||
b := NewAlgorithm(logger, eventsSvc, tt.algoFile, tt.args).(*binary)
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
b.stdout = &stdout
|
||||
|
||||
@@ -35,11 +35,11 @@ type docker struct {
|
||||
stdout io.Writer
|
||||
}
|
||||
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile, cmpID string) algorithm.Algorithm {
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string) algorithm.Algorithm {
|
||||
d := &docker{
|
||||
algoFile: algoFile,
|
||||
logger: logger,
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
|
||||
stdout: &logging.Stdout{Logger: logger},
|
||||
}
|
||||
|
||||
@@ -47,6 +47,8 @@ func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile, cmpID
|
||||
}
|
||||
|
||||
func (d *docker) Run() error {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a new Docker client.
|
||||
cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
|
||||
if err != nil {
|
||||
@@ -60,9 +62,8 @@ func (d *docker) Run() error {
|
||||
}
|
||||
defer imageFile.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
// Load the Docker image from the tar file.
|
||||
resp, err := cli.ImageLoad(ctx, imageFile, client.ImageLoadWithQuiet(true))
|
||||
resp, err := cli.ImageLoad(ctx, imageFile, true)
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not load Docker image from file: %v", err)
|
||||
}
|
||||
@@ -175,8 +176,3 @@ func writeToOut(readCloser io.ReadCloser, ioWriter io.Writer) error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *docker) Stop() error {
|
||||
// To be supported later.
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -18,7 +18,7 @@ func TestNewAlgorithm(t *testing.T) {
|
||||
eventsSvc := new(mocks.Service)
|
||||
algoFile := "/path/to/algo.tar"
|
||||
|
||||
algo := NewAlgorithm(logger, eventsSvc, algoFile, "")
|
||||
algo := NewAlgorithm(logger, eventsSvc, algoFile)
|
||||
|
||||
d, ok := algo.(*docker)
|
||||
assert.True(t, ok, "NewAlgorithm should return a *docker")
|
||||
|
||||
@@ -50,7 +50,6 @@ func (s *Stdout) Write(p []byte) (n int, err error) {
|
||||
type Stderr struct {
|
||||
Logger *slog.Logger
|
||||
EventSvc events.Service
|
||||
CmpID string
|
||||
}
|
||||
|
||||
// Write implements io.Writer.
|
||||
@@ -71,7 +70,9 @@ func (s *Stderr) Write(p []byte) (n int, err error) {
|
||||
s.Logger.Error(string(buf[:n]))
|
||||
}
|
||||
|
||||
s.EventSvc.SendEvent(s.CmpID, algorithmRun, warningStatus, json.RawMessage{})
|
||||
if err := s.EventSvc.SendEvent(algorithmRun, warningStatus, json.RawMessage{}); err != nil {
|
||||
return len(p), err
|
||||
}
|
||||
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
@@ -73,7 +73,7 @@ func TestStderrWrite(t *testing.T) {
|
||||
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)
|
||||
mockEventService.On("SendEvent", "AlgorithmRun", manager.Warning.String(), mock.Anything).Return(nil)
|
||||
|
||||
stderr := &Stderr{Logger: mglog.NewMock(), EventSvc: mockEventService}
|
||||
n, err := stderr.Write([]byte(tt.input))
|
||||
|
||||
@@ -39,13 +39,12 @@ type python struct {
|
||||
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 {
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, runtime, requirementsFile, algoFile string, args []string) algorithm.Algorithm {
|
||||
p := &python{
|
||||
algoFile: algoFile,
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
|
||||
stdout: &logging.Stdout{Logger: logger},
|
||||
requirementsFile: requirementsFile,
|
||||
args: args,
|
||||
@@ -86,15 +85,15 @@ func (p *python) Run() error {
|
||||
}
|
||||
|
||||
args := append([]string{p.algoFile}, p.args...)
|
||||
p.cmd = exec.Command(pythonPath, args...)
|
||||
p.cmd.Stderr = p.stderr
|
||||
p.cmd.Stdout = p.stdout
|
||||
cmd := exec.Command(pythonPath, args...)
|
||||
cmd.Stderr = p.stderr
|
||||
cmd.Stdout = p.stdout
|
||||
|
||||
if err := p.cmd.Start(); err != nil {
|
||||
if err := cmd.Start(); err != nil {
|
||||
return fmt.Errorf("error starting algorithm: %v", err)
|
||||
}
|
||||
|
||||
if err := p.cmd.Wait(); err != nil {
|
||||
if err := cmd.Wait(); err != nil {
|
||||
return fmt.Errorf("algorithm execution error: %v", err)
|
||||
}
|
||||
|
||||
@@ -104,19 +103,3 @@ func (p *python) Run() error {
|
||||
|
||||
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 err := p.cmd.Process.Kill(); err != nil {
|
||||
return fmt.Errorf("error stopping algorithm: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -50,7 +50,7 @@ func TestNewAlgorithm(t *testing.T) {
|
||||
algoFile := "algorithm.py"
|
||||
args := []string{"--arg1", "value1"}
|
||||
|
||||
algo := NewAlgorithm(logger, eventsSvc, runtime, requirementsFile, algoFile, args, "")
|
||||
algo := NewAlgorithm(logger, eventsSvc, runtime, requirementsFile, algoFile, args)
|
||||
|
||||
p, ok := algo.(*python)
|
||||
if !ok {
|
||||
|
||||
@@ -24,13 +24,12 @@ type wasm struct {
|
||||
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 {
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string) algorithm.Algorithm {
|
||||
return &wasm{
|
||||
algoFile: algoFile,
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
|
||||
stdout: &logging.Stdout{Logger: logger},
|
||||
args: args,
|
||||
}
|
||||
@@ -39,33 +38,17 @@ func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, args []string,
|
||||
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
|
||||
cmd := exec.Command(wasmRuntime, args...)
|
||||
cmd.Stderr = w.stderr
|
||||
cmd.Stdout = w.stdout
|
||||
|
||||
if err := w.cmd.Start(); err != nil {
|
||||
if err := cmd.Start(); err != nil {
|
||||
return fmt.Errorf("error starting algorithm: %v", err)
|
||||
}
|
||||
|
||||
if err := w.cmd.Wait(); err != nil {
|
||||
if err := 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 err := w.cmd.Process.Kill(); err != nil {
|
||||
return fmt.Errorf("error stopping algorithm: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -18,7 +18,7 @@ func TestNewAlgorithm(t *testing.T) {
|
||||
algoFile := "test.wasm"
|
||||
args := []string{"arg1", "arg2"}
|
||||
|
||||
algo := NewAlgorithm(logger, eventsSvc, args, algoFile, "")
|
||||
algo := NewAlgorithm(logger, eventsSvc, algoFile, args)
|
||||
|
||||
w, ok := algo.(*wasm)
|
||||
if !ok {
|
||||
@@ -54,7 +54,7 @@ func TestRunError(t *testing.T) {
|
||||
algoFile := "test.wasm"
|
||||
args := []string{"arg1", "arg2"}
|
||||
|
||||
w := NewAlgorithm(logger, eventsSvc, args, algoFile, "").(*wasm)
|
||||
w := NewAlgorithm(logger, eventsSvc, algoFile, args).(*wasm)
|
||||
|
||||
err := w.Run()
|
||||
if err == nil {
|
||||
|
||||
@@ -7,7 +7,6 @@ import (
|
||||
|
||||
"github.com/go-kit/kit/endpoint"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
config "github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
)
|
||||
|
||||
func algoEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
@@ -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, config.AttestationType(req.AttType))
|
||||
file, err := svc.Attestation(ctx, req.ReportData)
|
||||
if err != nil {
|
||||
return attestationRes{}, err
|
||||
}
|
||||
|
||||
@@ -9,7 +9,6 @@ import (
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/agent/mocks"
|
||||
config "github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"golang.org/x/crypto/sha3"
|
||||
)
|
||||
|
||||
@@ -142,11 +141,11 @@ func TestAttestationEndpoint(t *testing.T) {
|
||||
}{
|
||||
{
|
||||
name: "Success",
|
||||
req: attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: config.SNP},
|
||||
req: attestationReq{ReportData: sha3.Sum512([]byte("report data"))},
|
||||
},
|
||||
{
|
||||
name: "Service Error",
|
||||
req: attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: config.SNP},
|
||||
req: attestationReq{ReportData: sha3.Sum512([]byte("report data"))},
|
||||
expectedErr: true,
|
||||
},
|
||||
}
|
||||
@@ -154,9 +153,9 @@ func TestAttestationEndpoint(t *testing.T) {
|
||||
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()
|
||||
svc.On("Attestation", context.Background(), tt.req.ReportData).Return([]byte{}, errors.New("")).Once()
|
||||
} else {
|
||||
svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, nil).Once()
|
||||
svc.On("Attestation", context.Background(), tt.req.ReportData).Return([]byte{}, nil).Once()
|
||||
}
|
||||
endpoint := attestationEndpoint(svc)
|
||||
res, err := endpoint(context.Background(), tt.req)
|
||||
|
||||
@@ -4,10 +4,6 @@ package grpc
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
config "github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
)
|
||||
|
||||
type algoReq struct {
|
||||
@@ -42,16 +38,9 @@ func (req resultReq) validate() error {
|
||||
}
|
||||
|
||||
type attestationReq struct {
|
||||
TeeNonce [quoteprovider.Nonce]byte
|
||||
VtpmNonce [vtpm.Nonce]byte
|
||||
AttType config.AttestationType
|
||||
ReportData [64]byte
|
||||
}
|
||||
|
||||
func (req attestationReq) validate() error {
|
||||
switch req.AttType {
|
||||
case config.SNP, config.VTPM, config.SNPvTPM:
|
||||
return nil
|
||||
default:
|
||||
return errors.New("invalid attestation type in attestation request")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -11,9 +11,6 @@ import (
|
||||
|
||||
"github.com/go-kit/kit/transport/grpc"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
config "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"
|
||||
@@ -24,11 +21,6 @@ const (
|
||||
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")
|
||||
)
|
||||
|
||||
var _ agent.AgentServiceServer = (*grpcServer)(nil)
|
||||
|
||||
type grpcServer struct {
|
||||
@@ -104,20 +96,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: config.AttestationType(req.Type)}, nil
|
||||
return attestationReq{ReportData: [agent.ReportDataSize]byte(req.ReportData)}, nil
|
||||
}
|
||||
|
||||
func encodeAttestationResponse(_ context.Context, response interface{}) (interface{}, error) {
|
||||
|
||||
@@ -11,9 +11,6 @@ import (
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/agent/mocks"
|
||||
config "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"
|
||||
)
|
||||
@@ -154,12 +151,10 @@ func TestAttestation(t *testing.T) {
|
||||
mockStream := &MockAgentService_AttestationServer{ctx: context.Background()}
|
||||
mockStream.On("Send", mock.AnythingOfType("*agent.AttestationResponse")).Return(nil)
|
||||
|
||||
reportData := [quoteprovider.Nonce]byte{}
|
||||
vtpmNonce := [vtpm.Nonce]byte{}
|
||||
attestationType := config.SNP
|
||||
mockService.On("Attestation", mock.Anything, reportData, vtpmNonce, attestationType).Return([]byte("attestation data"), nil)
|
||||
reportData := [agent.ReportDataSize]byte{}
|
||||
mockService.On("Attestation", mock.Anything, reportData).Return([]byte("attestation data"), nil)
|
||||
|
||||
err := server.Attestation(&agent.AttestationRequest{TeeNonce: reportData[:]}, mockStream)
|
||||
err := server.Attestation(&agent.AttestationRequest{ReportData: reportData[:]}, mockStream)
|
||||
assert.NoError(t, err)
|
||||
|
||||
mockService.AssertExpectations(t)
|
||||
@@ -204,11 +199,11 @@ func TestEncodeResultResponse(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestDecodeAttestationRequest(t *testing.T) {
|
||||
nonce := [quoteprovider.Nonce]byte{}
|
||||
req := &agent.AttestationRequest{TeeNonce: nonce[:]}
|
||||
reportData := [agent.ReportDataSize]byte{}
|
||||
req := &agent.AttestationRequest{ReportData: reportData[:]}
|
||||
decoded, err := decodeAttestationRequest(context.Background(), req)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, attestationReq{TeeNonce: nonce}, decoded)
|
||||
assert.Equal(t, attestationReq{ReportData: reportData}, decoded)
|
||||
}
|
||||
|
||||
func TestEncodeAttestationResponse(t *testing.T) {
|
||||
|
||||
+2
-42
@@ -13,9 +13,6 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
config "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.Info(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))
|
||||
@@ -106,7 +66,7 @@ func (lm *loggingMiddleware) Result(ctx context.Context) (response []byte, err e
|
||||
return lm.svc.Result(ctx)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType config.AttestationType) (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,5 +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)
|
||||
return lm.svc.Attestation(ctx, reportData)
|
||||
}
|
||||
|
||||
+2
-35
@@ -12,9 +12,6 @@ import (
|
||||
|
||||
"github.com/go-kit/kit/metrics"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
config "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)
|
||||
@@ -92,11 +59,11 @@ func (ms *metricsMiddleware) Result(ctx context.Context) ([]byte, error) {
|
||||
return ms.svc.Result(ctx)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType config.AttestationType) ([]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)
|
||||
return ms.svc.Attestation(ctx, reportData)
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
var _ fmt.Stringer = (*Datasets)(nil)
|
||||
|
||||
type AgentConfig struct {
|
||||
LogLevel string `json:"log_level,omitempty"`
|
||||
Host string `json:"host,omitempty"`
|
||||
Port string `json:"port,omitempty"`
|
||||
CertFile string `json:"cert_file,omitempty"`
|
||||
@@ -29,6 +30,7 @@ type Computation struct {
|
||||
Datasets Datasets `json:"datasets,omitempty"`
|
||||
Algorithm Algorithm `json:"algorithm,omitempty"`
|
||||
ResultConsumers []ResultConsumer `json:"result_consumers,omitempty"`
|
||||
AgentConfig AgentConfig `json:"agent_config,omitempty"`
|
||||
}
|
||||
|
||||
type ResultConsumer struct {
|
||||
|
||||
@@ -106,6 +106,7 @@ func TestDecompressToContext(t *testing.T) {
|
||||
|
||||
func TestAgentConfigJSON(t *testing.T) {
|
||||
config := AgentConfig{
|
||||
LogLevel: "info",
|
||||
Host: "localhost",
|
||||
Port: "8080",
|
||||
CertFile: "cert.pem",
|
||||
|
||||
@@ -1,350 +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"
|
||||
"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) (cvms.Service_ProcessClient, error)
|
||||
}
|
||||
|
||||
// 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) (cvms.Service_ProcessClient, error)) (*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,
|
||||
}, 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)
|
||||
|
||||
stream, err := client.reconnectFn(ctx)
|
||||
if err != nil {
|
||||
client.logger.Error("Failed to reconnect", "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
client.mu.Lock()
|
||||
client.stream = stream
|
||||
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)
|
||||
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,
|
||||
Host: runReq.AgentConfig.Host,
|
||||
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
|
||||
}
|
||||
@@ -1,209 +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"
|
||||
"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
|
||||
}()
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (cvms.Service_ProcessClient, error) { return nil, nil })
|
||||
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()
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (cvms.Service_ProcessClient, error) { return nil, nil })
|
||||
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()
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (cvms.Service_ProcessClient, error) { return nil, nil })
|
||||
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)
|
||||
}
|
||||
@@ -1,5 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Package grpc contains implementation of kit service gRPC API.
|
||||
package grpc
|
||||
@@ -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
|
||||
}
|
||||
@@ -1,234 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery v2.43.2. 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 given 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 given 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
|
||||
}
|
||||
@@ -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
@@ -1,114 +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;
|
||||
}
|
||||
}
|
||||
|
||||
message ServerStreamMessage {
|
||||
oneof message {
|
||||
RunReqChunks runReqChunks = 1;
|
||||
ComputationRunReq runReq = 2;
|
||||
StopComputation stopComputation = 3;
|
||||
AgentStateReq agentStateReq = 4;
|
||||
}
|
||||
}
|
||||
|
||||
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 host = 2;
|
||||
string cert_file = 3;
|
||||
string key_file = 4;
|
||||
string client_ca_file = 5;
|
||||
string server_ca_file = 6;
|
||||
string log_level = 7;
|
||||
bool attested_tls = 8;
|
||||
}
|
||||
@@ -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",
|
||||
}
|
||||
@@ -1,96 +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"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
|
||||
"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
|
||||
}
|
||||
|
||||
func NewServer(logger *slog.Logger, svc agent.Service) AgentServer {
|
||||
return &agentServer{
|
||||
logger: logger,
|
||||
svc: svc,
|
||||
}
|
||||
}
|
||||
|
||||
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: cfg.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
|
||||
}
|
||||
|
||||
qp, err := quoteprovider.GetLeveledQuoteProvider()
|
||||
if err != nil {
|
||||
as.logger.Error(fmt.Sprintf("failed to create quote provider %s", err.Error()))
|
||||
return err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
as.gs = grpcserver.New(ctx, cancel, svcName, agentGrpcServerConfig, registerAgentServiceServer, as.logger, qp, authSvc)
|
||||
|
||||
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 {
|
||||
return as.gs.Stop()
|
||||
}
|
||||
@@ -1,130 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery v2.43.2. 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 given 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
|
||||
}
|
||||
+24
-19
@@ -4,38 +4,43 @@ package events
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"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 io.Writer
|
||||
}
|
||||
|
||||
type Service interface {
|
||||
SendEvent(cmpID, event, status string, details json.RawMessage)
|
||||
SendEvent(event, status string, details json.RawMessage) error
|
||||
}
|
||||
|
||||
func New(svc string, queue chan *cvms.ClientStreamMessage) (Service, error) {
|
||||
func New(svc, computationID string, conn io.Writer) (Service, error) {
|
||||
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 := EventsLogs{Message: &EventsLogs_AgentEvent{AgentEvent: &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
|
||||
}
|
||||
_, err = s.conn.Write(protoBody)
|
||||
return err
|
||||
}
|
||||
|
||||
+79
-36
@@ -3,8 +3,8 @@
|
||||
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.0
|
||||
// protoc v5.29.0
|
||||
// protoc-gen-go v1.34.2
|
||||
// protoc v5.28.1
|
||||
// source: agent/events/events.proto
|
||||
|
||||
package events
|
||||
@@ -25,22 +25,25 @@ const (
|
||||
)
|
||||
|
||||
type AgentEvent struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
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)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_events_events_proto_msgTypes[0]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *AgentEvent) String() string {
|
||||
@@ -51,7 +54,7 @@ func (*AgentEvent) ProtoMessage() {}
|
||||
|
||||
func (x *AgentEvent) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_events_events_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)
|
||||
@@ -109,20 +112,23 @@ func (x *AgentEvent) GetStatus() string {
|
||||
}
|
||||
|
||||
type AgentLog struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
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)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_events_events_proto_msgTypes[1]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *AgentLog) String() string {
|
||||
@@ -133,7 +139,7 @@ func (*AgentLog) ProtoMessage() {}
|
||||
|
||||
func (x *AgentLog) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_events_events_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)
|
||||
@@ -177,21 +183,24 @@ func (x *AgentLog) GetTimestamp() *timestamppb.Timestamp {
|
||||
}
|
||||
|
||||
type EventsLogs struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
// Types that are valid to be assigned to Message:
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
// Types that are assignable to Message:
|
||||
//
|
||||
// *EventsLogs_AgentLog
|
||||
// *EventsLogs_AgentEvent
|
||||
Message isEventsLogs_Message `protobuf_oneof:"message"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
Message isEventsLogs_Message `protobuf_oneof:"message"`
|
||||
}
|
||||
|
||||
func (x *EventsLogs) Reset() {
|
||||
*x = EventsLogs{}
|
||||
mi := &file_agent_events_events_proto_msgTypes[2]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_events_events_proto_msgTypes[2]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *EventsLogs) String() string {
|
||||
@@ -202,7 +211,7 @@ func (*EventsLogs) ProtoMessage() {}
|
||||
|
||||
func (x *EventsLogs) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_events_events_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)
|
||||
@@ -217,27 +226,23 @@ 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
|
||||
func (m *EventsLogs) GetMessage() isEventsLogs_Message {
|
||||
if m != nil {
|
||||
return m.Message
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *EventsLogs) GetAgentLog() *AgentLog {
|
||||
if x != nil {
|
||||
if x, ok := x.Message.(*EventsLogs_AgentLog); ok {
|
||||
return x.AgentLog
|
||||
}
|
||||
if x, ok := x.GetMessage().(*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
|
||||
}
|
||||
if x, ok := x.GetMessage().(*EventsLogs_AgentEvent); ok {
|
||||
return x.AgentEvent
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -337,6 +342,44 @@ func file_agent_events_events_proto_init() {
|
||||
if File_agent_events_events_proto != nil {
|
||||
return
|
||||
}
|
||||
if !protoimpl.UnsafeEnabled {
|
||||
file_agent_events_events_proto_msgTypes[0].Exporter = func(v any, i int) any {
|
||||
switch v := v.(*AgentEvent); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_events_events_proto_msgTypes[1].Exporter = func(v any, i int) any {
|
||||
switch v := v.(*AgentLog); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_events_events_proto_msgTypes[2].Exporter = func(v any, i int) any {
|
||||
switch v := v.(*EventsLogs); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
file_agent_events_events_proto_msgTypes[2].OneofWrappers = []any{
|
||||
(*EventsLogs_AgentLog)(nil),
|
||||
(*EventsLogs_AgentEvent)(nil),
|
||||
|
||||
+43
-17
@@ -3,36 +3,62 @@
|
||||
package events
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
type mockConn struct {
|
||||
writeErr error
|
||||
buf bytes.Buffer
|
||||
}
|
||||
|
||||
func (m *mockConn) Write(p []byte) (n int, err error) {
|
||||
if m.writeErr != nil {
|
||||
return 0, m.writeErr
|
||||
}
|
||||
return m.buf.Write(p)
|
||||
}
|
||||
|
||||
func TestSendEventSuccess(t *testing.T) {
|
||||
queue := make(chan *cvms.ClientStreamMessage, 1)
|
||||
svc, err := New("test_service", queue)
|
||||
mockConnection := &mockConn{}
|
||||
|
||||
svc, err := New("test_service", "12345", mockConnection)
|
||||
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)
|
||||
err = svc.SendEvent("test_event", "success", details)
|
||||
assert.NoError(t, err)
|
||||
|
||||
now := time.Now()
|
||||
eventTimestamp := msg.GetAgentEvent().GetTimestamp().AsTime()
|
||||
assert.WithinDuration(t, now, eventTimestamp, 1*time.Second)
|
||||
}()
|
||||
var writtenMessage EventsLogs
|
||||
err = proto.Unmarshal(mockConnection.buf.Bytes(), &writtenMessage)
|
||||
assert.NoError(t, err)
|
||||
|
||||
svc.SendEvent("testid", "test_event", "success", details)
|
||||
assert.Equal(t, "test_event", writtenMessage.GetAgentEvent().EventType)
|
||||
assert.Equal(t, "12345", writtenMessage.GetAgentEvent().ComputationId)
|
||||
assert.Equal(t, "test_service", writtenMessage.GetAgentEvent().Originator)
|
||||
assert.Equal(t, "success", writtenMessage.GetAgentEvent().Status)
|
||||
|
||||
time.Sleep(1 * time.Second)
|
||||
now := time.Now()
|
||||
eventTimestamp := writtenMessage.GetAgentEvent().GetTimestamp().AsTime()
|
||||
assert.WithinDuration(t, now, eventTimestamp, 1*time.Second)
|
||||
}
|
||||
|
||||
func TestSendEventFailure(t *testing.T) {
|
||||
mockConnection := &mockConn{writeErr: errors.New("write error")}
|
||||
|
||||
svc, err := New("test_service", "12345", mockConnection)
|
||||
assert.NoError(t, err)
|
||||
|
||||
details := json.RawMessage(`{"key": "value"}`)
|
||||
|
||||
err = svc.SendEvent("test_event", "failure", details)
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, "write error", err.Error())
|
||||
}
|
||||
|
||||
@@ -24,9 +24,22 @@ 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)
|
||||
// SendEvent provides a mock function with given fields: event, status, details
|
||||
func (_m *Service) SendEvent(event string, status string, details json.RawMessage) error {
|
||||
ret := _m.Called(event, status, details)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for SendEvent")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if rf, ok := ret.Get(0).(func(string, string, json.RawMessage) error); ok {
|
||||
r0 = rf(event, status, details)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
|
||||
return r0
|
||||
}
|
||||
|
||||
// Service_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
|
||||
@@ -35,27 +48,26 @@ type Service_SendEvent_Call struct {
|
||||
}
|
||||
|
||||
// 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 (_e *Service_Expecter) SendEvent(event interface{}, status interface{}, details interface{}) *Service_SendEvent_Call {
|
||||
return &Service_SendEvent_Call{Call: _e.mock.On("SendEvent", event, status, details)}
|
||||
}
|
||||
|
||||
func (_c *Service_SendEvent_Call) Run(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
|
||||
func (_c *Service_SendEvent_Call) Run(run func(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))
|
||||
run(args[0].(string), args[1].(string), args[2].(json.RawMessage))
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_SendEvent_Call) Return() *Service_SendEvent_Call {
|
||||
_c.Call.Return()
|
||||
func (_c *Service_SendEvent_Call) Return(_a0 error) *Service_SendEvent_Call {
|
||||
_c.Call.Return(_a0)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_SendEvent_Call) RunAndReturn(run func(string, string, string, json.RawMessage)) *Service_SendEvent_Call {
|
||||
func (_c *Service_SendEvent_Call) RunAndReturn(run func(string, string, json.RawMessage) error) *Service_SendEvent_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
+16
-157
@@ -6,11 +6,10 @@
|
||||
package mocks
|
||||
|
||||
import (
|
||||
agent "github.com/ultravioletrs/cocos/agent"
|
||||
config "github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
|
||||
context "context"
|
||||
|
||||
agent "github.com/ultravioletrs/cocos/agent"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
@@ -74,9 +73,9 @@ func (_c *Service_Algo_Call) RunAndReturn(run func(context.Context, agent.Algori
|
||||
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 config.AttestationType) ([]byte, error) {
|
||||
ret := _m.Called(ctx, reportData, nonce, attType)
|
||||
// Attestation provides a mock function with given fields: ctx, reportData
|
||||
func (_m *Service) Attestation(ctx context.Context, reportData [64]byte) ([]byte, error) {
|
||||
ret := _m.Called(ctx, reportData)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Attestation")
|
||||
@@ -84,19 +83,19 @@ func (_m *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [
|
||||
|
||||
var r0 []byte
|
||||
var r1 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, config.AttestationType) ([]byte, error)); ok {
|
||||
return rf(ctx, reportData, nonce, attType)
|
||||
if rf, ok := ret.Get(0).(func(context.Context, [64]byte) ([]byte, error)); ok {
|
||||
return rf(ctx, reportData)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, config.AttestationType) []byte); ok {
|
||||
r0 = rf(ctx, reportData, nonce, attType)
|
||||
if rf, ok := ret.Get(0).(func(context.Context, [64]byte) []byte); ok {
|
||||
r0 = rf(ctx, reportData)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).([]byte)
|
||||
}
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, [64]byte, [32]byte, config.AttestationType) error); ok {
|
||||
r1 = rf(ctx, reportData, nonce, attType)
|
||||
if rf, ok := ret.Get(1).(func(context.Context, [64]byte) error); ok {
|
||||
r1 = rf(ctx, reportData)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
@@ -112,15 +111,13 @@ type Service_Attestation_Call struct {
|
||||
// Attestation is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - reportData [64]byte
|
||||
// - nonce [32]byte
|
||||
// - attType config.AttestationType
|
||||
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 (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{}) *Service_Attestation_Call {
|
||||
return &Service_Attestation_Call{Call: _e.mock.On("Attestation", ctx, reportData)}
|
||||
}
|
||||
|
||||
func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType config.AttestationType)) *Service_Attestation_Call {
|
||||
func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte)) *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].(config.AttestationType))
|
||||
run(args[0].(context.Context), args[1].([64]byte))
|
||||
})
|
||||
return _c
|
||||
}
|
||||
@@ -130,7 +127,7 @@ func (_c *Service_Attestation_Call) Return(_a0 []byte, _a1 error) *Service_Attes
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Attestation_Call) RunAndReturn(run func(context.Context, [64]byte, [32]byte, config.AttestationType) ([]byte, error)) *Service_Attestation_Call {
|
||||
func (_c *Service_Attestation_Call) RunAndReturn(run func(context.Context, [64]byte) ([]byte, error)) *Service_Attestation_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -182,53 +179,6 @@ func (_c *Service_Data_Call) RunAndReturn(run func(context.Context, agent.Datase
|
||||
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)
|
||||
@@ -287,97 +237,6 @@ func (_c *Service_Result_Call) RunAndReturn(run func(context.Context) ([]byte, e
|
||||
return _c
|
||||
}
|
||||
|
||||
// State provides a mock function with given 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 {
|
||||
|
||||
+37
-112
@@ -23,9 +23,6 @@ import (
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
"github.com/ultravioletrs/cocos/agent/statemachine"
|
||||
"github.com/ultravioletrs/cocos/internal"
|
||||
config "github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"golang.org/x/crypto/sha3"
|
||||
)
|
||||
|
||||
@@ -72,6 +69,8 @@ const (
|
||||
)
|
||||
|
||||
const (
|
||||
// ReportDataSize is the size of the report data expected by the attestation service.
|
||||
ReportDataSize = 64
|
||||
algoFilePermission = 0o700
|
||||
)
|
||||
|
||||
@@ -100,50 +99,41 @@ var (
|
||||
ErrAllResultsConsumed = errors.New("all results have been consumed by declared consumers")
|
||||
// ErrAttestationFailed attestation failed.
|
||||
ErrAttestationFailed = errors.New("failed to get raw quote")
|
||||
// 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")
|
||||
)
|
||||
|
||||
// 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 config.AttestationType) ([]byte, error)
|
||||
State() string
|
||||
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.
|
||||
quoteProvider client.LeveledQuoteProvider // 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.
|
||||
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.
|
||||
quoteProvider client.QuoteProvider // Provider for generating attestation quotes.
|
||||
logger *slog.Logger // Logger for the agent service.
|
||||
resultsConsumed bool // Indicates if the results have been consumed.
|
||||
}
|
||||
|
||||
var _ Service = (*agentService)(nil)
|
||||
|
||||
// New instantiates the agent service implementation.
|
||||
func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, quoteProvider client.LeveledQuoteProvider, vmlp int) Service {
|
||||
func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp Computation, quoteProvider client.QuoteProvider) Service {
|
||||
sm := statemachine.NewStateMachine(Idle)
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
svc := &agentService{
|
||||
sm: sm,
|
||||
eventSvc: eventSvc,
|
||||
quoteProvider: quoteProvider,
|
||||
logger: logger,
|
||||
cancel: cancel,
|
||||
vmpl: vmlp,
|
||||
computation: cmp,
|
||||
}
|
||||
|
||||
transitions := []statemachine.Transition{
|
||||
@@ -151,6 +141,13 @@ func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, quot
|
||||
{From: ReceivingManifest, Event: ManifestReceived, To: ReceivingAlgorithm},
|
||||
}
|
||||
|
||||
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})
|
||||
}
|
||||
|
||||
transitions = append(transitions, []statemachine.Transition{
|
||||
{From: Running, Event: RunComplete, To: ConsumingResults},
|
||||
{From: Running, Event: RunFailed, To: Failed},
|
||||
@@ -161,6 +158,8 @@ func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, quot
|
||||
sm.AddTransition(t)
|
||||
}
|
||||
|
||||
sm.SetAction(Idle, svc.publishEvent(IdleState.String()))
|
||||
sm.SetAction(ReceivingManifest, svc.publishEvent(InProgress.String()))
|
||||
sm.SetAction(ReceivingAlgorithm, svc.publishEvent(InProgress.String()))
|
||||
sm.SetAction(ReceivingData, svc.publishEvent(InProgress.String()))
|
||||
sm.SetAction(Running, svc.runComputation)
|
||||
@@ -174,71 +173,11 @@ func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, quot
|
||||
}
|
||||
}()
|
||||
sm.SendEvent(Start)
|
||||
defer sm.SendEvent(ManifestReceived)
|
||||
|
||||
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
|
||||
}
|
||||
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})
|
||||
}
|
||||
|
||||
for _, t := range transitions {
|
||||
as.sm.AddTransition(t)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (as *agentService) StopComputation(ctx context.Context) error {
|
||||
as.mu.Lock()
|
||||
defer as.mu.Unlock()
|
||||
|
||||
as.cancel()
|
||||
|
||||
if err := as.algorithm.Stop(); err != nil {
|
||||
return fmt.Errorf("error stopping computation: %v", err)
|
||||
}
|
||||
|
||||
sm := statemachine.NewStateMachine(Idle)
|
||||
|
||||
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 = sm
|
||||
as.computation = Computation{}
|
||||
as.algorithm = nil
|
||||
as.result = nil
|
||||
as.runError = nil
|
||||
as.resultsConsumed = false
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
|
||||
if as.sm.GetState() != ReceivingAlgorithm {
|
||||
return ErrStateNotReady
|
||||
@@ -286,7 +225,7 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
|
||||
|
||||
switch algoType {
|
||||
case string(algorithm.AlgoTypeBin):
|
||||
as.algorithm = binary.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args, as.computation.ID)
|
||||
as.algorithm = binary.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args)
|
||||
case string(algorithm.AlgoTypePython):
|
||||
var requirementsFile string
|
||||
if len(algo.Requirements) > 0 {
|
||||
@@ -304,11 +243,11 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
|
||||
requirementsFile = fr.Name()
|
||||
}
|
||||
runtime := python.PythonRunTimeFromContext(ctx)
|
||||
as.algorithm = python.NewAlgorithm(as.logger, as.eventSvc, runtime, requirementsFile, f.Name(), args, as.computation.ID)
|
||||
as.algorithm = python.NewAlgorithm(as.logger, as.eventSvc, runtime, requirementsFile, f.Name(), args)
|
||||
case string(algorithm.AlgoTypeWasm):
|
||||
as.algorithm = wasm.NewAlgorithm(as.logger, as.eventSvc, args, f.Name(), as.computation.ID)
|
||||
as.algorithm = wasm.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args)
|
||||
case string(algorithm.AlgoTypeDocker):
|
||||
as.algorithm = docker.NewAlgorithm(as.logger, as.eventSvc, f.Name(), as.computation.ID)
|
||||
as.algorithm = docker.NewAlgorithm(as.logger, as.eventSvc, f.Name())
|
||||
}
|
||||
|
||||
if err := os.Mkdir(algorithm.DatasetsDir, 0o755); err != nil {
|
||||
@@ -402,29 +341,13 @@ func (as *agentService) Result(ctx context.Context) ([]byte, error) {
|
||||
return as.result, as.runError
|
||||
}
|
||||
|
||||
func (as *agentService) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType config.AttestationType) ([]byte, error) {
|
||||
switch attType {
|
||||
case config.SNP:
|
||||
rawQuote, err := as.quoteProvider.GetRawQuoteAtLevel(reportData, uint(as.vmpl))
|
||||
if err != nil {
|
||||
return []byte{}, err
|
||||
}
|
||||
return rawQuote, nil
|
||||
case config.VTPM:
|
||||
vTPMQuote, err := vtpm.Attest(reportData[:], nonce[:], false)
|
||||
if err != nil {
|
||||
return []byte{}, err
|
||||
}
|
||||
return vTPMQuote, nil
|
||||
case config.SNPvTPM:
|
||||
vTPMQuote, err := vtpm.Attest(reportData[:], nonce[:], true)
|
||||
if err != nil {
|
||||
return []byte{}, err
|
||||
}
|
||||
return vTPMQuote, nil
|
||||
default:
|
||||
return []byte{}, ErrAttestationType
|
||||
func (as *agentService) Attestation(ctx context.Context, reportData [ReportDataSize]byte) ([]byte, error) {
|
||||
rawQuote, err := as.quoteProvider.GetRawQuote(reportData)
|
||||
if err != nil {
|
||||
return []byte{}, err
|
||||
}
|
||||
|
||||
return rawQuote, nil
|
||||
}
|
||||
|
||||
func (as *agentService) runComputation(state statemachine.State) {
|
||||
@@ -477,6 +400,8 @@ func (as *agentService) runComputation(state statemachine.State) {
|
||||
|
||||
func (as *agentService) publishEvent(status string) statemachine.Action {
|
||||
return func(state statemachine.State) {
|
||||
as.eventSvc.SendEvent(as.computation.ID, state.String(), status, json.RawMessage{})
|
||||
if err := as.eventSvc.SendEvent(state.String(), status, json.RawMessage{}); err != nil {
|
||||
as.logger.Warn(err.Error())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+39
-37
@@ -22,7 +22,6 @@ import (
|
||||
smmocks "github.com/ultravioletrs/cocos/agent/statemachine/mocks"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
|
||||
mocks2 "github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"golang.org/x/crypto/sha3"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
@@ -36,7 +35,12 @@ var (
|
||||
const datasetFile = "iris.csv"
|
||||
|
||||
func TestAlgo(t *testing.T) {
|
||||
qp, err := quoteprovider.GetLeveledQuoteProvider()
|
||||
events := new(mocks.Service)
|
||||
|
||||
evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
defer evCall.Unset()
|
||||
|
||||
qp, err := quoteprovider.GetQuoteProvider()
|
||||
require.NoError(t, err)
|
||||
|
||||
algo, err := os.ReadFile(algoPath)
|
||||
@@ -116,15 +120,9 @@ func TestAlgo(t *testing.T) {
|
||||
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, qp, 0)
|
||||
|
||||
err := svc.InitComputation(ctx, testComputation(t))
|
||||
require.NoError(t, err)
|
||||
svc := New(ctx, mglog.NewMock(), events, testComputation(t), qp)
|
||||
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
|
||||
@@ -140,7 +138,12 @@ func TestAlgo(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestData(t *testing.T) {
|
||||
qp, err := quoteprovider.GetLeveledQuoteProvider()
|
||||
events := new(mocks.Service)
|
||||
|
||||
evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
defer evCall.Unset()
|
||||
|
||||
qp, err := quoteprovider.GetQuoteProvider()
|
||||
require.NoError(t, err)
|
||||
|
||||
algo, err := os.ReadFile(algoPath)
|
||||
@@ -206,9 +209,6 @@ func TestData(t *testing.T) {
|
||||
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)
|
||||
}
|
||||
@@ -216,16 +216,13 @@ func TestData(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
svc := New(ctx, mglog.NewMock(), events, qp, 0)
|
||||
|
||||
err := svc.InitComputation(ctx, testComputation(t))
|
||||
require.NoError(t, err)
|
||||
comp := testComputation(t)
|
||||
|
||||
svc := New(ctx, mglog.NewMock(), events, comp, qp)
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
|
||||
if tc.err != ErrStateNotReady {
|
||||
err = svc.Algo(ctx, alg)
|
||||
require.NoError(t, err)
|
||||
_ = svc.Algo(ctx, alg)
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
}
|
||||
err = svc.Data(ctx, tc.data)
|
||||
@@ -241,7 +238,12 @@ func TestData(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestResult(t *testing.T) {
|
||||
qp, err := quoteprovider.GetLeveledQuoteProvider()
|
||||
events := new(mocks.Service)
|
||||
|
||||
evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
defer evCall.Unset()
|
||||
|
||||
qp, err := quoteprovider.GetQuoteProvider()
|
||||
require.NoError(t, err)
|
||||
|
||||
cases := []struct {
|
||||
@@ -283,9 +285,6 @@ func TestResult(t *testing.T) {
|
||||
}
|
||||
|
||||
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),
|
||||
@@ -324,57 +323,55 @@ func TestResult(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAttestation(t *testing.T) {
|
||||
qp := new(mocks2.LeveledQuoteProvider)
|
||||
events := new(mocks.Service)
|
||||
qp := new(mocks2.QuoteProvider)
|
||||
|
||||
evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
defer evCall.Unset()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
reportData [quoteprovider.Nonce]byte
|
||||
nonce [vtpm.Nonce]byte
|
||||
reportData [ReportDataSize]byte
|
||||
rawQuote []uint8
|
||||
err error
|
||||
}{
|
||||
{
|
||||
name: "Test attestation successful",
|
||||
reportData: generateReportData(),
|
||||
nonce: [32]byte{},
|
||||
rawQuote: make([]uint8, 0),
|
||||
err: nil,
|
||||
},
|
||||
{
|
||||
name: "Test attestation failed",
|
||||
reportData: generateReportData(),
|
||||
nonce: [32]byte{},
|
||||
rawQuote: nil,
|
||||
err: ErrAttestationFailed,
|
||||
},
|
||||
}
|
||||
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 := qp.On("GetRawQuoteAtLevel", mock.Anything, mock.Anything).Return(tc.rawQuote, tc.err)
|
||||
getQuote := qp.On("GetRawQuote", mock.Anything).Return(tc.rawQuote, tc.err)
|
||||
if tc.err != ErrAttestationFailed {
|
||||
getQuote = qp.On("GetRawQuoteAtLevel", mock.Anything, mock.Anything).Return(tc.nonce, nil)
|
||||
getQuote = qp.On("GetRawQuote", mock.Anything).Return(tc.reportData, nil)
|
||||
}
|
||||
defer getQuote.Unset()
|
||||
|
||||
svc := New(ctx, mglog.NewMock(), events, qp, 0)
|
||||
svc := New(ctx, mglog.NewMock(), events, testComputation(t), qp)
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
_, err := svc.Attestation(ctx, tc.reportData, tc.nonce, 0)
|
||||
_, err := svc.Attestation(ctx, tc.reportData)
|
||||
assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func generateReportData() [quoteprovider.Nonce]byte {
|
||||
bytes := make([]byte, quoteprovider.Nonce)
|
||||
func generateReportData() [ReportDataSize]byte {
|
||||
bytes := make([]byte, ReportDataSize)
|
||||
_, err := rand.Read(bytes)
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to generate random bytes: %v", err)
|
||||
@@ -400,5 +397,10 @@ func testComputation(t *testing.T) Computation {
|
||||
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")}},
|
||||
AgentConfig: AgentConfig{
|
||||
Port: "7002",
|
||||
LogLevel: "debug",
|
||||
AttestedTls: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
Binary file not shown.
@@ -100,22 +100,3 @@ When defining the manifest dataset and algorithm checksums are required. This ca
|
||||
```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
|
||||
```
|
||||
+86
-432
@@ -3,11 +3,9 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"crypto"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -19,18 +17,11 @@ import (
|
||||
"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"
|
||||
sevVerify "github.com/google/go-sev-guest/verify"
|
||||
tpmAttest "github.com/google/go-tpm-tools/proto/attest"
|
||||
"github.com/google/go-tpm-tools/server"
|
||||
"github.com/google/go-tpm/legacy/tpm2"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/spf13/pflag"
|
||||
config "github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"google.golang.org/protobuf/encoding/protojson"
|
||||
"google.golang.org/protobuf/encoding/prototext"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/wrapperspb"
|
||||
)
|
||||
|
||||
@@ -51,12 +42,9 @@ const (
|
||||
size48 = 48
|
||||
size64 = 64
|
||||
attestationFilePath = "attestation.bin"
|
||||
vtpmFilePath = "../quote.dat"
|
||||
attestationJson = "attestation.json"
|
||||
sevProductNameMilan = "Milan"
|
||||
sevProductNameGenoa = "Genoa"
|
||||
FormatBinaryPB = "binarypb"
|
||||
FormatTextProto = "textproto"
|
||||
exampleJSONConfig = `
|
||||
{
|
||||
"rootOfTrust":{
|
||||
@@ -110,47 +98,29 @@ const (
|
||||
}
|
||||
}
|
||||
`
|
||||
SNP = "snp"
|
||||
VTPM = "vtpm"
|
||||
SNPvTPM = "snp-vtpm"
|
||||
)
|
||||
|
||||
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
|
||||
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}
|
||||
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
|
||||
getTextProtoAttestation bool
|
||||
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}
|
||||
getJsonAttestation bool
|
||||
errReportSize = errors.New("attestation contents too small")
|
||||
)
|
||||
|
||||
var errEmptyFile = errors.New("input file is empty")
|
||||
|
||||
var marshalOptions = prototext.MarshalOptions{
|
||||
Multiline: true,
|
||||
EmitASCII: true,
|
||||
}
|
||||
var unmarshalOptions = prototext.UnmarshalOptions{}
|
||||
|
||||
func (cli *CLI) NewAttestationCmd() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "attestation [command]",
|
||||
@@ -183,75 +153,31 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
|
||||
|
||||
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},
|
||||
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>`, SNP, VTPM, SNPvTPM),
|
||||
Args: cobra.ExactArgs(1),
|
||||
Use: "get",
|
||||
Short: "Retrieve attestation information from agent. Report data expected in hex enoded string of length 64 bytes.",
|
||||
Example: "get <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
|
||||
}
|
||||
|
||||
if err := cobra.OnlyValidArgs(cmd, args); err != nil {
|
||||
printError(cmd, "Bad attestation type: %v ❌ ", err)
|
||||
cmd.Println("Getting attestation")
|
||||
|
||||
reportData, err := hex.DecodeString(args[0])
|
||||
if err != nil {
|
||||
printError(cmd, "Error decoding report data: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
|
||||
attestationType := args[0]
|
||||
|
||||
attType := config.SNP
|
||||
switch attestationType {
|
||||
case SNP:
|
||||
cmd.Println("Fetching SEV-SNP attestation report")
|
||||
case VTPM:
|
||||
cmd.Println("Fetching vTPM report")
|
||||
attType = config.VTPM
|
||||
case SNPvTPM:
|
||||
cmd.Println("Fetching SEV-SNP and vTPM report")
|
||||
attType = config.SNPvTPM
|
||||
}
|
||||
|
||||
if (attType == config.VTPM || attType == config.SNPvTPM) && len(nonce) == 0 {
|
||||
msg := color.New(color.FgRed).Sprint("vTPM nonce must be defined for vTPM attestation ❌ ")
|
||||
if len(reportData) != agent.ReportDataSize {
|
||||
msg := color.New(color.FgRed).Sprintf("report data must be a hex encoded string of length %d bytes ❌ ", agent.ReportDataSize)
|
||||
cmd.Println(msg)
|
||||
return
|
||||
}
|
||||
|
||||
if (attType == config.SNP || attType == config.SNPvTPM) && len(teeNonce) == 0 {
|
||||
msg := color.New(color.FgRed).Sprint("TEE nonce must be defined for SEV-SNP attestation ❌ ")
|
||||
cmd.Println(msg)
|
||||
return
|
||||
}
|
||||
|
||||
var fixedReportData [quoteprovider.Nonce]byte
|
||||
if attType != config.VTPM {
|
||||
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 != config.SNP {
|
||||
if len(nonce) > 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
|
||||
}
|
||||
|
||||
copy(fixedVtpmNonceByte[:], nonce)
|
||||
}
|
||||
|
||||
filename := attestationFilePath
|
||||
if getTextProtoAttestation {
|
||||
if getJsonAttestation {
|
||||
filename = attestationJson
|
||||
}
|
||||
|
||||
@@ -261,7 +187,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
|
||||
return
|
||||
}
|
||||
|
||||
if err := cli.agentSDK.Attestation(cmd.Context(), fixedReportData, fixedVtpmNonceByte, int(attType), attestationFile); err != nil {
|
||||
if err := cli.agentSDK.Attestation(cmd.Context(), [agent.ReportDataSize]byte(reportData), attestationFile); err != nil {
|
||||
printError(cmd, "Failed to get attestation due to error: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
@@ -271,32 +197,16 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
|
||||
return
|
||||
}
|
||||
|
||||
if getTextProtoAttestation {
|
||||
if getJsonAttestation {
|
||||
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)
|
||||
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 ❌ ", ErrBadAttestation)
|
||||
}
|
||||
|
||||
result = []byte(marshalOptions.Format(&attvTPM))
|
||||
}
|
||||
|
||||
result, err = attesationToJSON(result)
|
||||
if err != nil {
|
||||
printError(cmd, "Error converting attestation to textproto: %v ❌ ", err)
|
||||
printError(cmd, "Error converting attestation to json: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -310,9 +220,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
|
||||
},
|
||||
}
|
||||
|
||||
cmd.Flags().BoolVarP(&getTextProtoAttestation, "textproto", "p", false, "Get attestation in textproto format")
|
||||
cmd.Flags().BytesHexVarP(&teeNonce, "tee", "e", []byte{}, "Define the nonce for the SNP attestation report (must be used with attestation type snp and snp-vtpm)")
|
||||
cmd.Flags().BytesHexVarP(&nonce, "vtpm", "t", []byte{}, "Define the nonce for the vTPM attestation report (must be used with attestation type vtpm and snp-vtpm)")
|
||||
cmd.Flags().BoolVarP(&getJsonAttestation, "json", "j", false, "Get attestation in json format")
|
||||
|
||||
return cmd
|
||||
}
|
||||
@@ -344,108 +252,52 @@ func isFileJSON(filename string) bool {
|
||||
|
||||
func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "validate",
|
||||
Short: "Validate and verify attestation information. You can choose from 3 modes: snp,vtpm and snp-vtpm.Default mode is snp.",
|
||||
Example: `Based on mode:
|
||||
validate <attestationreportfilepath> --report_data <reportdata> --product <product data> //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> --nonce <noncevalue> --format <formatvalue> --output <outputvalue>`,
|
||||
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) {
|
||||
cmd.Println("Checking attestation")
|
||||
|
||||
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")
|
||||
attestationFile = string(args[0])
|
||||
|
||||
if err := parseConfig(); err != nil {
|
||||
printError(cmd, "Error parsing config: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
if err := parseHashes(); err != nil {
|
||||
printError(cmd, "Error parsing hashes: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
if err := parseFiles(); err != nil {
|
||||
printError(cmd, "Error parsing files: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
// This format is the attestation report in AMD's specified ABI format, immediately
|
||||
// followed by the certificate table bytes.
|
||||
if len(attestation) < abi.ReportSize {
|
||||
msg := color.New(color.FgRed).Sprintf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes ❌ ", len(attestation), abi.ReportSize)
|
||||
cmd.Println(msg)
|
||||
return
|
||||
}
|
||||
if err := parseUints(); err != nil {
|
||||
printError(cmd, "Error parsing uints: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
cfg.Policy.Vmpl = wrapperspb.UInt32(0)
|
||||
|
||||
if err := validateInput(); err != nil {
|
||||
printError(cmd, "Error validating input: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 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-SNP mode: %v", err)
|
||||
}
|
||||
if err := cmd.MarkFlagRequired("product"); err != nil {
|
||||
return fmt.Errorf("failed to mark flag as required: %v ❌ ", err)
|
||||
}
|
||||
case "snp-vtpm":
|
||||
if err := cmd.MarkFlagRequired("nonce"); err != nil {
|
||||
return fmt.Errorf("failed to mark 'nonce' as required for vTPM mode: %v", err)
|
||||
}
|
||||
if err := cmd.MarkFlagRequired("format"); err != nil {
|
||||
return fmt.Errorf("failed to mark 'format' as required for vTPM mode: %v", err)
|
||||
}
|
||||
if err := cmd.MarkFlagRequired("output"); err != nil {
|
||||
return fmt.Errorf("failed to mark 'output' as required for vTPM mode: %v", err)
|
||||
}
|
||||
|
||||
case "vtpm":
|
||||
if err := cmd.MarkFlagRequired("nonce"); err != nil {
|
||||
return fmt.Errorf("failed to mark 'nonce' as required for vTPM mode: %v", err)
|
||||
}
|
||||
if err := cmd.MarkFlagRequired("format"); err != nil {
|
||||
return fmt.Errorf("failed to mark 'format' as required for vTPM mode: %v", err)
|
||||
}
|
||||
if err := cmd.MarkFlagRequired("output"); err != nil {
|
||||
return fmt.Errorf("failed to mark 'output' as required for vTPM mode: %v", err)
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("unknown mode: %s", mode)
|
||||
if err := quoteprovider.VerifyAndValidate(attestation, &cfg); err != nil {
|
||||
printError(cmd, "Attestation validation and verification failed with error: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
return nil
|
||||
cmd.Println("Attestation validation and verification is successful!")
|
||||
},
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
mode, _ := cmd.Flags().GetString("mode")
|
||||
switch mode {
|
||||
case "snp":
|
||||
return sevsnpverify(cmd, args)
|
||||
case "snp-vtpm":
|
||||
return vtpmSevSnpverify(args)
|
||||
case "vtpm":
|
||||
return vtpmverify(args)
|
||||
default:
|
||||
return fmt.Errorf("unknown mode: %s", mode)
|
||||
}
|
||||
},
|
||||
SilenceUsage: true,
|
||||
SilenceErrors: true,
|
||||
}
|
||||
cmd.Flags().StringVar(
|
||||
&mode,
|
||||
"mode",
|
||||
"snp", // default mode
|
||||
"The attestation validation mode. Example: sevsnp",
|
||||
)
|
||||
|
||||
// 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",
|
||||
)
|
||||
|
||||
cmd.Flags().BytesHexVar(
|
||||
&teeNonce,
|
||||
"tee-nonce",
|
||||
[]byte{},
|
||||
"hex encoded teenonce for hardware attestation, can be empty",
|
||||
)
|
||||
|
||||
// SEV-SNP FLAGS
|
||||
cmd.Flags().StringVar(
|
||||
&cfgString,
|
||||
"config",
|
||||
@@ -621,217 +473,19 @@ 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 {
|
||||
printError(cmd, "Failed to mark flag as required: %v ❌ ", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := cmd.MarkFlagRequired("product"); err != nil {
|
||||
printError(cmd, "Failed to mark flag as required: %v ❌ ", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
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]
|
||||
|
||||
return cli.measurement.Run(inputFile)
|
||||
},
|
||||
}
|
||||
|
||||
return igvmmeasureCmd
|
||||
}
|
||||
|
||||
func sevsnpverify(cmd *cobra.Command, args []string) error {
|
||||
cmd.Println("Checking attestation")
|
||||
|
||||
attestationFile = string(args[0])
|
||||
|
||||
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 := parseFiles(); err != nil {
|
||||
return fmt.Errorf("error parsing files: %v ❌ ", 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 {
|
||||
return fmt.Errorf("attestation too small: got 0x%x bytes, need at least 0x%x bytes", len(attestation), abi.ReportSize)
|
||||
}
|
||||
if err := parseUints(); err != nil {
|
||||
return fmt.Errorf("error parsing uints: %v ❌ ", err)
|
||||
}
|
||||
cfg.Policy.Vmpl = wrapperspb.UInt32(0)
|
||||
|
||||
if err := validateInput(); err != nil {
|
||||
return fmt.Errorf("error validating input: %v ❌ ", err)
|
||||
}
|
||||
|
||||
attestationPB, err := abi.ReportCertsToProto(attestation)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to convert attestation bytes to struct %v ❌ ", err)
|
||||
}
|
||||
|
||||
if err := quoteprovider.VerifyAndValidate(attestationPB, &cfg); 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 vtpmSevSnpverify(args []string) error {
|
||||
tpmAttestationFile = string(args[0])
|
||||
input, err := openInputFile()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if closer, ok := input.(*os.File); ok {
|
||||
defer closer.Close()
|
||||
}
|
||||
attestationBytes, err := io.ReadAll(input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
attestation := &tpmAttest.Attestation{}
|
||||
|
||||
if format == FormatBinaryPB {
|
||||
err = proto.Unmarshal(attestationBytes, attestation)
|
||||
} else if format == FormatTextProto {
|
||||
err = unmarshalOptions.Unmarshal(attestationBytes, attestation)
|
||||
} else {
|
||||
return fmt.Errorf("format should be either binarypb or textproto")
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("fail to unmarshal attestation report: %v", err)
|
||||
}
|
||||
|
||||
pub, err := tpm2.DecodePublic(attestation.GetAkPub())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cryptoPub, err := pub.Key()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var validateOpts interface{}
|
||||
switch attestation.GetTeeAttestation().(type) {
|
||||
case *tpmAttest.Attestation_SevSnpAttestation:
|
||||
if len(teeNonce) != 0 {
|
||||
validateOpts = &server.VerifySnpOpts{
|
||||
Validation: server.SevSnpDefaultValidateOpts(teeNonce),
|
||||
Verification: &sevVerify.Options{},
|
||||
}
|
||||
} else {
|
||||
validateOpts = &server.VerifySnpOpts{
|
||||
Validation: server.SevSnpDefaultValidateOpts(nonce),
|
||||
Verification: &sevVerify.Options{},
|
||||
}
|
||||
}
|
||||
default:
|
||||
validateOpts = nil
|
||||
}
|
||||
|
||||
ms, err := server.VerifyAttestation(attestation, server.VerifyOpts{Nonce: nonce, TrustedAKs: []crypto.PublicKey{cryptoPub}, TEEOpts: validateOpts})
|
||||
if err != nil {
|
||||
return fmt.Errorf("verifying attestation: %w", err)
|
||||
}
|
||||
out, err := marshalOptions.Marshal(ms)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
output, err := createOutputFile()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if closer, ok := output.(*os.File); ok {
|
||||
defer closer.Close()
|
||||
}
|
||||
if _, err := output.Write(out); err != nil {
|
||||
return fmt.Errorf("failed to write verified attestation report: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func vtpmverify(args []string) error {
|
||||
tpmAttestationFile = string(args[0])
|
||||
input, err := openInputFile()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if closer, ok := input.(*os.File); ok {
|
||||
defer closer.Close()
|
||||
}
|
||||
attestationBytes, err := io.ReadAll(input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
attestation := &tpmAttest.Attestation{}
|
||||
|
||||
if format == FormatBinaryPB {
|
||||
err = proto.Unmarshal(attestationBytes, attestation)
|
||||
} else if format == FormatTextProto {
|
||||
err = unmarshalOptions.Unmarshal(attestationBytes, attestation)
|
||||
} else {
|
||||
return fmt.Errorf("format should be either binarypb or textproto")
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("fail to unmarshal attestation report: %v", err)
|
||||
}
|
||||
|
||||
pub, err := tpm2.DecodePublic(attestation.GetAkPub())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cryptoPub, err := pub.Key()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
ms, err := server.VerifyAttestation(attestation, server.VerifyOpts{Nonce: nonce, TrustedAKs: []crypto.PublicKey{cryptoPub}, TEEOpts: nil})
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
out, err := marshalOptions.Marshal(ms)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
output, err := createOutputFile()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if closer, ok := output.(*os.File); ok {
|
||||
defer closer.Close()
|
||||
}
|
||||
if _, err := output.Write(out); err != nil {
|
||||
return fmt.Errorf("failed to write verified attestation report: %v", err)
|
||||
}
|
||||
|
||||
return 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
|
||||
/* {
|
||||
|
||||
+35
-211
@@ -14,12 +14,10 @@ import (
|
||||
"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/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
|
||||
)
|
||||
|
||||
@@ -34,8 +32,10 @@ func TestNewAttestationCmd(t *testing.T) {
|
||||
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.SetOutput(&buf)
|
||||
|
||||
reportData := bytes.Repeat([]byte{0x01}, agent.ReportDataSize)
|
||||
mockSDK.On("Attestation", mock.Anything, [agent.ReportDataSize]byte(reportData), mock.Anything).Return(nil)
|
||||
|
||||
cmd.SetArgs([]string{hex.EncodeToString(reportData)})
|
||||
err := cmd.Execute()
|
||||
@@ -46,10 +46,6 @@ func TestNewAttestationCmd(t *testing.T) {
|
||||
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))
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
args []string
|
||||
@@ -59,85 +55,57 @@ func TestNewGetAttestationCmd(t *testing.T) {
|
||||
expectedOut string
|
||||
}{
|
||||
{
|
||||
name: "successful SNP attestation retrieval",
|
||||
args: []string{"snp", "--tee", teeNonce},
|
||||
name: "successful attestation retrieval",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
|
||||
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", "-t", 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",
|
||||
name: "invalid report data (decoding error)",
|
||||
args: []string{"invalid"},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("error"),
|
||||
expectedErr: "Bad attestation type: invalid argument ",
|
||||
expectedErr: "Error decoding report data",
|
||||
},
|
||||
{
|
||||
name: "invalid report data size",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, 32))},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("error"),
|
||||
expectedErr: "report data must be a hex encoded string of length 64 bytes",
|
||||
},
|
||||
{
|
||||
name: "invalid report data hex",
|
||||
args: []string{"invalid"},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("error"),
|
||||
expectedErr: "Error decoding report data",
|
||||
},
|
||||
{
|
||||
name: "failed to get attestation",
|
||||
args: []string{"snp", "-e", teeNonce},
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("error"),
|
||||
expectedErr: "Failed to get attestation due to error",
|
||||
},
|
||||
{
|
||||
name: "Textproto report error",
|
||||
args: []string{"snp", "-e", teeNonce, "--textproto"},
|
||||
name: "JSON report error",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), "--json"},
|
||||
mockResponse: []byte("mock attestation"),
|
||||
mockError: nil,
|
||||
expectedErr: "Error converting attestation to textproto",
|
||||
expectedErr: "Error converting attestation to json",
|
||||
},
|
||||
{
|
||||
name: "successful Textproto report",
|
||||
args: []string{"snp", "-e", teeNonce, "--textproto"},
|
||||
name: "successful JSON report",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), "--json"},
|
||||
mockResponse: validattestation,
|
||||
mockError: nil,
|
||||
expectedOut: "Attestation result retrieved and saved successfully!",
|
||||
},
|
||||
{
|
||||
name: "connection error",
|
||||
args: []string{"snp", "-e", teeNonce},
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("failed to connect to agent"),
|
||||
expectedErr: "Failed to connect to agent",
|
||||
@@ -157,10 +125,10 @@ func TestNewGetAttestationCmd(t *testing.T) {
|
||||
}
|
||||
cmd := cli.NewGetAttestationCmd()
|
||||
var buf bytes.Buffer
|
||||
cmd.SetOut(&buf)
|
||||
cmd.SetOutput(&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)
|
||||
mockSDK.On("Attestation", mock.Anything, [agent.ReportDataSize]byte(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
|
||||
_, err := args.Get(2).(*os.File).Write(tc.mockResponse)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
|
||||
@@ -177,12 +145,12 @@ func TestNewGetAttestationCmd(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewValidateAttestationValidationCmdDefaults(t *testing.T) {
|
||||
func TestNewValidateAttestationValidationCmd(t *testing.T) {
|
||||
cli := &CLI{}
|
||||
cmd := cli.NewValidateAttestationValidationCmd()
|
||||
|
||||
assert.Equal(t, "validate", cmd.Use)
|
||||
assert.Equal(t, "Validate and verify attestation information. You can choose from 3 modes: snp,vtpm and snp-vtpm.Default mode is snp.", cmd.Short)
|
||||
assert.Equal(t, "Validate and verify attestation information. The report is provided as a file path.", 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())
|
||||
@@ -194,150 +162,6 @@ func TestNewValidateAttestationValidationCmdDefaults(t *testing.T) {
|
||||
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) error {
|
||||
args := m.Called(igvmBinaryPath)
|
||||
return 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()
|
||||
|
||||
+4
-4
@@ -12,7 +12,7 @@ import (
|
||||
"github.com/google/go-sev-guest/proto/check"
|
||||
"github.com/google/go-sev-guest/verify/trust"
|
||||
"github.com/spf13/cobra"
|
||||
config "github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/clients/grpc"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -27,14 +27,14 @@ func (cli *CLI) NewCABundleCmd(fileSavePath string) *cobra.Command {
|
||||
Example: "ca-bundle <path_to_platform_info_json>",
|
||||
Args: cobra.ExactArgs(1),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
attestationConfiguration := config.Config{SnpCheck: &check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}, PcrConfig: &config.PcrConfig{}}
|
||||
err := config.ReadAttestationPolicy(args[0], &attestationConfiguration)
|
||||
attestationConfiguration := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
|
||||
err := grpc.ReadAttestationPolicy(args[0], &attestationConfiguration)
|
||||
if err != nil {
|
||||
printError(cmd, "Error while reading manifest: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
|
||||
product := attestationConfiguration.SnpCheck.RootOfTrust.ProductLine
|
||||
product := attestationConfiguration.RootOfTrust.ProductLine
|
||||
|
||||
getter := trust.DefaultHTTPSGetter()
|
||||
caURL := kds.ProductCertChainURL(abi.VcekReportSigner, product)
|
||||
|
||||
@@ -131,7 +131,7 @@ func TestManifestChecksum(t *testing.T) {
|
||||
"name": "Example Computation",
|
||||
"description": "This is an example computation"
|
||||
}`,
|
||||
expectedSum: "a99683e4d22ba54cefa51aa49fb2e97a92b828c088395992ddff16a6236f3299",
|
||||
expectedSum: "868825367c32c4b6d621d5d95e2890f233d8554df2348ab743aac2663a936f08",
|
||||
},
|
||||
{
|
||||
name: "Invalid JSON",
|
||||
|
||||
-130
@@ -1,130 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package cli
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
"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"
|
||||
logLevel = "log-level"
|
||||
)
|
||||
|
||||
var (
|
||||
agentCVMServerUrl string
|
||||
agentCVMServerCA string
|
||||
agentCVMClientKey string
|
||||
agentCVMClientCrt string
|
||||
agentLogLevel string
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
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.SvmId, 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(&agentLogLevel, logLevel, "", "Agent Log level")
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
func (c *CLI) NewRemoveVMCmd() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "remove-vm",
|
||||
Short: "Remove a virtual machine",
|
||||
Example: `remove-vm <svm_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{SvmId: 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
|
||||
}
|
||||
+8
-30
@@ -6,36 +6,28 @@ import (
|
||||
"context"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/igvmmeasure"
|
||||
"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"
|
||||
)
|
||||
|
||||
var Verbose bool
|
||||
|
||||
type CLI struct {
|
||||
agentSDK sdk.SDK
|
||||
agentConfig grpc.AgentClientConfig
|
||||
managerConfig grpc.ManagerClientConfig
|
||||
client grpc.Client
|
||||
managerClient manager.ManagerServiceClient
|
||||
connectErr error
|
||||
measurement igvmmeasure.MeasurementProvider
|
||||
agentSDK sdk.SDK
|
||||
config grpc.AgentClientConfig
|
||||
client grpc.Client
|
||||
connectErr error
|
||||
}
|
||||
|
||||
func New(agentConfig grpc.AgentClientConfig, managerConfig grpc.ManagerClientConfig, measurement igvmmeasure.MeasurementProvider) *CLI {
|
||||
func New(config grpc.AgentClientConfig) *CLI {
|
||||
return &CLI{
|
||||
agentConfig: agentConfig,
|
||||
managerConfig: managerConfig,
|
||||
measurement: measurement,
|
||||
config: config,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *CLI) InitializeAgentSDK(cmd *cobra.Command) error {
|
||||
agentGRPCClient, agentClient, err := agent.NewAgentClient(context.Background(), c.agentConfig)
|
||||
func (c *CLI) InitializeSDK(cmd *cobra.Command) error {
|
||||
agentGRPCClient, agentClient, err := agent.NewAgentClient(context.Background(), c.config)
|
||||
if err != nil {
|
||||
c.connectErr = err
|
||||
return err
|
||||
@@ -47,20 +39,6 @@ func (c *CLI) InitializeAgentSDK(cmd *cobra.Command) error {
|
||||
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()
|
||||
}
|
||||
|
||||
+174
-99
@@ -3,165 +3,148 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"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/cenkalti/backoff/v4"
|
||||
"github.com/google/go-sev-guest/abi"
|
||||
"github.com/google/go-sev-guest/client"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"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/auth"
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
agentlogger "github.com/ultravioletrs/cocos/internal/logger"
|
||||
"github.com/ultravioletrs/cocos/internal/server"
|
||||
grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
|
||||
ackvsock "github.com/ultravioletrs/cocos/internal/vsock"
|
||||
managerevents "github.com/ultravioletrs/cocos/manager/events"
|
||||
"github.com/ultravioletrs/cocos/manager/qemu"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
|
||||
pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
|
||||
cvmsgrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/cvm"
|
||||
"golang.org/x/crypto/sha3"
|
||||
"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"
|
||||
retryInterval = 5 * time.Second
|
||||
)
|
||||
|
||||
type config struct {
|
||||
LogLevel string `env:"AGENT_LOG_LEVEL" envDefault:"debug"`
|
||||
Vmpl int `env:"AGENT_VMPL" envDefault:"2"`
|
||||
}
|
||||
|
||||
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())
|
||||
}
|
||||
|
||||
conn, err := dialVsock()
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
ackConn := ackvsock.NewAckWriter(conn)
|
||||
|
||||
var exitCode int
|
||||
defer mglog.ExitWithError(&exitCode)
|
||||
|
||||
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(ackConn, &slog.HandlerOptions{Level: level}, cfg.ID)
|
||||
logger := slog.New(handler)
|
||||
|
||||
eventSvc, err := events.New(svcName, eventsLogsQueue)
|
||||
eventSvc, err := events.New(svcName, cfg.ID, ackConn)
|
||||
if err != nil {
|
||||
logger.Error(fmt.Sprintf("failed to create events service %s", err.Error()))
|
||||
exitCode = 1
|
||||
return
|
||||
}
|
||||
|
||||
var qp client.LeveledQuoteProvider
|
||||
|
||||
if !sevGuesDeviceExists() {
|
||||
logger.Info("SEV-SNP device not found")
|
||||
qpMock := new(mocks.LeveledQuoteProvider)
|
||||
qpMock.On("GetRawQuoteAtLevel", mock.Anything, mock.Anything).Return([]uint8{}, errors.New("SEV-SNP device not found"))
|
||||
qp = qpMock
|
||||
} else {
|
||||
qp, err = quoteprovider.GetLeveledQuoteProvider()
|
||||
if err != nil {
|
||||
logger.Error(fmt.Sprintf("failed to create quote provider %s", err.Error()))
|
||||
exitCode = 1
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
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))
|
||||
qp, err := quoteprovider.GetQuoteProvider()
|
||||
if err != nil {
|
||||
logger.Error(fmt.Sprintf("failed to create quote provider %s", err.Error()))
|
||||
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) (cvms.Service_ProcessClient, error) {
|
||||
_, newClient, err := cvmsgrpc.NewCVMClient(cvmGrpcConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Don't defer close here as we want to keep the connection open
|
||||
|
||||
return newClient.Process(ctx)
|
||||
}
|
||||
|
||||
pc, err := cvmsClient.Process(ctx)
|
||||
if err != nil {
|
||||
if err := verifyManifest(cfg, qp); 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
|
||||
setDefaultValues(&cfg)
|
||||
|
||||
svc := newService(ctx, logger, eventSvc, cfg, qp)
|
||||
|
||||
agentGrpcServerConfig := server.AgentConfig{
|
||||
ServerConfig: server.ServerConfig{
|
||||
BaseConfig: server.BaseConfig{
|
||||
Host: cfg.AgentConfig.Host,
|
||||
Port: cfg.AgentConfig.Port,
|
||||
CertFile: cfg.AgentConfig.CertFile,
|
||||
KeyFile: cfg.AgentConfig.KeyFile,
|
||||
ServerCAFile: cfg.AgentConfig.ServerCAFile,
|
||||
ClientCAFile: cfg.AgentConfig.ClientCAFile,
|
||||
},
|
||||
},
|
||||
AttestedTLS: cfg.AgentConfig.AttestedTls,
|
||||
}
|
||||
|
||||
svc := newService(ctx, logger, eventSvc, qp, cfg.Vmpl)
|
||||
|
||||
if err := os.MkdirAll(storageDir, 0o755); err != nil {
|
||||
logger.Error(fmt.Sprintf("failed to create storage directory: %s", err))
|
||||
exitCode = 1
|
||||
return
|
||||
registerAgentServiceServer := func(srv *grpc.Server) {
|
||||
reflection.Register(srv)
|
||||
agent.RegisterAgentServiceServer(srv, agentgrpc.NewServer(svc))
|
||||
}
|
||||
|
||||
mc, err := cvmsapi.NewClient(pc, svc, eventsLogsQueue, logger, server.NewServer(logger, svc), storageDir, reconnectFn)
|
||||
authSvc, err := auth.New(cfg)
|
||||
if err != nil {
|
||||
logger.Error(err.Error())
|
||||
logger.Error(fmt.Sprintf("failed to create auth service %s", err.Error()))
|
||||
exitCode = 1
|
||||
return
|
||||
}
|
||||
|
||||
gs := grpcserver.New(ctx, cancel, svcName, agentGrpcServerConfig, registerAgentServiceServer, logger, qp, authSvc)
|
||||
|
||||
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()
|
||||
for {
|
||||
if _, err := io.Copy(io.Discard, conn); err != nil {
|
||||
log.Printf("vsock connection lost: %v, reconnecting...", err)
|
||||
conn.Close()
|
||||
conn, err = dialVsock()
|
||||
if err != nil {
|
||||
log.Fatal("failed to reconnect: ", err)
|
||||
}
|
||||
}
|
||||
time.Sleep(retryInterval)
|
||||
}
|
||||
})
|
||||
|
||||
g.Go(func() error {
|
||||
return mc.Process(ctx, cancel)
|
||||
return gs.Start()
|
||||
})
|
||||
|
||||
g.Go(func() error {
|
||||
return server.StopHandler(ctx, cancel, logger, svcName, gs)
|
||||
})
|
||||
|
||||
if err := g.Wait(); err != nil {
|
||||
@@ -169,8 +152,8 @@ func main() {
|
||||
}
|
||||
}
|
||||
|
||||
func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Service, qp client.LeveledQuoteProvider, vmpl int) agent.Service {
|
||||
svc := agent.New(ctx, logger, eventSvc, qp, vmpl)
|
||||
func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp agent.Computation, qp client.QuoteProvider) agent.Service {
|
||||
svc := agent.New(ctx, logger, eventSvc, cmp, qp)
|
||||
|
||||
svc = api.LoggingMiddleware(svc, logger)
|
||||
counter, latency := prometheus.MakeMetrics(svcName, "api")
|
||||
@@ -179,11 +162,103 @@ func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Servic
|
||||
return svc
|
||||
}
|
||||
|
||||
func sevGuesDeviceExists() bool {
|
||||
d, err := client.OpenDevice()
|
||||
func readConfig() (agent.Computation, error) {
|
||||
l, err := vsock.Listen(qemu.VsockConfigPort, nil)
|
||||
if err != nil {
|
||||
return false
|
||||
return agent.Computation{}, err
|
||||
}
|
||||
d.Close()
|
||||
return true
|
||||
defer l.Close()
|
||||
|
||||
conn, err := l.Accept()
|
||||
if err != nil {
|
||||
return agent.Computation{}, err
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
var buffer []byte
|
||||
for {
|
||||
chunk := make([]byte, 1024)
|
||||
n, err := conn.Read(chunk)
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
return agent.Computation{}, err
|
||||
}
|
||||
buffer = append(buffer, chunk[:n]...)
|
||||
}
|
||||
|
||||
ac := agent.Computation{
|
||||
AgentConfig: agent.AgentConfig{},
|
||||
}
|
||||
if err := json.Unmarshal(buffer, &ac); err != nil {
|
||||
return agent.Computation{}, err
|
||||
}
|
||||
return ac, nil
|
||||
}
|
||||
|
||||
func setDefaultValues(cfg *agent.Computation) {
|
||||
if cfg.AgentConfig.LogLevel == "" {
|
||||
cfg.AgentConfig.LogLevel = "info"
|
||||
}
|
||||
if cfg.AgentConfig.Port == "" {
|
||||
cfg.AgentConfig.Port = defSvcGRPCPort
|
||||
}
|
||||
}
|
||||
|
||||
func isTEE() bool {
|
||||
_, err := os.Stat("/dev/sev-guest")
|
||||
return !os.IsNotExist(err)
|
||||
}
|
||||
|
||||
func dialVsock() (*vsock.Conn, error) {
|
||||
var conn *vsock.Conn
|
||||
var err error
|
||||
|
||||
err = backoff.Retry(func() error {
|
||||
conn, err = vsock.Dial(vsock.Host, managerevents.ManagerVsockPort, nil)
|
||||
if err == nil {
|
||||
log.Println("vsock connection established")
|
||||
return nil
|
||||
}
|
||||
log.Printf("vsock connection failed, retrying in %s... Error: %v", retryInterval, err)
|
||||
return err
|
||||
}, backoff.NewExponentialBackOff())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func verifyManifest(cfg agent.Computation, qp client.QuoteProvider) error {
|
||||
if !isTEE() {
|
||||
return nil
|
||||
}
|
||||
|
||||
ar, err := qp.GetRawQuote(sha3.Sum512([]byte(cfg.ID)))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
arProto, err := abi.ReportCertsToProto(ar[:abi.ReportSize])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cfgBytes, err := json.Marshal(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
mcHash := sha3.Sum256(cfgBytes)
|
||||
|
||||
if arProto.Report.HostData == nil {
|
||||
return fmt.Errorf("manifest verification failed: HostData is nil")
|
||||
}
|
||||
if !bytes.Equal(arProto.Report.HostData, mcHash[:]) {
|
||||
return fmt.Errorf("manifest verification failed")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/agent/events/mocks"
|
||||
qpmocks "github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
|
||||
)
|
||||
|
||||
func TestSetDefaultValues(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input agent.Computation
|
||||
expected agent.Computation
|
||||
}{
|
||||
{
|
||||
name: "Empty config",
|
||||
input: agent.Computation{
|
||||
AgentConfig: agent.AgentConfig{},
|
||||
},
|
||||
expected: agent.Computation{
|
||||
AgentConfig: agent.AgentConfig{
|
||||
LogLevel: "info",
|
||||
Port: "7002",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Partial config",
|
||||
input: agent.Computation{
|
||||
AgentConfig: agent.AgentConfig{
|
||||
LogLevel: "debug",
|
||||
},
|
||||
},
|
||||
expected: agent.Computation{
|
||||
AgentConfig: agent.AgentConfig{
|
||||
LogLevel: "debug",
|
||||
Port: "7002",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
setDefaultValues(&tt.input)
|
||||
assert.Equal(t, tt.expected, tt.input)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewService(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := new(mocks.Service)
|
||||
eventSvc.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
cmp := agent.Computation{
|
||||
ID: "test-computation",
|
||||
AgentConfig: agent.AgentConfig{
|
||||
LogLevel: "info",
|
||||
Port: "7002",
|
||||
},
|
||||
}
|
||||
qp := new(qpmocks.QuoteProvider)
|
||||
|
||||
svc := newService(ctx, logger, eventSvc, cmp, qp)
|
||||
|
||||
assert.NotNil(t, svc)
|
||||
}
|
||||
|
||||
func TestVerifyManifest(t *testing.T) {
|
||||
cfg := agent.Computation{
|
||||
ID: "test-computation",
|
||||
AgentConfig: agent.AgentConfig{
|
||||
LogLevel: "info",
|
||||
Port: "7002",
|
||||
},
|
||||
}
|
||||
|
||||
mockQP := new(qpmocks.QuoteProvider)
|
||||
mockQP.On("GetRawQuote", mock.Anything).Return([]byte{}, nil)
|
||||
|
||||
err := verifyManifest(cfg, mockQP)
|
||||
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
+8
-28
@@ -14,23 +14,20 @@ import (
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/spf13/pflag"
|
||||
"github.com/ultravioletrs/cocos/cli"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/igvmmeasure"
|
||||
"github.com/ultravioletrs/cocos/pkg/clients/grpc"
|
||||
cmd "github.com/virtee/sev-snp-measure-go/sevsnpmeasure/cmd"
|
||||
)
|
||||
|
||||
const (
|
||||
svcName = "cli"
|
||||
envPrefixAgentGRPC = "AGENT_GRPC_"
|
||||
envPrefixManagerGRPC = "MANAGER_GRPC_"
|
||||
completion = "completion"
|
||||
filePermision = 0o755
|
||||
cocosDirectory = ".cocos"
|
||||
svcName = "cli"
|
||||
envPrefixAgentGRPC = "AGENT_GRPC_"
|
||||
completion = "completion"
|
||||
filePermision = 0o755
|
||||
cocosDirectory = ".cocos"
|
||||
)
|
||||
|
||||
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() {
|
||||
@@ -101,23 +98,9 @@ func main() {
|
||||
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
|
||||
}
|
||||
cliSVC := cli.New(agentGRPCConfig)
|
||||
|
||||
measurement, err := igvmmeasure.NewIgvmMeasurement(cfg.IgvmBinaryPath, os.Stderr, os.Stdout)
|
||||
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 {
|
||||
if err := cliSVC.InitializeSDK(rootCmd); err == nil {
|
||||
defer cliSVC.Close()
|
||||
}
|
||||
|
||||
@@ -136,8 +119,6 @@ func main() {
|
||||
rootCmd.AddCommand(attestationPolicyCmd)
|
||||
rootCmd.AddCommand(keysCmd)
|
||||
rootCmd.AddCommand(cliSVC.NewCABundleCmd(directoryCachePath))
|
||||
rootCmd.AddCommand(cliSVC.NewCreateVMCmd())
|
||||
rootCmd.AddCommand(cliSVC.NewRemoveVMCmd())
|
||||
|
||||
// Attestation commands
|
||||
attestationCmd.AddCommand(cliSVC.NewGetAttestationCmd())
|
||||
@@ -145,7 +126,6 @@ func main() {
|
||||
|
||||
// measure.
|
||||
rootCmd.AddCommand(cmd.NewRootCmd())
|
||||
rootCmd.AddCommand(cliSVC.NewMeasureCmd(cfg.IgvmBinaryPath))
|
||||
|
||||
// Flags
|
||||
keysCmd.PersistentFlags().StringVarP(
|
||||
|
||||
+47
-15
@@ -10,24 +10,25 @@ import (
|
||||
"log/slog"
|
||||
"net/url"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"syscall"
|
||||
|
||||
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/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/events"
|
||||
"github.com/ultravioletrs/cocos/manager/qemu"
|
||||
"github.com/ultravioletrs/cocos/manager/tracing"
|
||||
pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
|
||||
managergrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/manager"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/reflection"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -91,33 +92,64 @@ func main() {
|
||||
args := qemuCfg.ConstructQemuArgs()
|
||||
logger.Info(strings.Join(args, " "))
|
||||
|
||||
managerGRPCConfig := server.ServerConfig{}
|
||||
managerGRPCConfig := pkggrpc.ManagerClientConfig{}
|
||||
if err := env.ParseWithOptions(&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(logger, tracer, qemuCfg, cfg.AttestationPolicyBinary, cfg.EosVersion)
|
||||
managerGRPCClient, managerClient, err := managergrpc.NewManagerClient(managerGRPCConfig)
|
||||
if err != nil {
|
||||
logger.Error(err.Error())
|
||||
exitCode = 1
|
||||
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 *manager.ClientStreamMessage, clientBufferSize)
|
||||
svc, err := newService(logger, tracer, qemuCfg, eventsChan, cfg.AttestationPolicyBinary, cfg.EosVersion)
|
||||
if err != nil {
|
||||
logger.Error(err.Error())
|
||||
exitCode = 1
|
||||
return
|
||||
}
|
||||
|
||||
gs := grpcserver.New(ctx, cancel, svcName, managerGRPCConfig, registerManagerServiceServer, logger, nil, nil)
|
||||
eventsSvc, err := events.New(logger, svc.ReportBrokenConnection, eventsChan)
|
||||
if err != nil {
|
||||
logger.Error(err.Error())
|
||||
exitCode = 1
|
||||
return
|
||||
}
|
||||
|
||||
go eventsSvc.Listen(ctx)
|
||||
|
||||
mc := managerapi.NewClient(pc, svc, eventsChan, logger)
|
||||
|
||||
g.Go(func() error {
|
||||
return gs.Start()
|
||||
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()
|
||||
}
|
||||
})
|
||||
|
||||
g.Go(func() error {
|
||||
return server.StopHandler(ctx, cancel, logger, svcName, gs)
|
||||
return mc.Process(ctx, cancel)
|
||||
})
|
||||
|
||||
if err := g.Wait(); err != nil {
|
||||
@@ -125,8 +157,8 @@ func main() {
|
||||
}
|
||||
}
|
||||
|
||||
func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, attestationPolicyPath string, eosVersion string) (manager.Service, error) {
|
||||
svc, err := manager.New(qemuCfg, attestationPolicyPath, logger, qemu.NewVM, eosVersion)
|
||||
func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, eventsChan chan *manager.ClientStreamMessage, attestationPolicyPath string, eosVersion string) (manager.Service, error) {
|
||||
svc, err := manager.New(qemuCfg, attestationPolicyPath, logger, eventsChan, qemu.NewVM, eosVersion)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
+6
-7
@@ -10,8 +10,7 @@ MANAGER_ATTESTATION_POLICY_BINARY=../../build
|
||||
MANAGER_GRPC_CLIENT_CERT=
|
||||
MANAGER_GRPC_CLIENT_KEY=
|
||||
MANAGER_GRPC_SERVER_CA_CERTS=
|
||||
MANAGER_GRPC_PORT=6101
|
||||
MANAGER_GRPC_HOST=0.0.0.0
|
||||
MANAGER_GRPC_URL=localhost:7001
|
||||
MANAGER_GRPC_TIMEOUT=60s
|
||||
MANAGER_EOS_VERSION=""
|
||||
|
||||
@@ -22,13 +21,13 @@ MANAGER_QEMU_MAX_MEMORY=30G
|
||||
MANAGER_QEMU_OVMF_CODE_IF=pflash
|
||||
MANAGER_QEMU_OVMF_CODE_FORMAT=raw
|
||||
MANAGER_QEMU_OVMF_CODE_UNIT=0
|
||||
MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/edk2/x64/OVMF_CODE.fd
|
||||
MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/OVMF/x64/OVMF_CODE.fd
|
||||
MANAGER_QEMU_OVMF_VERSION=edk2-stable202408
|
||||
MANAGER_QEMU_OVMF_CODE_READONLY=on
|
||||
MANAGER_QEMU_OVMF_VARS_IF=pflash
|
||||
MANAGER_QEMU_OVMF_VARS_FORMAT=raw
|
||||
MANAGER_QEMU_OVMF_VARS_UNIT=1
|
||||
MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/edk2/x64/OVMF_VARS.fd
|
||||
MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/OVMF/x64/OVMF_VARS.fd
|
||||
MANAGER_QEMU_NETDEV_ID=vmnic
|
||||
MANAGER_QEMU_HOST_FWD_AGENT=7020
|
||||
MANAGER_QEMU_GUEST_FWD_AGENT=7002
|
||||
@@ -36,8 +35,8 @@ MANAGER_QEMU_VIRTIO_NET_PCI_DISABLE_LEGACY=on
|
||||
MANAGER_QEMU_VIRTIO_NET_PCI_IOMMU_PLATFORM=true
|
||||
MANAGER_QEMU_VIRTIO_NET_PCI_ADDR=0x2
|
||||
MANAGER_QEMU_VIRTIO_NET_PCI_ROMFILE=
|
||||
MANAGER_QEMU_DISK_IMG_KERNEL_FILE=/etc/cocos/bzImage
|
||||
MANAGER_QEMU_DISK_IMG_ROOTFS_FILE=/etc/cocos/rootfs.cpio.gz
|
||||
MANAGER_QEMU_DISK_IMG_KERNEL_FILE=/home/sammyk/Documents/cocos-ai/cmd/manager/img/bzImage
|
||||
MANAGER_QEMU_DISK_IMG_ROOTFS_FILE=/home/sammyk/Documents/cocos-ai/cmd/manager/img/rootfs.cpio.gz
|
||||
MANAGER_QEMU_SEV_ID=sev0
|
||||
MANAGER_QEMU_SEV_CBITPOS=51
|
||||
MANAGER_QEMU_SEV_REDUCED_PHYS_BITS=1
|
||||
@@ -49,13 +48,13 @@ MANAGER_QEMU_BIN_PATH=qemu-system-x86_64
|
||||
MANAGER_QEMU_USE_SUDO=true
|
||||
MANAGER_QEMU_ENABLE_SEV=false
|
||||
MANAGER_QEMU_ENABLE_SEV_SNP=false
|
||||
MANAGER_QEMU_IGVM_FILE=/etc/cocos/coconut-qemu.igvm
|
||||
MANAGER_QEMU_ENABLE_KVM=true
|
||||
MANAGER_QEMU_MACHINE=q35
|
||||
MANAGER_QEMU_CPU=EPYC
|
||||
MANAGER_QEMU_SMP_COUNT=4
|
||||
MANAGER_QEMU_SMP_MAXCPUS=16
|
||||
MANAGER_QEMU_MEM_ID=ram1
|
||||
MANAGER_QEMU_KERNEL_HASH=false
|
||||
MANAGER_QEMU_NO_GRAPHIC=true
|
||||
MANAGER_QEMU_MONITOR=pty
|
||||
MANAGER_QEMU_HOST_FWD_RANGE=6100-6200
|
||||
|
||||
@@ -1,31 +1,30 @@
|
||||
module github.com/ultravioletrs/cocos
|
||||
|
||||
go 1.23.2
|
||||
go 1.23.0
|
||||
|
||||
require (
|
||||
github.com/absmach/magistrala v0.15.1
|
||||
github.com/caarlos0/env/v11 v11.3.1
|
||||
github.com/caarlos0/env/v11 v11.2.2
|
||||
github.com/cenkalti/backoff/v4 v4.3.0
|
||||
github.com/fatih/color v1.18.0
|
||||
github.com/go-kit/kit v0.13.0
|
||||
github.com/gofrs/uuid v4.4.0+incompatible
|
||||
github.com/google/go-sev-guest v0.12.1
|
||||
github.com/google/go-tdx-guest v0.3.2-0.20241009005452-097ee70d0843 // indirect
|
||||
github.com/google/go-sev-guest v0.11.1
|
||||
github.com/mdlayher/vsock v1.2.1
|
||||
github.com/spf13/cobra v1.9.1
|
||||
github.com/spf13/pflag v1.0.6
|
||||
github.com/spf13/cobra v1.8.1
|
||||
github.com/spf13/pflag v1.0.5
|
||||
github.com/stretchr/testify v1.10.0
|
||||
github.com/virtee/sev-snp-measure-go v0.0.0-20240530153610-e6e8dc9b6877
|
||||
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.59.0
|
||||
go.opentelemetry.io/otel/trace v1.34.0
|
||||
golang.org/x/crypto v0.35.0
|
||||
golang.org/x/sync v0.11.0
|
||||
google.golang.org/grpc v1.70.0
|
||||
google.golang.org/protobuf v1.36.5
|
||||
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.57.0
|
||||
go.opentelemetry.io/otel/trace v1.32.0
|
||||
golang.org/x/crypto v0.30.0
|
||||
golang.org/x/sync v0.10.0
|
||||
google.golang.org/grpc v1.68.1
|
||||
google.golang.org/protobuf v1.35.2
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/Microsoft/go-winio v0.6.2 // indirect
|
||||
github.com/cenkalti/backoff/v4 v4.3.0 // indirect
|
||||
github.com/containerd/log v0.1.0 // indirect
|
||||
github.com/distribution/reference v0.6.0 // indirect
|
||||
github.com/docker/go-connections v0.5.0 // indirect
|
||||
@@ -33,12 +32,6 @@ require (
|
||||
github.com/felixge/httpsnoop v1.0.4 // indirect
|
||||
github.com/gofrs/uuid/v5 v5.3.0 // indirect
|
||||
github.com/gogo/protobuf v1.3.2 // indirect
|
||||
github.com/golang/protobuf v1.5.4 // indirect
|
||||
github.com/google/certificate-transparency-go v1.1.2 // indirect
|
||||
github.com/google/gce-tcb-verifier v0.2.3-0.20240905212129-12f728a62786 // indirect
|
||||
github.com/google/go-attestation v0.5.1 // indirect
|
||||
github.com/google/go-eventlog v0.0.2-0.20241003021507-01bb555f7cba // indirect
|
||||
github.com/google/go-tspi v0.3.0 // indirect
|
||||
github.com/mattn/go-colorable v0.1.13 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/moby/docker-image-spec v1.3.1 // indirect
|
||||
@@ -46,13 +39,11 @@ require (
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
|
||||
github.com/opencontainers/go-digest v1.0.0 // indirect
|
||||
github.com/opencontainers/image-spec v1.1.0 // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.57.0 // indirect
|
||||
go.opentelemetry.io/otel v1.34.0 // indirect
|
||||
go.opentelemetry.io/otel v1.32.0 // indirect
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.32.0 // indirect
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.32.0 // indirect
|
||||
go.opentelemetry.io/otel/sdk v1.32.0 // indirect
|
||||
golang.org/x/exp v0.0.0-20240909161429-701f63a606c0 // indirect
|
||||
gotest.tools/v3 v3.5.1 // indirect
|
||||
)
|
||||
|
||||
@@ -60,19 +51,19 @@ require (
|
||||
github.com/beorn7/perks v1.0.1 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
|
||||
github.com/docker/docker v28.0.1+incompatible
|
||||
github.com/docker/docker v27.4.0+incompatible
|
||||
github.com/go-kit/log v0.2.1 // indirect
|
||||
github.com/go-logfmt/logfmt v0.6.0 // indirect
|
||||
github.com/go-logr/logr v1.4.2 // indirect
|
||||
github.com/go-logr/stdr v1.2.2 // indirect
|
||||
github.com/google/go-configfs-tsm v0.3.3-0.20240919001351-b4b5b84fdcbc // indirect
|
||||
github.com/google/go-tpm v0.9.3
|
||||
github.com/google/go-tpm-tools v0.4.4
|
||||
github.com/golang/protobuf v1.5.4 // indirect
|
||||
github.com/google/go-configfs-tsm v0.2.2 // indirect
|
||||
github.com/google/logger v1.1.1
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.23.0 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
github.com/mdlayher/socket v0.4.1 // indirect
|
||||
github.com/pborman/uuid v1.2.1 // indirect
|
||||
github.com/pkg/errors v0.9.1 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
|
||||
github.com/prometheus/client_golang v1.20.5 // indirect
|
||||
@@ -80,18 +71,16 @@ require (
|
||||
github.com/prometheus/common v0.59.1 // indirect
|
||||
github.com/prometheus/procfs v0.15.1 // indirect
|
||||
github.com/stretchr/objx v0.5.2 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.34.0 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.32.0 // indirect
|
||||
go.opentelemetry.io/proto/otlp v1.3.1 // indirect
|
||||
go.uber.org/multierr v1.11.0 // indirect
|
||||
golang.org/x/net v0.34.0 // indirect
|
||||
golang.org/x/sys v0.30.0 // indirect
|
||||
golang.org/x/term v0.29.0
|
||||
golang.org/x/text v0.22.0 // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20250204164813-702378808489 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20250207221924-e9438ea467c6 // indirect
|
||||
golang.org/x/net v0.31.0 // indirect
|
||||
golang.org/x/sys v0.28.0 // indirect
|
||||
golang.org/x/term v0.27.0
|
||||
golang.org/x/text v0.21.0 // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20241104194629-dd2ea8efbc28 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20241104194629-dd2ea8efbc28 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
)
|
||||
|
||||
replace github.com/virtee/sev-snp-measure-go => github.com/sammyoina/sev-snp-measure-go v0.0.0-20241202151803-ef189f0ff825
|
||||
|
||||
replace github.com/google/go-tpm-tools => github.com/danko-miladinovic/go-tpm-tools v0.0.0-20250228160324-1ebcfd79567c
|
||||
replace github.com/virtee/sev-snp-measure-go => github.com/sammyoina/sev-snp-measure-go v0.0.0-20241107163739-38915ab517c7
|
||||
|
||||
@@ -1,89 +0,0 @@
|
||||
#### memory config
|
||||
MEMORY_SIZE=2048M
|
||||
MEMORY_SLOTS=5
|
||||
MAX_MEMORY=30G
|
||||
|
||||
|
||||
#### ovmf code config
|
||||
OVMF_CODE_IF=pflash
|
||||
OVMF_CODE_FORMAT=raw
|
||||
OVMF_CODE_UNIT=0
|
||||
OVMF_CODE_FILE=/usr/share/OVMF/OVMF_CODE.fd
|
||||
OVMF_CODE_READONLY=on
|
||||
OVMF_VERSION=
|
||||
|
||||
#### ovmf vars config
|
||||
OVMF_VARS_IF=pflash
|
||||
OVMF_VARS_FORMAT=raw
|
||||
OVMF_VARS_UNIT=1
|
||||
OVMF_VARS_FILE=/usr/share/OVMF/OVMF_VARS.fd
|
||||
|
||||
#### net dev config
|
||||
NET_DEV_ID=vmnic
|
||||
NET_DEV_HOST_FWD_AGENT=7020
|
||||
NET_DEV_GUEST_FWD_AGENT=7002
|
||||
|
||||
#### Virtio Net Pci Config
|
||||
VIRTIO_NET_PCI_DISABLE_LEGACY=on
|
||||
VIRTIO_NET_PCI_IOMMU_PLATFORM=true
|
||||
VIRTIO_NET_PCI_ADDR=0x2
|
||||
VIRTIO_NET_PCI_ROMFILE=
|
||||
|
||||
#### Disk image config
|
||||
DISK_IMG_KERNEL_FILE=
|
||||
DISK_IMG_ROOTFS_FILE=
|
||||
|
||||
KERNEL_COMMAND_LINE="quiet console=null"
|
||||
|
||||
#### Sev Config
|
||||
SEV_ID=sev0
|
||||
SEV_CBIT_POS=51
|
||||
SEV_REDUCED_PHYS_BITS=1
|
||||
SEV_HOST_DATA=
|
||||
|
||||
#### VSock Config
|
||||
VSOCK_ID=vhost-vsock-pci0
|
||||
VSOCK_GUEST_CID=3
|
||||
|
||||
BIN_PATH=qemu-system-x86_64
|
||||
USE_SUDO=false
|
||||
ENABLE_SEV=false
|
||||
ENABLE_SEV_SNP=false
|
||||
|
||||
ENABLE_KVM=true
|
||||
|
||||
MACHINE=q35
|
||||
CPU=EPYC
|
||||
SMP_COUNT=8
|
||||
SMP_MAXCPUS=64
|
||||
MEM_ID=ram1
|
||||
|
||||
KERNEL_HASH=false
|
||||
|
||||
NO_GRAPHIC=true
|
||||
MONITOR=pty
|
||||
HOST_FWD_RANGE=6100-6200
|
||||
CERTS_MOUNT=/etc/cocos/certs
|
||||
ENV_MOUNT=/etc/cocos/environment
|
||||
|
||||
COCOS_AGENT_VERSION=v0.3.1
|
||||
|
||||
#### Base image URL and names
|
||||
BASE_IMAGE_URL=https://cloud-images.ubuntu.com/noble/current/noble-server-cloudimg-amd64.img
|
||||
BASE_IMAGE=ubuntu-base.qcow2
|
||||
CUSTOM_IMAGE=ubuntu-custom.qcow2
|
||||
|
||||
#### Paths for OVMF firmware
|
||||
OVMF_CODE=/usr/share/ovmf/x64/OVMF_CODE.4m.fd
|
||||
OVMF_VARS=/usr/share/ovmf/x64/OVMF_VARS.4m.fd
|
||||
|
||||
#### VM parameters
|
||||
VM_NAME=cocos-vm
|
||||
RAM=16G
|
||||
DISK_SIZE=10G # Size for root filesystem
|
||||
QEMU_BINARY=qemu-system-x86_64
|
||||
|
||||
AGENT_GRPC_SERVER_CERT=/etc/cocos/certs/server.pem
|
||||
AGENT_GRPC_SERVER_KEY=/etc/cocos/certs/key.pem
|
||||
AGENT_GRPC_SERVER_CA_CERTS=/etc/cocos/ca.pem
|
||||
AGENT_GRPC_CLIENT_CA_CERTS=/etc/cocos/ca.pem
|
||||
@@ -1,114 +0,0 @@
|
||||
# Agent Cloud Init Setup
|
||||
|
||||
## Overview
|
||||
|
||||
The `hal/cloud` directory contains essential files required for setting up a virtual machine (VM) with cloud-init. This setup ensures the automated installation of dependencies, configuration of the environment, and deployment of the Cocos agent as a systemd service.
|
||||
|
||||
### Directory Contents
|
||||
|
||||
- **`config.yaml`**: This YAML file provides configuration instructions for the cloud image.
|
||||
- **`meta-data`**: Contains VM metadata, such as instance-specific details and identifiers.
|
||||
- **`qemu.sh`**: A Bash script for downloading and configuring a cloud image, running QEMU to simulate a VM with the cloud-init configuration.
|
||||
- **`.env`**: Contains environment variables for starting the VM in different modes, configuring disk space, memory allocation, and other parameters.
|
||||
|
||||
## Configuration
|
||||
|
||||
### Preparing the Cloud-Config File
|
||||
|
||||
The `config.yaml` file defines system configurations, including user creation, package installations, file management, and command execution.
|
||||
Ensure that the cloud-config file is set up with the following configurations:
|
||||
|
||||
- **User Credentials**: Specify the default username and password.
|
||||
- **Certificates and Keys**: Certificate files for agent for secure communication.
|
||||
- **Environment Variables**: Configuration parameters required by the system.
|
||||
|
||||
The `config.yaml` file is divided into multiple sections, each addressing a specific aspect of the setup process.
|
||||
|
||||
### 1. User Configuration
|
||||
|
||||
This section creates a default user with specific permissions and configurations:
|
||||
|
||||
- Creates a user named **`cocos_user`**.
|
||||
- Adds `cocos_user` to the `sudo` and `docker` groups.
|
||||
- Sets a default password (should be changed for production use).
|
||||
- Configures the user’s shell as `/bin/bash`.
|
||||
|
||||
### 2. Package Installation
|
||||
|
||||
Installs essential system packages required for various operations:
|
||||
|
||||
- **`curl`**: For downloading files from the web.
|
||||
- **`make`**: A utility for building software.
|
||||
- **`git`**: Version control system for managing code repositories.
|
||||
- **`python3` and `python3-dev`**: Required for running Python-based tools.
|
||||
- **`net-tools`**: Provides networking utilities such as `ifconfig` and `route`.
|
||||
|
||||
### 3. File Management (write_files)
|
||||
|
||||
Creates and configures critical files required for the setup:
|
||||
|
||||
- **Certificates**: Cert files (`cert.pem`, `ca.pem`, `key.pem`) located at `/etc/cocos/certs/`.
|
||||
- **Environment Variables**: An env file stored at `/etc/cocos/environment`.
|
||||
- **Systemd Service File**: Cocos agent service configuration file at `/etc/systemd/system/cocos-agent.service` for managing the Cocos agent.
|
||||
- **Agent Scripts**:
|
||||
- `agent_setup.sh`: Configures network interfaces and resizes the root filesystem.
|
||||
- `agent_start_script.sh`: Sets up Docker and starts the Cocos agent.
|
||||
|
||||
### 4. Execution of Commands (runcmd)
|
||||
|
||||
A sequence of commands is executed to finalize the setup:
|
||||
|
||||
- Creates necessary directories: `/cocos`, `/cocos_init`, `/var/log/cocos`, `/etc/cocos`.
|
||||
- Downloads and installs the Cocos agent binary.
|
||||
- Installs **Wasmtime** and configures its environment variables.
|
||||
- Installs **Docker** and adds `cocos_user` to the Docker group.
|
||||
- Reloads systemd and enables the Cocos agent service.
|
||||
|
||||
## Running the Agent
|
||||
|
||||
To test the cloud-init configuration, execute the `qemu.sh` script to bring up a VM using QEMU:
|
||||
|
||||
```bash
|
||||
sudo ./qemu.sh
|
||||
```
|
||||
|
||||
**Important:** The script must be executed as root.
|
||||
|
||||
Once the QEMU boots the VM, the Cocos agent will run as a systemd service. The service is configured to start automatically on boot and restart in case of failure.
|
||||
|
||||
## Debugging and Monitoring
|
||||
|
||||
For troubleshooting and monitoring the Cocos agent service, use the following commands within the VM:
|
||||
|
||||
### Manually Start the Service
|
||||
|
||||
To manually start the agent service, execute:
|
||||
|
||||
```bash
|
||||
sudo systemctl start cocos-agent.service
|
||||
```
|
||||
|
||||
### Verify Service Status
|
||||
|
||||
To check if the service is running properly, use:
|
||||
|
||||
```bash
|
||||
sudo systemctl status cocos-agent.service
|
||||
```
|
||||
|
||||
### View Service Logs
|
||||
|
||||
To inspect logs generated by the agent service, execute:
|
||||
|
||||
```bash
|
||||
journalctl -u cocos-agent.service
|
||||
```
|
||||
|
||||
### Check Standard Output and Error Logs
|
||||
|
||||
To check logs stored in the system, use:
|
||||
|
||||
```bash
|
||||
cat /var/log/cocos/agent.stdout.log
|
||||
cat /var/log/cocos/agent.stderr.log
|
||||
```
|
||||
@@ -1,174 +0,0 @@
|
||||
#cloud-config
|
||||
package_update: true
|
||||
package_upgrade: false
|
||||
|
||||
users:
|
||||
- default
|
||||
- name: cocos_user
|
||||
gecos: Default User
|
||||
groups:
|
||||
- sudo
|
||||
- docker # Add cocos user to the docker group
|
||||
sudo:
|
||||
- ALL=(ALL:ALL) ALL
|
||||
shell: /bin/bash
|
||||
|
||||
chpasswd:
|
||||
list: |
|
||||
cocos_user:password
|
||||
expire: False
|
||||
|
||||
ssh_pwauth: True
|
||||
|
||||
packages:
|
||||
- curl
|
||||
- make
|
||||
- git
|
||||
- python3
|
||||
- python3-dev
|
||||
- net-tools # Add net-tools to install the 'route' command
|
||||
|
||||
write_files:
|
||||
- path: /etc/cocos/certs/cert.pem
|
||||
content: |
|
||||
# Add certificate content here
|
||||
permissions: "0644"
|
||||
|
||||
- path: /etc/cocos/certs/ca.pem
|
||||
content: |
|
||||
# Add CA certificate content here
|
||||
permissions: "0644"
|
||||
|
||||
- path: /etc/cocos/certs/key.pem
|
||||
content: |
|
||||
# Add private key content here
|
||||
permissions: "0600"
|
||||
|
||||
- path: /etc/cocos/environment
|
||||
content: |
|
||||
# Add environment variables here
|
||||
permissions: "0644"
|
||||
|
||||
- path: /etc/systemd/system/cocos-agent.service
|
||||
content: |
|
||||
[Unit]
|
||||
Description=Cocos AI agent
|
||||
After=network.target
|
||||
Before=docker.service
|
||||
|
||||
[Service]
|
||||
WorkingDirectory=/cocos
|
||||
StandardOutput=file:/var/log/cocos/agent.stdout
|
||||
StandardError=file:/var/log/cocos/agent.stderr
|
||||
EnvironmentFile=/etc/cocos/environment
|
||||
ExecStartPre=/cocos_init/agent_setup.sh
|
||||
ExecStart=/cocos_init/agent_start_script.sh
|
||||
Restart=always
|
||||
|
||||
[Install]
|
||||
WantedBy=default.target
|
||||
permissions: "0644"
|
||||
|
||||
# Agent setup script
|
||||
- path: /cocos_init/agent_setup.sh
|
||||
content: |
|
||||
#!/bin/sh
|
||||
|
||||
WORK_DIR="/cocos"
|
||||
|
||||
# IFACES are all network interfaces excluding lo (LOOPBACK) and sit interfaces
|
||||
IFACES=$(ip link show | grep -vE 'LOOPBACK|sit*' | awk -F': ' '{print $2}')
|
||||
|
||||
# This for loop brings up all network interfaces in IFACES and dhclient obtains an IP address for the every interface
|
||||
for IFACE in $IFACES; do
|
||||
STATE=$(ip link show $IFACE | grep DOWN)
|
||||
if [ -n "$STATE" ]; then
|
||||
ip link set $IFACE up
|
||||
fi
|
||||
|
||||
IP_ADDR=$(ip addr show $IFACE | grep 'inet ')
|
||||
if [ -z "$IP_ADDR" ]; then
|
||||
dhclient $IFACE
|
||||
fi
|
||||
done
|
||||
|
||||
if [ ! -d "$WORK_DIR" ]; then
|
||||
mkdir -p $WORK_DIR
|
||||
fi
|
||||
|
||||
# Resize the root filesystem to 100% of available space
|
||||
ROOT_DEV=$(findmnt / -o SOURCE -n) # Get the root filesystem device
|
||||
resize2fs "$ROOT_DEV" && echo "Root filesystem resized successfully" || echo "Failed to resize root filesystem"
|
||||
permissions: "0755"
|
||||
|
||||
# Agent start script
|
||||
- path: /cocos_init/agent_start_script.sh
|
||||
content: |
|
||||
#!/bin/sh
|
||||
|
||||
# Change the docker.service file to allow Docker to run in RAM
|
||||
mkdir -p /etc/systemd/system/docker.service.d
|
||||
|
||||
# Create or overwrite the override.conf file with the new Environment variable
|
||||
tee /etc/systemd/system/docker.service.d/override.conf > /dev/null <<EOF
|
||||
[Service]
|
||||
Environment=DOCKER_RAMDISK=true
|
||||
EOF
|
||||
|
||||
systemctl daemon-reload
|
||||
|
||||
NUM_OF_PERMITED_IFACE=1
|
||||
|
||||
NUM_OF_IFACE=$(ip route | grep -Eo 'dev [a-z0-9]+' | awk '{ print $2 }' | grep -v '^docker' | sort | uniq | wc -l)
|
||||
|
||||
if [ $NUM_OF_IFACE -gt $NUM_OF_PERMITED_IFACE ]; then
|
||||
echo "More than one network interface in the VM"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
DEFAULT_IFACE=$(route | grep '^default' | grep -o '[^ ]*$')
|
||||
AGENT_GRPC_HOST=$(ip -4 addr show $DEFAULT_IFACE | grep inet | awk '{print $2}' | cut -d/ -f1)
|
||||
|
||||
export AGENT_GRPC_HOST
|
||||
|
||||
exec /bin/cocos-agent
|
||||
permissions: "0755"
|
||||
|
||||
runcmd:
|
||||
# Create necessary directories
|
||||
- mkdir -p /cocos
|
||||
- mkdir -p /cocos_init
|
||||
- mkdir -p /var/log/cocos
|
||||
- mkdir -p /etc/cocos
|
||||
|
||||
# Download the cocos-agent binary
|
||||
- echo "[ COCOS AGENT SETUP ] Downloading the cocos-agent binary..."
|
||||
- curl -L -O -J https://github.com/smithjilks/cocos/releases/download/v1.0.0/cocos-agent --progress-bar && echo "[ COCOS AGENT SETUP ] cocos-agent binary downloaded successfully" || echo "Failed to download cocos-agent binary"
|
||||
|
||||
# Install the agent binary
|
||||
- echo "[ COCOS AGENT SETUP ] Installing cocos-agent binary..."
|
||||
- install -D -m 0755 cocos-agent /bin/cocos-agent && echo "[ COCOS AGENT SETUP ] cocos-agent binary installed successfully" || echo "[ COCOS AGENT SETUP ] Failed to install cocos-agent binary"
|
||||
|
||||
# Install Wasmtime
|
||||
- echo "Installing Wasmtime runtime..."
|
||||
- curl https://wasmtime.dev/install.sh -sSf | bash && echo "Wasmtime installed successfully" || echo "Failed to install Wasmtime"
|
||||
- echo "Configuring Wasmtime environment variables..."
|
||||
- echo "export WASMTIME_HOME=$HOME/.wasmtime" >> /etc/profile.d/wasm_env.sh
|
||||
- echo "export PATH=\$WASMTIME_HOME/bin:\$PATH" >> /etc/profile.d/wasm_env.sh
|
||||
- . /etc/profile.d/wasm_env.sh && echo "Wasmtime environment variables configured successfully" || echo "Failed to configure Wasmtime environment variables"
|
||||
|
||||
# Install Docker
|
||||
- echo "Starting Docker installation..."
|
||||
- curl -fsSL https://get.docker.com -o get-docker.sh && echo "Docker install script downloaded successfully" || echo "Failed to download Docker install script"
|
||||
- sh ./get-docker.sh && echo "Docker installed successfully" || echo "Failed to install Docker"
|
||||
- usermod -aG docker cocos_user && echo "Added cocos_user to the docker group" || echo "Failed to add cocos_user to the docker group"
|
||||
|
||||
# Reload systemd and enable the service
|
||||
- echo "[ COCOS AGENT SETUP ] Reloading systemd daemon..."
|
||||
- systemctl daemon-reload && echo "[ COCOS AGENT SETUP ] Systemd daemon reloaded successfully" || echo "[ COCOS AGENT SETUP ] Failed to reload systemd daemon"
|
||||
- echo "[ COCOS AGENT SETUP ] Enabling cocos-agent.service..."
|
||||
- systemctl enable cocos-agent.service && echo "[ COCOS AGENT SETUP ] cocos-agent.service enabled successfully" || echo "[ COCOS AGENT SETUP ] Failed to enable cocos-agent.service"
|
||||
- echo "[ COCOS AGENT SETUP ] Starting cocos-agent.service..."
|
||||
- systemctl start cocos-agent.service && echo "[ COCOS AGENT SETUP ] cocos-agent.service started successfully" || echo "[ COCOS AGENT SETUP ] Failed to start cocos-agent.service"
|
||||
|
||||
final_message: "Cocos agent setup complete. Verify logs to confirm successful service startup."
|
||||
@@ -1,2 +0,0 @@
|
||||
instance-id: iid-cocos-vm
|
||||
local-hostname: cocos-vm
|
||||
@@ -1,124 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Source environment variables
|
||||
source ./.env
|
||||
|
||||
# Required commands
|
||||
REQUIRED_CMDS=("wget" "cloud-localds" "$QEMU_BINARY" "qemu-img")
|
||||
|
||||
# Check for required commands
|
||||
for cmd in "${REQUIRED_CMDS[@]}"; do
|
||||
if ! command -v "$cmd" &> /dev/null; then
|
||||
echo "Error: $cmd is not installed. Please install it and try again."
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
|
||||
# Ensure script is run as root
|
||||
if [[ $EUID -ne 0 ]]; then
|
||||
echo "Error: This script must be run as root."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Create the root filesystem image if it doesn't exist
|
||||
if [ ! -f "$BASE_IMAGE" ]; then
|
||||
echo "Downloading base Ubuntu image..."
|
||||
wget -q "$BASE_IMAGE_URL" -O "$BASE_IMAGE" --show-progress
|
||||
fi
|
||||
|
||||
# Create custom image
|
||||
echo "Creating custom QEMU image..."
|
||||
qemu-img create -f qcow2 -b "$BASE_IMAGE" -F qcow2 "$CUSTOM_IMAGE" "$DISK_SIZE"
|
||||
|
||||
# Cloud-init configuration files
|
||||
CLOUD_CONFIG="config.yaml"
|
||||
META_DATA="meta-data"
|
||||
SEED_IMAGE="seed.img"
|
||||
|
||||
# Create seed image for cloud-init
|
||||
echo "Creating seed image..."
|
||||
cloud-localds "$SEED_IMAGE" "$CLOUD_CONFIG" "$META_DATA"
|
||||
|
||||
# Construct QEMU arguments from environment variables
|
||||
construct_qemu_args() {
|
||||
args=()
|
||||
|
||||
args+=("-name" "$VM_NAME")
|
||||
|
||||
# Virtualization (Enable KVM)
|
||||
if [ "$ENABLE_KVM" == "true" ]; then
|
||||
args+=("-enable-kvm")
|
||||
fi
|
||||
|
||||
# Machine, CPU, RAM
|
||||
if [ -n "$MACHINE" ]; then
|
||||
args+=("-machine" "$MACHINE")
|
||||
fi
|
||||
|
||||
if [ -n "$CPU" ]; then
|
||||
args+=("-cpu" "$CPU")
|
||||
fi
|
||||
|
||||
args+=("-boot" "d")
|
||||
args+=("-smp" "$SMP_COUNT,maxcpus=$SMP_MAXCPUS")
|
||||
args+=("-m" "$MEMORY_SIZE,slots=$MEMORY_SLOTS,maxmem=$MAX_MEMORY")
|
||||
|
||||
# OVMF (if applicable)
|
||||
if [ "$ENABLE_SEV_SNP" != "true" ]; then
|
||||
args+=("-drive" "if=$OVMF_CODE_IF,format=$OVMF_CODE_FORMAT,unit=$OVMF_CODE_UNIT,file=$OVMF_CODE,readonly=$OVMF_CODE_READONLY")
|
||||
args+=("-drive" "if=$OVMF_VARS_IF,format=$OVMF_VARS_FORMAT,unit=$OVMF_VARS_UNIT,file=$OVMF_VARS")
|
||||
fi
|
||||
|
||||
# Network configuration
|
||||
args+=("-netdev" "user,id=$NET_DEV_ID,hostfwd=tcp::$NET_DEV_HOST_FWD_AGENT-:$NET_DEV_GUEST_FWD_AGENT")
|
||||
args+=("-device" "virtio-net-pci,disable-legacy=$VIRTIO_NET_PCI_DISABLE_LEGACY,iommu_platform=$VIRTIO_NET_PCI_IOMMU_PLATFORM,netdev=$NET_DEV_ID,addr=$VIRTIO_NET_PCI_ADDR,romfile=$VIRTIO_NET_PCI_ROMFILE")
|
||||
args+=("-device" "vhost-vsock-pci,id=$VSOCK_ID,guest-cid=$VSOCK_GUEST_CID")
|
||||
|
||||
# SEV (if enabled)
|
||||
if [ "$ENABLE_SEV" == "true" ] || [ "$ENABLE_SEV_SNP" == "true" ]; then
|
||||
sev_type="sev-guest"
|
||||
kernel_hash=""
|
||||
host_data=""
|
||||
|
||||
args+=("-machine" "confidential-guest-support=$SEV_ID,memory-backend=$MEM_ID")
|
||||
|
||||
if [ "$ENABLE_SEV_SNP" == "true" ]; then
|
||||
args+=("-bios" "$OVMF_CODE_FILE")
|
||||
sev_type="sev-snp-guest"
|
||||
|
||||
if [ -n "$SEV_HOST_DATA" ]; then
|
||||
host_data=",host-data=$SEV_HOST_DATA"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "$ENABLE_KERNEL_HASH" == "true" ]; then
|
||||
kernel_hash=",kernel-hashes=on"
|
||||
fi
|
||||
|
||||
args+=("-object" "memory-backend-memfd,id=$MEM_ID,size=$MEMORY_SIZE,share=true,prealloc=false")
|
||||
args+=("-object" "$sev_type,id=$SEV_ID,cbitpos=$SEV_CBIT_POS,reduced-phys-bits=$SEV_REDUCED_PHYS_BITS$kernel_hash$host_data")
|
||||
fi
|
||||
|
||||
# Disk image configuration
|
||||
args+=("-drive" "file=$SEED_IMAGE,media=cdrom")
|
||||
args+=("-drive" "file=$CUSTOM_IMAGE,if=none,id=disk0,format=qcow2")
|
||||
args+=("-device" "virtio-scsi-pci,id=scsi,disable-legacy=on,iommu_platform=true")
|
||||
args+=("-device" "scsi-hd,drive=disk0")
|
||||
|
||||
# Display options
|
||||
if [ "$NO_GRAPHIC" == "true" ]; then
|
||||
args+=("-nographic")
|
||||
fi
|
||||
|
||||
args+=("-monitor" "$MONITOR")
|
||||
args+=("-no-reboot")
|
||||
args+=("-vnc" ":9")
|
||||
|
||||
echo "${args[@]}"
|
||||
}
|
||||
|
||||
qemu_args=$(construct_qemu_args)
|
||||
|
||||
echo "Running QEMU with the following arguments: $qemu_args"
|
||||
echo "Starting QEMU VM..."
|
||||
$QEMU_BINARY $qemu_args
|
||||
@@ -65,14 +65,3 @@ CONFIG_PREEMPT_DYNAMIC=n
|
||||
CONFIG_DEBUG_PREEMPT=n
|
||||
CONFIG_CGROUP_MISC=y
|
||||
CONFIG_X86_CPUID=y
|
||||
|
||||
CONFIG_NET_9P=y
|
||||
CONFIG_NET_9P_VIRTIO=y
|
||||
CONFIG_9P_FS=y
|
||||
CONFIG_9P_FS_POSIX_ACL=y
|
||||
CONFIG_9P_FS_SECURITY=y
|
||||
|
||||
# TCG TPM
|
||||
CONFIG_TCG_TPM=y
|
||||
CONFIG_TCG_TPM2_HMAC=y
|
||||
CONFIG_TCG_PLATFORM=y
|
||||
|
||||
@@ -6,23 +6,6 @@ set -e
|
||||
# Add a console on tty1
|
||||
if [ -e ${TARGET_DIR}/etc/inittab ]; then
|
||||
grep -qE '^tty1::' ${TARGET_DIR}/etc/inittab || \
|
||||
sed -i '/GENERIC_SERIAL/a\
|
||||
sed -i '/GENERIC_SERIAL/a\
|
||||
tty1::respawn:/sbin/getty -L tty1 0 vt100 # QEMU graphical window' ${TARGET_DIR}/etc/inittab
|
||||
fi
|
||||
|
||||
# Create the mount points
|
||||
# Create the mount points
|
||||
mkdir -p ${TARGET_DIR}/etc/certs
|
||||
mkdir -p ${TARGET_DIR}/etc/cocos
|
||||
|
||||
# Ensure /etc/fstab exists
|
||||
if [ ! -f "${TARGET_DIR}/etc/fstab" ]; then
|
||||
touch "${TARGET_DIR}/etc/fstab"
|
||||
fi
|
||||
|
||||
# Add the 9p entries to /etc/fstab
|
||||
grep -q "certs_share /etc/certs" ${TARGET_DIR}/etc/fstab || \
|
||||
echo "certs_share /etc/certs 9p trans=virtio,version=9p2000.L,cache=mmap 0 0" >> "${TARGET_DIR}/etc/fstab"
|
||||
|
||||
grep -q "env_share /etc/cocos" ${TARGET_DIR}/etc/fstab || \
|
||||
echo "env_share /etc/cocos 9p trans=virtio,version=9p2000.L,cache=mmap 0 0" >> "${TARGET_DIR}/etc/fstab"
|
||||
|
||||
@@ -14,8 +14,6 @@ BR2_SYSTEM_BIN_SH_BASH=y
|
||||
BR2_TARGET_ROOTFS_CPIO=y
|
||||
BR2_TARGET_ROOTFS_CPIO_FULL=y
|
||||
BR2_TARGET_ROOTFS_CPIO_GZIP=y
|
||||
BR2_TARGET_ROOTFS_OVERLAY="overlay"
|
||||
BR2_PACKAGE_9PFS=y
|
||||
|
||||
# Image
|
||||
BR2_ROOTFS_POST_BUILD_SCRIPT="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/post-build.sh"
|
||||
@@ -27,14 +25,14 @@ BR2_ROOTFS_POST_SCRIPT_ARGS="$(BR2_DEFCONFIG)"
|
||||
# Linux headers same as kernel
|
||||
BR2_PACKAGE_HOST_LINUX_HEADERS_CUSTOM_6_11=y
|
||||
BR2_TOOLCHAIN_HEADERS_LATEST=y
|
||||
BR2_TOOLCHAIN_HEADERS_AT_LEAST="6.11-rc7"
|
||||
BR2_TOOLCHAIN_HEADERS_AT_LEAST="6.12-rc6"
|
||||
|
||||
# Kernel
|
||||
BR2_LINUX_KERNEL=y
|
||||
BR2_LINUX_KERNEL_CUSTOM_GIT=y
|
||||
BR2_LINUX_KERNEL_CUSTOM_REPO_URL="https://github.com/coconut-svsm/linux.git"
|
||||
BR2_LINUX_KERNEL_CUSTOM_REPO_VERSION="svsm"
|
||||
BR2_LINUX_KERNEL_VERSION="svsm"
|
||||
BR2_LINUX_KERNEL_CUSTOM_REPO_URL="https://github.com/torvalds/linux.git"
|
||||
BR2_LINUX_KERNEL_CUSTOM_REPO_VERSION="v6.12-rc6"
|
||||
BR2_LINUX_KERNEL_VERSION="v6.12-rc6"
|
||||
BR2_LINUX_KERNEL_PATCH=""
|
||||
BR2_LINUX_KERNEL_USE_CUSTOM_CONFIG=y
|
||||
BR2_LINUX_KERNEL_CUSTOM_CONFIG_FILE="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/linux.config"
|
||||
|
||||
@@ -8,7 +8,8 @@ WorkingDirectory=/cocos
|
||||
StandardOutput=file:/var/log/cocos/agent.stdout
|
||||
StandardError=file:/var/log/cocos/agent.stderr
|
||||
|
||||
EnvironmentFile=/etc/cocos/environment
|
||||
Environment=AGENT_GRPC_PORT=7002
|
||||
Environment=AGENT_LOG_LEVEL=info
|
||||
|
||||
ExecStartPre=/cocos_init/agent_setup.sh
|
||||
ExecStart=/cocos_init/agent_start_script.sh
|
||||
|
||||
@@ -7,9 +7,8 @@ import (
|
||||
"io"
|
||||
"log/slog"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
"google.golang.org/protobuf/encoding/protojson"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
)
|
||||
|
||||
@@ -19,17 +18,16 @@ type handler struct {
|
||||
opts slog.HandlerOptions
|
||||
w io.Writer
|
||||
cmpID string
|
||||
queue chan *cvms.ClientStreamMessage
|
||||
}
|
||||
|
||||
func NewProtoHandler(conn io.Writer, opts *slog.HandlerOptions, queue chan *cvms.ClientStreamMessage) slog.Handler {
|
||||
func NewProtoHandler(conn io.Writer, opts *slog.HandlerOptions, cmpID string) slog.Handler {
|
||||
if opts == nil {
|
||||
opts = &slog.HandlerOptions{}
|
||||
}
|
||||
h := &handler{
|
||||
opts: *opts,
|
||||
w: conn,
|
||||
queue: queue,
|
||||
cmpID: cmpID,
|
||||
}
|
||||
|
||||
return h
|
||||
@@ -71,18 +69,7 @@ func (h *handler) Handle(_ context.Context, r slog.Record) error {
|
||||
},
|
||||
}
|
||||
|
||||
h.queue <- &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_AgentLog{
|
||||
AgentLog: &cvms.AgentLog{
|
||||
Timestamp: timestamp,
|
||||
Message: chunk,
|
||||
Level: level,
|
||||
ComputationId: h.cmpID,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
b, err := protojson.Marshal(&agentLog)
|
||||
b, err := proto.Marshal(&agentLog)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -91,11 +78,6 @@ func (h *handler) Handle(_ context.Context, r slog.Record) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = h.w.Write([]byte("\n"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -106,8 +88,7 @@ func (h *handler) WithAttrs(attrs []slog.Attr) slog.Handler {
|
||||
}
|
||||
|
||||
func (h *handler) WithGroup(name string) slog.Handler {
|
||||
h.cmpID = name
|
||||
return h
|
||||
panic("unimplemented")
|
||||
}
|
||||
|
||||
func (h *handler) Close() error {
|
||||
|
||||
@@ -11,7 +11,6 @@ import (
|
||||
|
||||
"github.com/absmach/magistrala/pkg/errors"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
)
|
||||
|
||||
type failedWriter struct{}
|
||||
@@ -22,14 +21,14 @@ func (f *failedWriter) Write(p []byte) (n int, err error) {
|
||||
|
||||
// TestNewProtoHandler tests the initialization of the ProtoHandler.
|
||||
func TestNewProtoHandler(t *testing.T) {
|
||||
handler := NewProtoHandler(io.Discard, nil, make(chan *cvms.ClientStreamMessage))
|
||||
handler := NewProtoHandler(io.Discard, nil, "testCmpID")
|
||||
|
||||
assert.NotNil(t, handler, "Handler should not be nil")
|
||||
}
|
||||
|
||||
// TestHandleMessageSuccess tests the handling of a message when the write succeeds.
|
||||
func TestHandleMessageSuccess(t *testing.T) {
|
||||
handler := NewProtoHandler(io.Discard, nil, make(chan *cvms.ClientStreamMessage, 1))
|
||||
handler := NewProtoHandler(io.Discard, nil, "testCmpID")
|
||||
record := slog.Record{
|
||||
Time: time.Now(),
|
||||
Message: "Test message",
|
||||
@@ -43,7 +42,7 @@ func TestHandleMessageSuccess(t *testing.T) {
|
||||
|
||||
// TestHandleMessageFailure tests the caching mechanism when the write fails.
|
||||
func TestHandleMessageFailure(t *testing.T) {
|
||||
protohandler := NewProtoHandler(&failedWriter{}, nil, make(chan *cvms.ClientStreamMessage, 1))
|
||||
protohandler := NewProtoHandler(&failedWriter{}, nil, "testCmpID")
|
||||
record := slog.Record{
|
||||
Time: time.Now(),
|
||||
Message: "Test message",
|
||||
@@ -57,7 +56,7 @@ func TestHandleMessageFailure(t *testing.T) {
|
||||
|
||||
// TestEnabled tests that the handler enables logging based on level.
|
||||
func TestEnabled(t *testing.T) {
|
||||
handler := NewProtoHandler(io.Discard, nil, make(chan *cvms.ClientStreamMessage, 1))
|
||||
handler := NewProtoHandler(io.Discard, nil, "testCmpID")
|
||||
|
||||
assert.True(t, handler.Enabled(context.Background(), slog.LevelInfo), "Logging should be enabled for LevelInfo")
|
||||
assert.False(t, handler.Enabled(context.Background(), slog.LevelDebug), "Logging should be disabled for LevelDebug by default")
|
||||
@@ -67,7 +66,7 @@ func TestEnabled(t *testing.T) {
|
||||
func TestCloseStopsRetry(t *testing.T) {
|
||||
mockWriter := io.Discard
|
||||
|
||||
handler := NewProtoHandler(mockWriter, nil, make(chan *cvms.ClientStreamMessage, 1)).(*handler)
|
||||
handler := NewProtoHandler(mockWriter, nil, "testCmpID").(*handler)
|
||||
|
||||
time.Sleep(2 * time.Second)
|
||||
err := handler.Close()
|
||||
|
||||
@@ -25,7 +25,6 @@ import (
|
||||
"github.com/ultravioletrs/cocos/agent/auth"
|
||||
"github.com/ultravioletrs/cocos/internal/server"
|
||||
"github.com/ultravioletrs/cocos/pkg/atls"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/credentials"
|
||||
@@ -52,7 +51,7 @@ type Server struct {
|
||||
server.BaseServer
|
||||
server *grpc.Server
|
||||
registerService serviceRegister
|
||||
quoteProvider client.LeveledQuoteProvider
|
||||
quoteProvider client.QuoteProvider
|
||||
authSvc auth.Authenticator
|
||||
health *health.Server
|
||||
}
|
||||
@@ -61,7 +60,7 @@ type serviceRegister func(srv *grpc.Server)
|
||||
|
||||
var _ server.Server = (*Server)(nil)
|
||||
|
||||
func New(ctx context.Context, cancel context.CancelFunc, name string, config server.ServerConfiguration, registerService serviceRegister, logger *slog.Logger, qp client.LeveledQuoteProvider, authSvc auth.Authenticator) server.Server {
|
||||
func New(ctx context.Context, cancel context.CancelFunc, name string, config server.ServerConfiguration, registerService serviceRegister, logger *slog.Logger, qp client.QuoteProvider, authSvc auth.Authenticator) server.Server {
|
||||
base := config.GetBaseConfig()
|
||||
listenFullAddress := fmt.Sprintf("%s:%s", base.Host, base.Port)
|
||||
return &Server{
|
||||
@@ -302,19 +301,5 @@ func generateCertificatesForATLS() ([]byte, []byte, error) {
|
||||
Bytes: privateKeyBytes,
|
||||
})
|
||||
|
||||
cert, err := x509.ParseCertificate(certDERBytes)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to parse certificate: %w", err)
|
||||
}
|
||||
|
||||
pubKeyDER, err := x509.MarshalPKIXPublicKey(cert.PublicKey)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to marshal public key to DER format: %w", err)
|
||||
}
|
||||
|
||||
if err := vtpm.ExtendPCR(vtpm.PCR15, pubKeyDER); err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to extend vTPM PCR with public key: %w", err)
|
||||
}
|
||||
|
||||
return certBytes, keyBytes, nil
|
||||
}
|
||||
|
||||
@@ -22,7 +22,6 @@ import (
|
||||
authmocks "github.com/ultravioletrs/cocos/agent/auth/mocks"
|
||||
"github.com/ultravioletrs/cocos/internal/server"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/test/bufconn"
|
||||
)
|
||||
@@ -31,28 +30,6 @@ const bufSize = 1024 * 1024
|
||||
|
||||
var lis *bufconn.Listener
|
||||
|
||||
type DummyRWC struct{}
|
||||
|
||||
// Read fills p with byte(len(p)) and returns len(p).
|
||||
func (l *DummyRWC) Read(p []byte) (int, error) {
|
||||
n := len(p)
|
||||
// Fill each byte in p with the value of n as a byte.
|
||||
for i := range p {
|
||||
p[i] = byte(n)
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// Write simply returns len(p) indicating that all bytes were written.
|
||||
func (l *DummyRWC) Write(p []byte) (int, error) {
|
||||
// In this simple implementation, we ignore the data.
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (l *DummyRWC) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func init() {
|
||||
lis = bufconn.Listen(bufSize)
|
||||
}
|
||||
@@ -70,7 +47,7 @@ func TestNew(t *testing.T) {
|
||||
},
|
||||
}
|
||||
logger := slog.Default()
|
||||
qp := new(mocks.LeveledQuoteProvider)
|
||||
qp := new(mocks.QuoteProvider)
|
||||
authSvc := new(authmocks.Authenticator)
|
||||
|
||||
srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
|
||||
@@ -120,7 +97,7 @@ func TestServerStartWithTLSFile(t *testing.T) {
|
||||
|
||||
logBuffer := &ThreadSafeBuffer{}
|
||||
logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
|
||||
qp := new(mocks.LeveledQuoteProvider)
|
||||
qp := new(mocks.QuoteProvider)
|
||||
authSvc := new(authmocks.Authenticator)
|
||||
|
||||
srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
|
||||
@@ -167,7 +144,7 @@ func TestServerStartWithmTLSFile(t *testing.T) {
|
||||
|
||||
logBuffer := &ThreadSafeBuffer{}
|
||||
logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
|
||||
qp := new(mocks.LeveledQuoteProvider)
|
||||
qp := new(mocks.QuoteProvider)
|
||||
authSvc := new(authmocks.Authenticator)
|
||||
|
||||
srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
|
||||
@@ -207,7 +184,7 @@ func TestServerStop(t *testing.T) {
|
||||
}
|
||||
buf := &ThreadSafeBuffer{}
|
||||
logger := slog.New(slog.NewTextHandler(buf, &slog.HandlerOptions{Level: slog.LevelDebug}))
|
||||
qp := new(mocks.LeveledQuoteProvider)
|
||||
qp := new(mocks.QuoteProvider)
|
||||
authSvc := new(authmocks.Authenticator)
|
||||
|
||||
srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
|
||||
@@ -282,8 +259,6 @@ func (b *ThreadSafeBuffer) String() string {
|
||||
}
|
||||
|
||||
func TestServerInitializationAndStartup(t *testing.T) {
|
||||
vtpm.ExternalTPM = &DummyRWC{}
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
config server.AgentConfig
|
||||
@@ -399,7 +374,7 @@ func TestServerInitializationAndStartup(t *testing.T) {
|
||||
|
||||
logBuffer := &ThreadSafeBuffer{}
|
||||
logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
|
||||
qp := new(mocks.LeveledQuoteProvider)
|
||||
qp := new(mocks.QuoteProvider)
|
||||
authSvc := new(authmocks.Authenticator)
|
||||
|
||||
srv := New(ctx, cancel, "TestServer", tc.config, func(srv *grpc.Server) {}, logger, qp, authSvc)
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package vsock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
const (
|
||||
maxRetries = 3
|
||||
retryDelay = time.Second
|
||||
maxMessageSize = 1 << 20 // 1 MB
|
||||
ackTimeout = 5 * time.Second
|
||||
maxConcurrent = 100
|
||||
)
|
||||
|
||||
type MessageStatus int
|
||||
|
||||
const (
|
||||
StatusPending MessageStatus = iota
|
||||
StatusSent
|
||||
StatusAcknowledged
|
||||
StatusFailed
|
||||
)
|
||||
|
||||
type Message struct {
|
||||
ID uint32
|
||||
Content []byte
|
||||
Status MessageStatus
|
||||
Retries int
|
||||
}
|
||||
|
||||
type AckWriter struct {
|
||||
conn net.Conn
|
||||
pendingMessages chan *Message
|
||||
messageStore sync.Map // map[uint32]*Message
|
||||
nextID uint32
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
func NewAckWriter(conn net.Conn) io.WriteCloser {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
aw := &AckWriter{
|
||||
conn: conn,
|
||||
pendingMessages: make(chan *Message, maxConcurrent),
|
||||
nextID: 1,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
aw.wg.Add(2)
|
||||
go aw.sendMessages()
|
||||
go aw.handleAcknowledgments()
|
||||
return aw
|
||||
}
|
||||
|
||||
func (aw *AckWriter) Write(p []byte) (int, error) {
|
||||
if len(p) > maxMessageSize {
|
||||
return 0, fmt.Errorf("message size exceeds maximum allowed size of %d bytes", maxMessageSize)
|
||||
}
|
||||
|
||||
messageID := atomic.AddUint32(&aw.nextID, 1)
|
||||
message := &Message{
|
||||
ID: messageID,
|
||||
Content: make([]byte, len(p)),
|
||||
Status: StatusPending,
|
||||
}
|
||||
copy(message.Content, p)
|
||||
|
||||
aw.messageStore.Store(messageID, message)
|
||||
select {
|
||||
case aw.pendingMessages <- message:
|
||||
return len(p), nil
|
||||
case <-aw.ctx.Done():
|
||||
return 0, fmt.Errorf("writer is closed")
|
||||
}
|
||||
}
|
||||
|
||||
func (aw *AckWriter) sendMessages() {
|
||||
defer aw.wg.Done()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-aw.ctx.Done():
|
||||
return
|
||||
case msg := <-aw.pendingMessages:
|
||||
if err := aw.sendWithRetry(msg); err != nil {
|
||||
log.Printf("Failed to send message %d after all retries: %v", msg.ID, err)
|
||||
msg.Status = StatusFailed
|
||||
aw.messageStore.Store(msg.ID, msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (aw *AckWriter) sendWithRetry(msg *Message) error {
|
||||
for msg.Retries < maxRetries {
|
||||
if err := aw.writeMessage(msg.ID, msg.Content); err != nil {
|
||||
msg.Retries++
|
||||
msg.Status = StatusPending
|
||||
log.Printf("Error writing message %d (attempt %d): %v", msg.ID, msg.Retries, err)
|
||||
time.Sleep(retryDelay)
|
||||
continue
|
||||
}
|
||||
msg.Status = StatusSent
|
||||
aw.messageStore.Store(msg.ID, msg)
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("max retries reached")
|
||||
}
|
||||
|
||||
func (aw *AckWriter) writeMessage(messageID uint32, p []byte) error {
|
||||
if err := binary.Write(aw.conn, binary.LittleEndian, messageID); err != nil {
|
||||
return fmt.Errorf("failed to write message ID: %w", err)
|
||||
}
|
||||
|
||||
messageLen := uint32(len(p))
|
||||
if err := binary.Write(aw.conn, binary.LittleEndian, messageLen); err != nil {
|
||||
return fmt.Errorf("failed to write message length: %w", err)
|
||||
}
|
||||
|
||||
if _, err := aw.conn.Write(p); err != nil {
|
||||
return fmt.Errorf("failed to write message content: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (aw *AckWriter) handleAcknowledgments() {
|
||||
defer aw.wg.Done()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-aw.ctx.Done():
|
||||
return
|
||||
default:
|
||||
var ackID uint32
|
||||
if err := binary.Read(aw.conn, binary.LittleEndian, &ackID); err != nil {
|
||||
if err == io.EOF {
|
||||
log.Println("Connection closed, stopping acknowledgment handler")
|
||||
return
|
||||
}
|
||||
log.Printf("Error reading ACK: %v", err)
|
||||
time.Sleep(retryDelay)
|
||||
continue
|
||||
}
|
||||
|
||||
if msg, ok := aw.messageStore.Load(ackID); ok {
|
||||
m := msg.(*Message)
|
||||
m.Status = StatusAcknowledged
|
||||
aw.messageStore.Store(ackID, m)
|
||||
|
||||
// Clean up old messages periodically
|
||||
go aw.cleanupOldMessages(ackID)
|
||||
} else {
|
||||
log.Printf("Received ACK for unknown message ID: %d", ackID)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (aw *AckWriter) cleanupOldMessages(currentID uint32) {
|
||||
aw.messageStore.Range(func(key, value interface{}) bool {
|
||||
msgID := key.(uint32)
|
||||
msg := value.(*Message)
|
||||
|
||||
// Clean up acknowledged messages that are old
|
||||
if msg.Status == StatusAcknowledged && msgID < currentID-maxConcurrent {
|
||||
aw.messageStore.Delete(msgID)
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
func (aw *AckWriter) Close() error {
|
||||
aw.cancel()
|
||||
aw.wg.Wait()
|
||||
return aw.conn.Close()
|
||||
}
|
||||
|
||||
type Reader interface {
|
||||
Read() ([]byte, error)
|
||||
ReadProto(msg proto.Message) error
|
||||
}
|
||||
|
||||
type AckReader struct {
|
||||
conn net.Conn
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func NewAckReader(conn net.Conn) Reader {
|
||||
return &AckReader{
|
||||
conn: conn,
|
||||
ctx: context.Background(),
|
||||
}
|
||||
}
|
||||
|
||||
func (ar *AckReader) ReadProto(msg proto.Message) error {
|
||||
data, err := ar.Read()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read proto message: %w", err)
|
||||
}
|
||||
return proto.Unmarshal(data, msg)
|
||||
}
|
||||
|
||||
func (ar *AckReader) Read() ([]byte, error) {
|
||||
var messageID uint32
|
||||
if err := binary.Read(ar.conn, binary.LittleEndian, &messageID); err != nil {
|
||||
return nil, fmt.Errorf("error reading message ID: %w", err)
|
||||
}
|
||||
|
||||
var messageLen uint32
|
||||
if err := binary.Read(ar.conn, binary.LittleEndian, &messageLen); err != nil {
|
||||
return nil, fmt.Errorf("error reading message length: %w", err)
|
||||
}
|
||||
|
||||
if messageLen > maxMessageSize {
|
||||
return nil, fmt.Errorf("message size %d exceeds maximum allowed size of %d bytes", messageLen, maxMessageSize)
|
||||
}
|
||||
|
||||
data := make([]byte, messageLen)
|
||||
if _, err := io.ReadFull(ar.conn, data); err != nil {
|
||||
return nil, fmt.Errorf("error reading message content: %w", err)
|
||||
}
|
||||
|
||||
if err := ar.sendAck(messageID); err != nil {
|
||||
return nil, fmt.Errorf("error sending ACK: %w", err)
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func (ar *AckReader) sendAck(messageID uint32) error {
|
||||
return binary.Write(ar.conn, binary.LittleEndian, messageID)
|
||||
}
|
||||
@@ -0,0 +1,337 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package vsock
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
// MockConn implements net.Conn for testing purposes.
|
||||
type MockConn struct {
|
||||
ReadData []byte
|
||||
WrittenData []byte
|
||||
ReadErr error
|
||||
WriteErr error
|
||||
closed bool
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func (m *MockConn) Read(b []byte) (n int, err error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.closed {
|
||||
return 0, io.EOF
|
||||
}
|
||||
if len(m.ReadData) == 0 {
|
||||
return 0, io.EOF // Ensure we handle this case more predictably
|
||||
}
|
||||
if m.ReadErr != nil {
|
||||
return 0, m.ReadErr
|
||||
}
|
||||
n = copy(b, m.ReadData)
|
||||
m.ReadData = m.ReadData[n:]
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (m *MockConn) Write(b []byte) (n int, err error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.closed {
|
||||
return 0, errors.New("connection closed")
|
||||
}
|
||||
if m.WriteErr != nil {
|
||||
return 0, m.WriteErr
|
||||
}
|
||||
m.WrittenData = append(m.WrittenData, b...)
|
||||
return len(b), nil
|
||||
}
|
||||
|
||||
func (m *MockConn) Close() error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// Implement other net.Conn methods with empty implementations.
|
||||
func (m *MockConn) LocalAddr() net.Addr { return nil }
|
||||
func (m *MockConn) RemoteAddr() net.Addr { return nil }
|
||||
func (m *MockConn) SetDeadline(t time.Time) error { return nil }
|
||||
func (m *MockConn) SetReadDeadline(t time.Time) error { return nil }
|
||||
func (m *MockConn) SetWriteDeadline(t time.Time) error { return nil }
|
||||
|
||||
func TestAckReader_Read(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
wantErr bool
|
||||
}{
|
||||
{"Valid message", []byte("Hello, World!"), false},
|
||||
{"Empty message", []byte{}, false},
|
||||
{"Message at max size", make([]byte, maxMessageSize), false},
|
||||
{"Message exceeds max size", make([]byte, maxMessageSize+1), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockConn := &MockConn{}
|
||||
ar := NewAckReader(mockConn)
|
||||
|
||||
// Prepare mock data
|
||||
messageID := uint32(1)
|
||||
messageLen := uint32(len(tt.data))
|
||||
mockData := make([]byte, 8+len(tt.data))
|
||||
binary.LittleEndian.PutUint32(mockData[:4], messageID)
|
||||
binary.LittleEndian.PutUint32(mockData[4:8], messageLen)
|
||||
copy(mockData[8:], tt.data)
|
||||
mockConn.ReadData = mockData
|
||||
|
||||
data, err := ar.Read()
|
||||
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("AckReader.Read() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
|
||||
if !tt.wantErr {
|
||||
if !bytes.Equal(data, tt.data) {
|
||||
t.Errorf("AckReader.Read() got = %v, want %v", data, tt.data)
|
||||
}
|
||||
|
||||
// Check if ACK was sent
|
||||
if len(mockConn.WrittenData) != 4 {
|
||||
t.Errorf("AckReader.Read() did not send ACK")
|
||||
} else {
|
||||
ackID := binary.LittleEndian.Uint32(mockConn.WrittenData)
|
||||
if ackID != messageID {
|
||||
t.Errorf("AckReader.Read() sent wrong ACK ID, got %d, want %d", ackID, messageID)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAckReader_ReadProto(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
msg *manager.ClientStreamMessage
|
||||
wantErr bool
|
||||
}{
|
||||
{"Valid proto message", &manager.ClientStreamMessage{}, false},
|
||||
{"Empty proto message", &manager.ClientStreamMessage{}, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockConn := &MockConn{}
|
||||
ar := NewAckReader(mockConn)
|
||||
|
||||
// Prepare mock data
|
||||
protoData, _ := proto.Marshal(tt.msg)
|
||||
messageID := uint32(1)
|
||||
messageLen := uint32(len(protoData))
|
||||
mockData := make([]byte, 8+len(protoData))
|
||||
binary.LittleEndian.PutUint32(mockData[:4], messageID)
|
||||
binary.LittleEndian.PutUint32(mockData[4:8], messageLen)
|
||||
copy(mockData[8:], protoData)
|
||||
mockConn.ReadData = mockData
|
||||
|
||||
receivedMsg := &manager.ClientStreamMessage{}
|
||||
err := ar.ReadProto(receivedMsg)
|
||||
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("AckReader.ReadProto() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
|
||||
if !tt.wantErr {
|
||||
if receivedMsg.Message != tt.msg.Message {
|
||||
t.Errorf("AckReader.ReadProto() got = %v, want %v", receivedMsg, tt.msg)
|
||||
}
|
||||
|
||||
// Check if ACK was sent
|
||||
if len(mockConn.WrittenData) != 4 {
|
||||
t.Errorf("AckReader.ReadProto() did not send ACK")
|
||||
} else {
|
||||
ackID := binary.LittleEndian.Uint32(mockConn.WrittenData)
|
||||
if ackID != messageID {
|
||||
t.Errorf("AckReader.ReadProto() sent wrong ACK ID, got %d, want %d", ackID, messageID)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAckWriter(t *testing.T) {
|
||||
mockConn := &MockConn{}
|
||||
writer := NewAckWriter(mockConn)
|
||||
|
||||
if _, ok := writer.(io.Writer); !ok {
|
||||
t.Errorf("NewAckWriter() did not return an io.Writer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAckReader(t *testing.T) {
|
||||
mockConn := &MockConn{}
|
||||
reader := NewAckReader(mockConn)
|
||||
|
||||
assert.NotNil(t, reader)
|
||||
}
|
||||
|
||||
func TestAckWriter_Close(t *testing.T) {
|
||||
mockConn := &MockConn{}
|
||||
aw := NewAckWriter(mockConn)
|
||||
|
||||
err := aw.Close()
|
||||
if err != nil {
|
||||
t.Errorf("AckWriter.Close() error = %v, wantErr %v", err, nil)
|
||||
}
|
||||
|
||||
if !mockConn.closed {
|
||||
t.Errorf("AckWriter.Close() did not close the connection")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAckWriter_Write(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input []byte
|
||||
expectErr bool
|
||||
expectedError string
|
||||
}{
|
||||
{
|
||||
name: "Message exceeds max size",
|
||||
input: make([]byte, maxMessageSize+1),
|
||||
expectErr: true,
|
||||
expectedError: "message size exceeds maximum allowed size",
|
||||
},
|
||||
{
|
||||
name: "Write succeeds",
|
||||
input: []byte("Hello, world!"),
|
||||
expectErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockConn := &MockConn{
|
||||
mu: sync.Mutex{},
|
||||
}
|
||||
|
||||
writer := NewAckWriter(mockConn)
|
||||
defer writer.Close()
|
||||
|
||||
if tt.expectErr {
|
||||
writer.(*AckWriter).ctx.Done()
|
||||
}
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
n, err := writer.Write(tt.input)
|
||||
|
||||
if tt.expectErr {
|
||||
assert.Error(t, err)
|
||||
if tt.expectedError != "" {
|
||||
assert.Contains(t, err.Error(), tt.expectedError)
|
||||
}
|
||||
assert.Zero(t, n)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, len(tt.input), n)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAckWriter_CleanupOldMessages(t *testing.T) {
|
||||
mockConn := &MockConn{}
|
||||
writer := NewAckWriter(mockConn).(*AckWriter)
|
||||
defer writer.Close()
|
||||
|
||||
for i := uint32(1); i <= maxConcurrent+10; i++ {
|
||||
msg := &Message{
|
||||
ID: i,
|
||||
Content: []byte("test"),
|
||||
Status: StatusAcknowledged,
|
||||
}
|
||||
writer.messageStore.Store(i, msg)
|
||||
}
|
||||
|
||||
writer.cleanupOldMessages(maxConcurrent + 11)
|
||||
|
||||
var count int
|
||||
writer.messageStore.Range(func(key, value interface{}) bool {
|
||||
count++
|
||||
return true
|
||||
})
|
||||
|
||||
assert.LessOrEqual(t, count, maxConcurrent)
|
||||
}
|
||||
|
||||
func TestAckReader_LargeMessage(t *testing.T) {
|
||||
mockConn := &MockConn{}
|
||||
reader := NewAckReader(mockConn)
|
||||
|
||||
largeMessage := make([]byte, maxMessageSize-1)
|
||||
for i := range largeMessage {
|
||||
largeMessage[i] = byte(i % 256)
|
||||
}
|
||||
|
||||
messageID := uint32(1)
|
||||
messageLen := uint32(len(largeMessage))
|
||||
mockData := make([]byte, 8+len(largeMessage))
|
||||
binary.LittleEndian.PutUint32(mockData[:4], messageID)
|
||||
binary.LittleEndian.PutUint32(mockData[4:8], messageLen)
|
||||
copy(mockData[8:], largeMessage)
|
||||
mockConn.ReadData = mockData
|
||||
|
||||
data, err := reader.Read()
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, largeMessage, data)
|
||||
|
||||
assert.Equal(t, 4, len(mockConn.WrittenData))
|
||||
ackID := binary.LittleEndian.Uint32(mockConn.WrittenData)
|
||||
assert.Equal(t, messageID, ackID)
|
||||
}
|
||||
|
||||
func TestAckWriter_FailedSends(t *testing.T) {
|
||||
mockConn := &MockConn{
|
||||
WriteErr: errors.New("write error"),
|
||||
}
|
||||
writer := NewAckWriter(mockConn).(*AckWriter)
|
||||
defer writer.Close()
|
||||
|
||||
// Add some messages to the channel
|
||||
for i := 0; i < 5; i++ {
|
||||
msg := &Message{
|
||||
ID: uint32(i + 1),
|
||||
Content: []byte(fmt.Sprintf("Message %d", i+1)),
|
||||
Status: StatusPending,
|
||||
}
|
||||
writer.pendingMessages <- msg
|
||||
}
|
||||
|
||||
// Wait for the messages to be sent
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Check that the messages were marked as failed
|
||||
writer.messageStore.Range(func(key, value interface{}) bool {
|
||||
msg := value.(*Message)
|
||||
assert.Equal(t, StatusFailed, msg.Status)
|
||||
return true
|
||||
})
|
||||
}
|
||||
+18
-6
@@ -11,7 +11,7 @@ The service is configured using the environment variables from the following tab
|
||||
| COCOS_JAEGER_URL | The URL for the Jaeger tracing endpoint. | http://localhost:4318 |
|
||||
| COCOS_JAEGER_TRACE_RATIO | The ratio of traces to sample. | 1.0 |
|
||||
| MANAGER_INSTANCE_ID | The instance ID for the manager service. | |
|
||||
| MANAGER_ATTESTATION_POLICY_BINARY | The file path for the attestation policy and igvmmeassure binaries. | ../../build |
|
||||
| MANAGER_ATTESTATION_POLICY_BINARY | The file path for the attestation policy binary. | ../../build |
|
||||
| MANAGER_GRPC_CLIENT_CERT | The file path for the client certificate. | |
|
||||
| MANAGER_GRPC_CLIENT_KEY | The file path for the client private key. | |
|
||||
| MANAGER_GRPC_SERVER_CA_CERTS | The file path for the server CA certificate(s). | |
|
||||
@@ -44,10 +44,7 @@ The service is configured using the environment variables from the following tab
|
||||
| MANAGER_QEMU_SEV_ID | The ID for the Secure Encrypted Virtualization (SEV) device. | sev0 |
|
||||
| MANAGER_QEMU_SEV_CBITPOS | The position of the C-bit in the physical address. | 51 |
|
||||
| MANAGER_QEMU_SEV_REDUCED_PHYS_BITS | The number of reduced physical address bits for SEV. | 1 |
|
||||
| MANAGER_QEMU_ENABLE_HOST_DATA | Enable additional data for the SEV host. | false |
|
||||
| MANAGER_QEMU_HOST_DATA | Additional data for the SEV host. | |
|
||||
| MANAGER_QEMU_IGVM_ID | The ID of the IGVM file. | igvm0 |
|
||||
| MANAGER_QEMU_IGVM_FILE | The file path to the IGVM file. | /root/coconut-qemu.igvm |
|
||||
| MANAGER_QEMU_VSOCK_ID | The ID for the virtual socket device. | vhost-vsock-pci0 |
|
||||
| MANAGER_QEMU_VSOCK_GUEST_CID | The guest-side CID (Context ID) for the virtual socket device. | 3 |
|
||||
| MANAGER_QEMU_VSOCK_VNC | Whether to enable the virtual socket device for VNC. | 0 |
|
||||
@@ -61,6 +58,7 @@ The service is configured using the environment variables from the following tab
|
||||
| MANAGER_QEMU_SMP_COUNT | The number of virtual CPUs. | 4 |
|
||||
| MANAGER_QEMU_SMP_MAXCPUS | The maximum number of virtual CPUs. | 64 |
|
||||
| MANAGER_QEMU_MEM_ID | The ID for the memory device. | ram1 |
|
||||
| MANAGER_QEMU_KERNEL_HASH | Whether to enable kernel hash verification. | false |
|
||||
| MANAGER_QEMU_NO_GRAPHIC | Whether to disable the graphical display. | true |
|
||||
| MANAGER_QEMU_MONITOR | The type of monitor to use. | pty |
|
||||
| MANAGER_QEMU_HOST_FWD_RANGE | The range of host ports to forward. | 6100-6200 |
|
||||
@@ -193,7 +191,7 @@ NB: we set environment variables that we will use in the shell process where we
|
||||
|
||||
## Deployment
|
||||
|
||||
To start the service, execute the following shell script (note a server needs to be running see [here](../test/cvms/README.md)):
|
||||
To start the service, execute the following shell script (note a server needs to be running see [here](../test/computations/README.md)):
|
||||
|
||||
```bash
|
||||
# Download the latest version of the service
|
||||
@@ -234,7 +232,21 @@ MANAGER_QEMU_ENABLE_SEV=false \
|
||||
MANAGER_QEMU_ENABLE_SEV_SNP=true \
|
||||
MANAGER_QEMU_SEV_CBITPOS=51 \
|
||||
MANAGER_QEMU_BIN_PATH=<path to QEMU binary> \
|
||||
MANAGER_QEMU_IGVM_FILE=<path to IGVM file> \
|
||||
MANAGER_QEMU_QEMU_OVMF_CODE_FILE=<path to OVMF.fd Amd Sev built package> \
|
||||
./build/cocos-manager
|
||||
```
|
||||
|
||||
To include the kernel hash into the measurement of the attestation report (SEV or SEV-SNP), start manager like this
|
||||
|
||||
```sh
|
||||
MANAGER_GRPC_URL=localhost:7001 \
|
||||
MANAGER_LOG_LEVEL=debug \
|
||||
MANAGER_QEMU_ENABLE_SEV=false \
|
||||
MANAGER_QEMU_ENABLE_SEV_SNP=true \
|
||||
MANAGER_QEMU_SEV_CBITPOS=51 \
|
||||
MANAGER_QEMU_KERNEL_HASH=true \
|
||||
MANAGER_QEMU_BIN_PATH=<path to QEMU binary> \
|
||||
MANAGER_QEMU_QEMU_OVMF_CODE_FILE=<path to OVMF.fd Amd Sev built package> \
|
||||
./build/cocos-manager
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package manager
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strconv"
|
||||
|
||||
"github.com/ultravioletrs/cocos/pkg/manager"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
)
|
||||
|
||||
var (
|
||||
errFailedToParseCID = fmt.Errorf("failed to parse computation ID")
|
||||
errComputationNotFound = fmt.Errorf("computation not found")
|
||||
)
|
||||
|
||||
func (ms *managerService) computationIDFromAddress(address string) (string, error) {
|
||||
re := regexp.MustCompile(`vm\((\d+)\)`)
|
||||
matches := re.FindStringSubmatch(address)
|
||||
|
||||
if len(matches) > 1 {
|
||||
cid, err := strconv.Atoi(matches[1])
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return ms.findComputationID(cid)
|
||||
}
|
||||
return "", errFailedToParseCID
|
||||
}
|
||||
|
||||
func (ms *managerService) findComputationID(cid int) (string, error) {
|
||||
ms.mu.Lock()
|
||||
defer ms.mu.Unlock()
|
||||
for cmpID, vm := range ms.vms {
|
||||
if vm.GetCID() == cid {
|
||||
return cmpID, nil
|
||||
}
|
||||
}
|
||||
|
||||
return "", errComputationNotFound
|
||||
}
|
||||
|
||||
func (ms *managerService) reportBrokenConnection(cmpID string) {
|
||||
ms.eventsChan <- &ClientStreamMessage{
|
||||
Message: &ClientStreamMessage_AgentEvent{
|
||||
AgentEvent: &AgentEvent{
|
||||
EventType: ms.vms[cmpID].State(),
|
||||
ComputationId: cmpID,
|
||||
Status: manager.Disconnected.String(),
|
||||
Timestamp: timestamppb.Now(),
|
||||
Originator: "manager",
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (ms *managerService) ReportBrokenConnection(addr string) {
|
||||
cmpID, err := ms.computationIDFromAddress(addr)
|
||||
if err != nil {
|
||||
ms.logger.Warn(err.Error())
|
||||
return
|
||||
}
|
||||
ms.reportBrokenConnection(cmpID)
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package manager
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/ultravioletrs/cocos/manager/qemu"
|
||||
"github.com/ultravioletrs/cocos/manager/vm"
|
||||
"github.com/ultravioletrs/cocos/pkg/manager"
|
||||
)
|
||||
|
||||
func TestComputationIDFromAddress(t *testing.T) {
|
||||
ms := &managerService{
|
||||
vms: map[string]vm.VM{
|
||||
"comp1": qemu.NewVM(qemu.VMInfo{Config: qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 3}}}, func(event interface{}) error { return nil }, "comp1"),
|
||||
"comp2": qemu.NewVM(qemu.VMInfo{Config: qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 5}}}, func(event interface{}) error { return nil }, "comp2"),
|
||||
},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
address string
|
||||
want string
|
||||
wantErr bool
|
||||
}{
|
||||
{"Valid address", "vm(3)", "comp1", false},
|
||||
{"Invalid address", "invalid", "", true},
|
||||
{"Non-existent CID", "vm(10)", "", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := ms.computationIDFromAddress(tt.address)
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, tt.want, got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestReportBrokenConnection(t *testing.T) {
|
||||
ms := &managerService{
|
||||
eventsChan: make(chan *ClientStreamMessage, 1),
|
||||
vms: map[string]vm.VM{
|
||||
"comp1": qemu.NewVM(qemu.VMInfo{Config: qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 3}}}, func(event interface{}) error { return nil }, "comp1"),
|
||||
},
|
||||
}
|
||||
|
||||
ms.reportBrokenConnection("comp1")
|
||||
|
||||
select {
|
||||
case msg := <-ms.eventsChan:
|
||||
assert.Equal(t, "comp1", msg.GetAgentEvent().ComputationId)
|
||||
assert.Equal(t, manager.Disconnected.String(), msg.GetAgentEvent().Status)
|
||||
assert.Equal(t, "manager", msg.GetAgentEvent().Originator)
|
||||
default:
|
||||
t.Error("Expected message in eventsChan, but none received")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
// 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/manager"
|
||||
"github.com/ultravioletrs/cocos/manager/qemu"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
var (
|
||||
errTerminationFromServer = errors.New("server requested client termination")
|
||||
errCorruptedManifest = errors.New("received manifest may be corrupted")
|
||||
sendTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
type ManagerClient struct {
|
||||
stream manager.ManagerService_ProcessClient
|
||||
svc manager.Service
|
||||
messageQueue chan *manager.ClientStreamMessage
|
||||
logger *slog.Logger
|
||||
runReqManager *runRequestManager
|
||||
}
|
||||
|
||||
// NewClient returns new gRPC client instance.
|
||||
func NewClient(stream manager.ManagerService_ProcessClient, svc manager.Service, messageQueue chan *manager.ClientStreamMessage, logger *slog.Logger) ManagerClient {
|
||||
return ManagerClient{
|
||||
stream: stream,
|
||||
svc: svc,
|
||||
messageQueue: messageQueue,
|
||||
logger: logger,
|
||||
runReqManager: newRunRequestManager(),
|
||||
}
|
||||
}
|
||||
|
||||
func (client ManagerClient) Process(ctx context.Context, cancel context.CancelFunc) 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 ManagerClient) 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 ManagerClient) processIncomingMessage(ctx context.Context, req *manager.ServerStreamMessage) error {
|
||||
switch mes := req.Message.(type) {
|
||||
case *manager.ServerStreamMessage_RunReqChunks:
|
||||
return client.handleRunReqChunks(ctx, mes)
|
||||
case *manager.ServerStreamMessage_TerminateReq:
|
||||
return client.handleTerminateReq(mes)
|
||||
case *manager.ServerStreamMessage_StopComputation:
|
||||
go client.handleStopComputation(ctx, mes)
|
||||
case *manager.ServerStreamMessage_AttestationPolicyReq:
|
||||
go client.handleAttestationPolicyReq(ctx, mes)
|
||||
case *manager.ServerStreamMessage_SvmInfoReq:
|
||||
go client.handleSVMInfoReq(ctx, mes)
|
||||
default:
|
||||
return errors.New("unknown message type")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (client *ManagerClient) handleRunReqChunks(ctx context.Context, mes *manager.ServerStreamMessage_RunReqChunks) error {
|
||||
buffer, complete := client.runReqManager.addChunk(mes.RunReqChunks.Id, mes.RunReqChunks.Data, mes.RunReqChunks.IsLast)
|
||||
|
||||
if complete {
|
||||
var runReq manager.ComputationRunReq
|
||||
if err := proto.Unmarshal(buffer, &runReq); err != nil {
|
||||
return errors.Wrap(err, errCorruptedManifest)
|
||||
}
|
||||
|
||||
go client.executeRun(ctx, &runReq)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (client ManagerClient) executeRun(ctx context.Context, runReq *manager.ComputationRunReq) {
|
||||
port, err := client.svc.Run(ctx, runReq)
|
||||
if err != nil {
|
||||
client.logger.Warn(err.Error())
|
||||
return
|
||||
}
|
||||
runRes := &manager.ClientStreamMessage_RunRes{
|
||||
RunRes: &manager.RunResponse{
|
||||
AgentPort: port,
|
||||
ComputationId: runReq.Id,
|
||||
},
|
||||
}
|
||||
client.sendMessage(&manager.ClientStreamMessage{Message: runRes})
|
||||
}
|
||||
|
||||
func (client ManagerClient) handleTerminateReq(mes *manager.ServerStreamMessage_TerminateReq) error {
|
||||
return errors.Wrap(errTerminationFromServer, errors.New(mes.TerminateReq.Message))
|
||||
}
|
||||
|
||||
func (client ManagerClient) handleStopComputation(ctx context.Context, mes *manager.ServerStreamMessage_StopComputation) {
|
||||
msg := &manager.ClientStreamMessage_StopComputationRes{
|
||||
StopComputationRes: &manager.StopComputationResponse{
|
||||
ComputationId: mes.StopComputation.ComputationId,
|
||||
},
|
||||
}
|
||||
if err := client.svc.Stop(ctx, mes.StopComputation.ComputationId); err != nil {
|
||||
msg.StopComputationRes.Message = err.Error()
|
||||
}
|
||||
client.sendMessage(&manager.ClientStreamMessage{Message: msg})
|
||||
}
|
||||
|
||||
func (client ManagerClient) handleAttestationPolicyReq(ctx context.Context, mes *manager.ServerStreamMessage_AttestationPolicyReq) {
|
||||
res, err := client.svc.FetchAttestationPolicy(ctx, mes.AttestationPolicyReq.Id)
|
||||
if err != nil {
|
||||
client.logger.Warn(err.Error())
|
||||
return
|
||||
}
|
||||
info := &manager.ClientStreamMessage_AttestationPolicy{
|
||||
AttestationPolicy: &manager.AttestationPolicy{
|
||||
Info: res,
|
||||
Id: mes.AttestationPolicyReq.Id,
|
||||
},
|
||||
}
|
||||
client.sendMessage(&manager.ClientStreamMessage{Message: info})
|
||||
}
|
||||
|
||||
func (client ManagerClient) handleSVMInfoReq(ctx context.Context, mes *manager.ServerStreamMessage_SvmInfoReq) {
|
||||
ovmfVersion, cpuNum, cpuType, eosVersion := client.svc.ReturnSVMInfo(ctx)
|
||||
info := &manager.ClientStreamMessage_SvmInfo{
|
||||
SvmInfo: &manager.SVMInfo{
|
||||
OvmfVersion: ovmfVersion,
|
||||
CpuNum: int32(cpuNum),
|
||||
CpuType: cpuType,
|
||||
KernelCmd: qemu.KernelCommandLine,
|
||||
EosVersion: eosVersion,
|
||||
Id: mes.SvmInfoReq.Id,
|
||||
},
|
||||
}
|
||||
client.sendMessage(&manager.ClientStreamMessage{Message: info})
|
||||
}
|
||||
|
||||
func (client ManagerClient) handleOutgoingMessages(ctx context.Context) error {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case mes := <-client.messageQueue:
|
||||
if err := client.stream.Send(mes); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (client ManagerClient) sendMessage(mes *manager.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
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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/ultravioletrs/cocos/manager"
|
||||
"github.com/ultravioletrs/cocos/manager/mocks"
|
||||
"github.com/ultravioletrs/cocos/manager/qemu"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
type mockStream struct {
|
||||
mock.Mock
|
||||
grpc.ClientStream
|
||||
}
|
||||
|
||||
func (m *mockStream) Recv() (*manager.ServerStreamMessage, error) {
|
||||
args := m.Called()
|
||||
return args.Get(0).(*manager.ServerStreamMessage), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *mockStream) Send(msg *manager.ClientStreamMessage) error {
|
||||
args := m.Called(msg)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func TestManagerClient_Process1(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setupMocks func(mockStream *mockStream, mockSvc *mocks.Service)
|
||||
expectError bool
|
||||
errorMsg string
|
||||
}{
|
||||
{
|
||||
name: "Stop computation",
|
||||
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
|
||||
mockStream.On("Recv").Return(&manager.ServerStreamMessage{
|
||||
Message: &manager.ServerStreamMessage_StopComputation{
|
||||
StopComputation: &manager.StopComputation{},
|
||||
},
|
||||
}, nil)
|
||||
mockStream.On("Send", mock.Anything).Return(nil)
|
||||
mockSvc.On("Stop", mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
expectError: true,
|
||||
errorMsg: "context deadline exceeded",
|
||||
},
|
||||
{
|
||||
name: "Terminate request",
|
||||
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
|
||||
mockStream.On("Recv").Return(&manager.ServerStreamMessage{
|
||||
Message: &manager.ServerStreamMessage_TerminateReq{
|
||||
TerminateReq: &manager.Terminate{},
|
||||
},
|
||||
}, nil)
|
||||
},
|
||||
expectError: true,
|
||||
errorMsg: errTerminationFromServer.Error(),
|
||||
},
|
||||
{
|
||||
name: "Attestation Policy request",
|
||||
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
|
||||
mockStream.On("Recv").Return(&manager.ServerStreamMessage{
|
||||
Message: &manager.ServerStreamMessage_AttestationPolicyReq{
|
||||
AttestationPolicyReq: &manager.AttestationPolicyReq{},
|
||||
},
|
||||
}, nil)
|
||||
mockStream.On("Send", mock.Anything).Return(nil).Once()
|
||||
mockSvc.On("FetchAttestationPolicy", mock.Anything, mock.Anything).Return(nil, assert.AnError)
|
||||
},
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "Run request chunks",
|
||||
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
|
||||
mockStream.On("Recv").Return(&manager.ServerStreamMessage{
|
||||
Message: &manager.ServerStreamMessage_RunReqChunks{
|
||||
RunReqChunks: &manager.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) {
|
||||
mockStream.On("Recv").Return(&manager.ServerStreamMessage{}, assert.AnError)
|
||||
},
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
messageQueue := make(chan *manager.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
|
||||
client := NewClient(mockStream, mockSvc, messageQueue, logger)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
tc.setupMocks(mockStream, mockSvc)
|
||||
|
||||
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)
|
||||
messageQueue := make(chan *manager.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
|
||||
client := NewClient(mockStream, mockSvc, messageQueue, logger)
|
||||
|
||||
runReq := &manager.ComputationRunReq{
|
||||
Id: "test-id",
|
||||
}
|
||||
runReqBytes, _ := proto.Marshal(runReq)
|
||||
|
||||
chunk1 := &manager.ServerStreamMessage_RunReqChunks{
|
||||
RunReqChunks: &manager.RunReqChunks{
|
||||
Id: "chunk-1",
|
||||
Data: runReqBytes[:len(runReqBytes)/2],
|
||||
IsLast: false,
|
||||
},
|
||||
}
|
||||
chunk2 := &manager.ServerStreamMessage_RunReqChunks{
|
||||
RunReqChunks: &manager.RunReqChunks{
|
||||
Id: "chunk-1",
|
||||
Data: runReqBytes[len(runReqBytes)/2:],
|
||||
IsLast: true,
|
||||
},
|
||||
}
|
||||
|
||||
mockSvc.On("Run", mock.Anything, mock.AnythingOfType("*manager.ComputationRunReq")).Return("8080", 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.(*manager.ClientStreamMessage_RunRes)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "8080", runRes.RunRes.AgentPort)
|
||||
assert.Equal(t, "test-id", runRes.RunRes.ComputationId)
|
||||
}
|
||||
|
||||
func TestManagerClient_handleTerminateReq(t *testing.T) {
|
||||
client := ManagerClient{}
|
||||
|
||||
terminateReq := &manager.ServerStreamMessage_TerminateReq{
|
||||
TerminateReq: &manager.Terminate{
|
||||
Message: "Test termination",
|
||||
},
|
||||
}
|
||||
|
||||
err := client.handleTerminateReq(terminateReq)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "Test termination")
|
||||
assert.True(t, errors.Contains(err, errTerminationFromServer))
|
||||
}
|
||||
|
||||
func TestManagerClient_handleStopComputation(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
messageQueue := make(chan *manager.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
|
||||
client := NewClient(mockStream, mockSvc, messageQueue, logger)
|
||||
|
||||
stopReq := &manager.ServerStreamMessage_StopComputation{
|
||||
StopComputation: &manager.StopComputation{
|
||||
ComputationId: "test-comp-id",
|
||||
},
|
||||
}
|
||||
|
||||
mockSvc.On("Stop", mock.Anything, "test-comp-id").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.(*manager.ClientStreamMessage_StopComputationRes)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "test-comp-id", stopRes.StopComputationRes.ComputationId)
|
||||
assert.Empty(t, stopRes.StopComputationRes.Message)
|
||||
}
|
||||
|
||||
func TestManagerClient_handleAttestationPolicyReq(t *testing.T) {
|
||||
t.Run("success", func(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
messageQueue := make(chan *manager.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
|
||||
client := NewClient(mockStream, mockSvc, messageQueue, logger)
|
||||
|
||||
infoReq := &manager.ServerStreamMessage_AttestationPolicyReq{
|
||||
AttestationPolicyReq: &manager.AttestationPolicyReq{
|
||||
Id: "test-info-id",
|
||||
},
|
||||
}
|
||||
|
||||
mockSvc.On("FetchAttestationPolicy", context.Background(), infoReq.AttestationPolicyReq.Id).Return([]byte("test-attestation-policy"), nil)
|
||||
|
||||
client.handleAttestationPolicyReq(context.Background(), infoReq)
|
||||
|
||||
// Wait for the goroutine to finish
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
mockSvc.AssertExpectations(t)
|
||||
assert.Len(t, messageQueue, 1)
|
||||
|
||||
msg := <-messageQueue
|
||||
infoRes, ok := msg.Message.(*manager.ClientStreamMessage_AttestationPolicy)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "test-info-id", infoRes.AttestationPolicy.Id)
|
||||
assert.Equal(t, []byte("test-attestation-policy"), infoRes.AttestationPolicy.Info)
|
||||
})
|
||||
t.Run("error", func(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
messageQueue := make(chan *manager.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
|
||||
client := NewClient(mockStream, mockSvc, messageQueue, logger)
|
||||
|
||||
infoReq := &manager.ServerStreamMessage_AttestationPolicyReq{
|
||||
AttestationPolicyReq: &manager.AttestationPolicyReq{
|
||||
Id: "test-info-id",
|
||||
},
|
||||
}
|
||||
|
||||
mockSvc.On("FetchAttestationPolicy", context.Background(), infoReq.AttestationPolicyReq.Id).Return(nil, assert.AnError)
|
||||
|
||||
client.handleAttestationPolicyReq(context.Background(), infoReq)
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
mockSvc.AssertExpectations(t)
|
||||
assert.Len(t, messageQueue, 0)
|
||||
})
|
||||
}
|
||||
|
||||
func TestManagerClient_handleSVMInfoReq(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
messageQueue := make(chan *manager.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
|
||||
client := NewClient(mockStream, mockSvc, messageQueue, logger)
|
||||
|
||||
mockSvc.On("ReturnSVMInfo", context.Background()).Return("edk2-stable202408", 4, "EPYC", "")
|
||||
|
||||
client.handleSVMInfoReq(context.Background(), &manager.ServerStreamMessage_SvmInfoReq{SvmInfoReq: &manager.SVMInfoReq{Id: "test-svm-info-id"}})
|
||||
|
||||
// Wait for the goroutine to finish
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
mockSvc.AssertExpectations(t)
|
||||
assert.Len(t, messageQueue, 1)
|
||||
|
||||
msg := <-messageQueue
|
||||
infoRes, ok := msg.Message.(*manager.ClientStreamMessage_SvmInfo)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "edk2-stable202408", infoRes.SvmInfo.OvmfVersion)
|
||||
assert.Equal(t, int32(4), infoRes.SvmInfo.CpuNum)
|
||||
assert.Equal(t, "EPYC", infoRes.SvmInfo.CpuType)
|
||||
assert.Equal(t, "", infoRes.SvmInfo.EosVersion)
|
||||
assert.Equal(t, qemu.KernelCommandLine, infoRes.SvmInfo.KernelCmd)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
+101
-40
@@ -3,11 +3,17 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
"google.golang.org/protobuf/types/known/emptypb"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"google.golang.org/grpc/credentials"
|
||||
"google.golang.org/grpc/peer"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -15,58 +21,113 @@ var (
|
||||
ErrUnexpectedMsg = errors.New("unknown message type")
|
||||
)
|
||||
|
||||
const (
|
||||
bufferSize = 1024 * 1024 // 1 MB
|
||||
runReqTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
type SendFunc func(*manager.ServerStreamMessage) error
|
||||
|
||||
type grpcServer struct {
|
||||
manager.UnimplementedManagerServiceServer
|
||||
svc manager.Service
|
||||
incoming chan *manager.ClientStreamMessage
|
||||
svc Service
|
||||
}
|
||||
|
||||
type Service interface {
|
||||
Run(ctx context.Context, ipAddress string, sendMessage SendFunc, authInfo credentials.AuthInfo)
|
||||
}
|
||||
|
||||
// NewServer returns new AuthServiceServer instance.
|
||||
func NewServer(svc manager.Service) manager.ManagerServiceServer {
|
||||
func NewServer(incoming chan *manager.ClientStreamMessage, svc Service) manager.ManagerServiceServer {
|
||||
return &grpcServer{
|
||||
svc: svc,
|
||||
incoming: incoming,
|
||||
svc: svc,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *grpcServer) CreateVm(ctx context.Context, req *manager.CreateReq) (*manager.CreateRes, error) {
|
||||
port, id, err := s.svc.CreateVM(ctx, req)
|
||||
func (s *grpcServer) Process(stream manager.ManagerService_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 *manager.ServerStreamMessage) error {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
switch m := msg.Message.(type) {
|
||||
case *manager.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 manager.ManagerService_ProcessServer, runReq *manager.ComputationRunReq) error {
|
||||
data, err := proto.Marshal(runReq)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return err
|
||||
}
|
||||
|
||||
return &manager.CreateRes{
|
||||
ForwardedPort: port,
|
||||
SvmId: id,
|
||||
}, nil
|
||||
}
|
||||
dataBuffer := bytes.NewBuffer(data)
|
||||
buf := make([]byte, bufferSize)
|
||||
|
||||
func (s *grpcServer) RemoveVm(ctx context.Context, req *manager.RemoveReq) (*emptypb.Empty, error) {
|
||||
if err := s.svc.RemoveVM(ctx, req.SvmId); err != nil {
|
||||
return nil, err
|
||||
for {
|
||||
n, err := dataBuffer.Read(buf)
|
||||
isLast := false
|
||||
|
||||
if err == io.EOF {
|
||||
isLast = true
|
||||
} else if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
chunk := &manager.ServerStreamMessage{
|
||||
Message: &manager.ServerStreamMessage_RunReqChunks{
|
||||
RunReqChunks: &manager.RunReqChunks{
|
||||
Id: runReq.Id,
|
||||
Data: buf[:n],
|
||||
IsLast: isLast,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
if err := stream.Send(chunk); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if isLast {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return &emptypb.Empty{}, nil
|
||||
}
|
||||
|
||||
func (s *grpcServer) SVMInfo(ctx context.Context, req *manager.SVMInfoReq) (*manager.SVMInfoRes, error) {
|
||||
ovmf, cpunum, cputype, eosversion := s.svc.ReturnSVMInfo(ctx)
|
||||
|
||||
return &manager.SVMInfoRes{
|
||||
OvmfVersion: ovmf,
|
||||
CpuNum: int32(cpunum),
|
||||
CpuType: cputype,
|
||||
EosVersion: eosversion,
|
||||
Id: req.Id,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *grpcServer) AttestationPolicy(ctx context.Context, req *manager.AttestationPolicyReq) (*manager.AttestationPolicyRes, error) {
|
||||
policy, err := s.svc.FetchAttestationPolicy(ctx, req.Id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &manager.AttestationPolicyRes{
|
||||
Info: policy,
|
||||
Id: req.Id,
|
||||
}, nil
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -10,24 +10,24 @@ import (
|
||||
"github.com/absmach/magistrala/pkg/errors"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
"google.golang.org/grpc/credentials"
|
||||
"google.golang.org/grpc/peer"
|
||||
)
|
||||
|
||||
type mockServerStream struct {
|
||||
mock.Mock
|
||||
cvms.Service_ProcessServer
|
||||
manager.ManagerService_ProcessServer
|
||||
}
|
||||
|
||||
func (m *mockServerStream) Send(msg *cvms.ServerStreamMessage) error {
|
||||
func (m *mockServerStream) Send(msg *manager.ServerStreamMessage) error {
|
||||
args := m.Called(msg)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *mockServerStream) Recv() (*cvms.ClientStreamMessage, error) {
|
||||
func (m *mockServerStream) Recv() (*manager.ClientStreamMessage, error) {
|
||||
args := m.Called()
|
||||
return args.Get(0).(*cvms.ClientStreamMessage), args.Error(1)
|
||||
return args.Get(0).(*manager.ClientStreamMessage), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *mockServerStream) Context() context.Context {
|
||||
@@ -44,7 +44,7 @@ func (m *mockService) Run(ctx context.Context, ipAddress string, sendMessage Sen
|
||||
}
|
||||
|
||||
func TestNewServer(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage)
|
||||
incoming := make(chan *manager.ClientStreamMessage)
|
||||
mockSvc := new(mockService)
|
||||
|
||||
server := NewServer(incoming, mockSvc)
|
||||
@@ -56,19 +56,19 @@ func TestNewServer(t *testing.T) {
|
||||
func TestGrpcServer_Process(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
recvReturn *cvms.ClientStreamMessage
|
||||
recvReturn *manager.ClientStreamMessage
|
||||
recvError error
|
||||
expectedError string
|
||||
}{
|
||||
{
|
||||
name: "Process with context deadline exceeded",
|
||||
recvReturn: &cvms.ClientStreamMessage{},
|
||||
recvReturn: &manager.ClientStreamMessage{},
|
||||
recvError: nil,
|
||||
expectedError: "context deadline exceeded",
|
||||
},
|
||||
{
|
||||
name: "Process with Recv error",
|
||||
recvReturn: &cvms.ClientStreamMessage{},
|
||||
recvReturn: &manager.ClientStreamMessage{},
|
||||
recvError: errors.New("recv error"),
|
||||
expectedError: "recv error",
|
||||
},
|
||||
@@ -76,7 +76,7 @@ func TestGrpcServer_Process(t *testing.T) {
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage, 1)
|
||||
incoming := make(chan *manager.ClientStreamMessage, 1)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
@@ -111,13 +111,13 @@ func TestGrpcServer_Process(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage)
|
||||
incoming := make(chan *manager.ClientStreamMessage)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
mockStream := new(mockServerStream)
|
||||
|
||||
runReq := &cvms.ComputationRunReq{
|
||||
runReq := &manager.ComputationRunReq{
|
||||
Id: "test-id",
|
||||
}
|
||||
|
||||
@@ -125,10 +125,10 @@ func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
|
||||
for i := range largePayload {
|
||||
largePayload[i] = byte(i % 256)
|
||||
}
|
||||
runReq.Algorithm = &cvms.Algorithm{}
|
||||
runReq.Algorithm = &manager.Algorithm{}
|
||||
runReq.Algorithm.UserKey = largePayload
|
||||
|
||||
mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(nil).Times(4)
|
||||
mockStream.On("Send", mock.AnythingOfType("*manager.ServerStreamMessage")).Return(nil).Times(4)
|
||||
|
||||
err := server.sendRunReqInChunks(mockStream, runReq)
|
||||
|
||||
@@ -139,7 +139,7 @@ func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
|
||||
assert.Equal(t, 4, len(calls))
|
||||
|
||||
for i, call := range calls {
|
||||
msg := call.Arguments[0].(*cvms.ServerStreamMessage)
|
||||
msg := call.Arguments[0].(*manager.ServerStreamMessage)
|
||||
chunk := msg.GetRunReqChunks()
|
||||
|
||||
assert.NotNil(t, chunk)
|
||||
@@ -174,9 +174,9 @@ func TestGrpcServer_ProcessWithMockService(t *testing.T) {
|
||||
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 := &manager.ComputationRunReq{Id: "test-run-id"}
|
||||
err := sendFunc(&manager.ServerStreamMessage{
|
||||
Message: &manager.ServerStreamMessage_RunReq{
|
||||
RunReq: runReq,
|
||||
},
|
||||
})
|
||||
@@ -184,17 +184,34 @@ func TestGrpcServer_ProcessWithMockService(t *testing.T) {
|
||||
}).
|
||||
Return()
|
||||
|
||||
mockStream.On("Send", mock.MatchedBy(func(msg *cvms.ServerStreamMessage) bool {
|
||||
mockStream.On("Send", mock.MatchedBy(func(msg *manager.ServerStreamMessage) bool {
|
||||
chunks := msg.GetRunReqChunks()
|
||||
return chunks != nil && chunks.Id == "test-run-id"
|
||||
})).Return(nil)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Terminate 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)
|
||||
err := sendFunc(&manager.ServerStreamMessage{
|
||||
Message: &manager.ServerStreamMessage_TerminateReq{
|
||||
TerminateReq: &manager.Terminate{},
|
||||
},
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
}).Return()
|
||||
|
||||
mockStream.On("Send", mock.AnythingOfType("*manager.ServerStreamMessage")).Return(nil)
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage, 10)
|
||||
incoming := make(chan *manager.ClientStreamMessage, 10)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
@@ -214,7 +231,7 @@ func TestGrpcServer_ProcessWithMockService(t *testing.T) {
|
||||
})
|
||||
|
||||
mockStream.On("Context").Return(peerCtx)
|
||||
mockStream.On("Recv").Return(&cvms.ClientStreamMessage{}, nil).Maybe()
|
||||
mockStream.On("Recv").Return(&manager.ClientStreamMessage{}, nil).Maybe()
|
||||
|
||||
tt.setupMockFn(mockSvc, mockStream)
|
||||
|
||||
@@ -234,18 +251,18 @@ func TestGrpcServer_ProcessWithMockService(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestGrpcServer_sendRunReqInChunksError(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage)
|
||||
incoming := make(chan *manager.ClientStreamMessage)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
mockStream := new(mockServerStream)
|
||||
|
||||
runReq := &cvms.ComputationRunReq{
|
||||
runReq := &manager.ComputationRunReq{
|
||||
Id: "test-id",
|
||||
}
|
||||
|
||||
// Simulate an error when sending
|
||||
mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(errors.New("send error")).Once()
|
||||
mockStream.On("Send", mock.AnythingOfType("*manager.ServerStreamMessage")).Return(errors.New("send error")).Once()
|
||||
|
||||
err := server.sendRunReqInChunks(mockStream, runReq)
|
||||
|
||||
@@ -255,7 +272,7 @@ func TestGrpcServer_sendRunReqInChunksError(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestGrpcServer_ProcessMissingPeerInfo(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage)
|
||||
incoming := make(chan *manager.ClientStreamMessage)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
+10
-6
@@ -27,9 +27,9 @@ func LoggingMiddleware(svc manager.Service, logger *slog.Logger) manager.Service
|
||||
return &loggingMiddleware{logger, svc}
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) CreateVM(ctx context.Context, req *manager.CreateReq) (agentAddr string, id string, err error) {
|
||||
func (lm *loggingMiddleware) Run(ctx context.Context, mc *manager.ComputationRunReq) (agentAddr string, err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method CreateVM for id %s on port %s took %s to complete", id, agentAddr, time.Since(begin))
|
||||
message := fmt.Sprintf("Method Run for computation took %s to complete", time.Since(begin))
|
||||
if err != nil {
|
||||
lm.logger.Warn(fmt.Sprintf("%s with error: %s.", message, err))
|
||||
return
|
||||
@@ -37,12 +37,12 @@ func (lm *loggingMiddleware) CreateVM(ctx context.Context, req *manager.CreateRe
|
||||
lm.logger.Info(message)
|
||||
}(time.Now())
|
||||
|
||||
return lm.svc.CreateVM(ctx, req)
|
||||
return lm.svc.Run(ctx, mc)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) RemoveVM(ctx context.Context, id string) (err error) {
|
||||
func (lm *loggingMiddleware) Stop(ctx context.Context, computationID string) (err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method RemoveVM for vm %s took %s to complete", id, time.Since(begin))
|
||||
message := fmt.Sprintf("Method Stop for computation took %s to complete", time.Since(begin))
|
||||
if err != nil {
|
||||
lm.logger.Warn(fmt.Sprintf("%s with error: %s.", message, err))
|
||||
return
|
||||
@@ -50,7 +50,7 @@ func (lm *loggingMiddleware) RemoveVM(ctx context.Context, id string) (err error
|
||||
lm.logger.Info(message)
|
||||
}(time.Now())
|
||||
|
||||
return lm.svc.RemoveVM(ctx, id)
|
||||
return lm.svc.Stop(ctx, computationID)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) FetchAttestationPolicy(ctx context.Context, cmpId string) (body []byte, err error) {
|
||||
@@ -67,6 +67,10 @@ func (lm *loggingMiddleware) FetchAttestationPolicy(ctx context.Context, cmpId s
|
||||
return lm.svc.FetchAttestationPolicy(ctx, cmpId)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) ReportBrokenConnection(addr string) {
|
||||
lm.svc.ReportBrokenConnection(addr)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) ReturnSVMInfo(ctx context.Context) (string, int, string, string) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method ReturnSVMInfo for computation took %s to complete", time.Since(begin))
|
||||
|
||||
@@ -32,22 +32,22 @@ func MetricsMiddleware(svc manager.Service, counter metrics.Counter, latency met
|
||||
}
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) CreateVM(ctx context.Context, req *manager.CreateReq) (string, string, error) {
|
||||
func (ms *metricsMiddleware) Run(ctx context.Context, mc *manager.ComputationRunReq) (string, error) {
|
||||
defer func(begin time.Time) {
|
||||
ms.counter.With("method", "Run").Add(1)
|
||||
ms.latency.With("method", "Run").Observe(time.Since(begin).Seconds())
|
||||
}(time.Now())
|
||||
|
||||
return ms.svc.CreateVM(ctx, req)
|
||||
return ms.svc.Run(ctx, mc)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) RemoveVM(ctx context.Context, computationID string) error {
|
||||
func (ms *metricsMiddleware) Stop(ctx context.Context, computationID string) error {
|
||||
defer func(begin time.Time) {
|
||||
ms.counter.With("method", "Stop").Add(1)
|
||||
ms.latency.With("method", "Stop").Observe(time.Since(begin).Seconds())
|
||||
}(time.Now())
|
||||
|
||||
return ms.svc.RemoveVM(ctx, computationID)
|
||||
return ms.svc.Stop(ctx, computationID)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) FetchAttestationPolicy(ctx context.Context, cmpId string) ([]byte, error) {
|
||||
@@ -59,6 +59,10 @@ func (ms *metricsMiddleware) FetchAttestationPolicy(ctx context.Context, cmpId s
|
||||
return ms.svc.FetchAttestationPolicy(ctx, cmpId)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) ReportBrokenConnection(addr string) {
|
||||
ms.svc.ReportBrokenConnection(addr)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) ReturnSVMInfo(ctx context.Context) (string, int, string, string) {
|
||||
defer func(begin time.Time) {
|
||||
ms.counter.With("method", "ReturnSVMInfo").Add(1)
|
||||
|
||||
@@ -7,8 +7,6 @@
|
||||
package manager
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
@@ -18,7 +16,6 @@ import (
|
||||
|
||||
"github.com/google/go-sev-guest/proto/check"
|
||||
"github.com/ultravioletrs/cocos/manager/qemu"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/igvmmeasure"
|
||||
"github.com/virtee/sev-snp-measure-go/cpuid"
|
||||
"github.com/virtee/sev-snp-measure-go/guest"
|
||||
"github.com/virtee/sev-snp-measure-go/vmmtypes"
|
||||
@@ -70,33 +67,17 @@ func (ms *managerService) FetchAttestationPolicy(_ context.Context, computationI
|
||||
return nil, err
|
||||
}
|
||||
case vmi.Config.EnableSEVSNP:
|
||||
igvmMeasurementBinaryPath := fmt.Sprintf("%s/igvmmeasure", ms.attestationPolicyBinaryPath)
|
||||
|
||||
var stdoutBuffer bytes.Buffer
|
||||
var stderrBuffer bytes.Buffer
|
||||
|
||||
stdout := bufio.NewWriter(&stdoutBuffer)
|
||||
stderr := bufio.NewWriter(&stderrBuffer)
|
||||
|
||||
igvmMeasurement, err := igvmmeasure.NewIgvmMeasurement(igvmMeasurementBinaryPath, stderr, stdout)
|
||||
measurement, err = guest.CalcLaunchDigest(guest.SEV_SNP, vmi.Config.SMPCount, uint64(cpuid.CpuSigs[vmi.Config.CPU]), vmi.Config.OVMFCodeConfig.File, vmi.Config.KernelFile, vmi.Config.RootFsFile, strconv.Quote(qemu.KernelCommandLine), defGuestFeatures, "", vmmtypes.QEMU, false, "", 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
err = igvmMeasurement.Run(ms.qemuCfg.IGVMConfig.File)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
measurement = stdoutBuffer.Bytes()
|
||||
}
|
||||
|
||||
if measurement != nil {
|
||||
attestationPolicy.Policy.Measurement = measurement
|
||||
}
|
||||
|
||||
if vmi.Config.SevConfig.EnableHostData {
|
||||
hostData, err := base64.StdEncoding.DecodeString(vmi.Config.SevConfig.HostData)
|
||||
if vmi.Config.HostData != "" {
|
||||
hostData, err := base64.StdEncoding.DecodeString(vmi.Config.HostData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -57,10 +57,9 @@ func TestFetchAttestationPolicy(t *testing.T) {
|
||||
binaryBehavior: "success",
|
||||
vmConfig: qemu.VMInfo{
|
||||
Config: qemu.Config{
|
||||
EnableSEV: true,
|
||||
EnableSEVSNP: false,
|
||||
SMPCount: 2,
|
||||
CPU: "EPYC",
|
||||
EnableSEV: true,
|
||||
SMPCount: 2,
|
||||
CPU: "EPYC",
|
||||
OVMFCodeConfig: qemu.OVMFCodeConfig{
|
||||
File: "/path/to/OVMF_CODE.fd",
|
||||
},
|
||||
@@ -69,6 +68,23 @@ func TestFetchAttestationPolicy(t *testing.T) {
|
||||
},
|
||||
expectedError: "open /path/to/OVMF_CODE.fd: no such file or directory",
|
||||
},
|
||||
{
|
||||
name: "Valid SEV-SNP configuration",
|
||||
computationId: "sev-snp-computation",
|
||||
binaryBehavior: "success",
|
||||
vmConfig: qemu.VMInfo{
|
||||
Config: qemu.Config{
|
||||
EnableSEVSNP: true,
|
||||
SMPCount: 4,
|
||||
CPU: "EPYC-v2",
|
||||
OVMFCodeConfig: qemu.OVMFCodeConfig{
|
||||
File: "/path/to/OVMF_CODE_SNP.fd",
|
||||
},
|
||||
},
|
||||
LaunchTCB: 0,
|
||||
},
|
||||
expectedError: "open /path/to/OVMF_CODE_SNP.fd: no such file or director",
|
||||
},
|
||||
{
|
||||
name: "Invalid computation ID",
|
||||
computationId: "non-existent",
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package events
|
||||
|
||||
import "context"
|
||||
|
||||
type Listener interface {
|
||||
Listen(ctx context.Context)
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package events
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net"
|
||||
|
||||
"github.com/mdlayher/vsock"
|
||||
agentevents "github.com/ultravioletrs/cocos/agent/events"
|
||||
internalvsock "github.com/ultravioletrs/cocos/internal/vsock"
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
const ManagerVsockPort = 9997
|
||||
|
||||
type ReportBrokenConnectionFunc func(address string)
|
||||
|
||||
type events struct {
|
||||
reportBrokenConnection ReportBrokenConnectionFunc
|
||||
lis net.Listener
|
||||
logger *slog.Logger
|
||||
eventsChan chan *manager.ClientStreamMessage
|
||||
}
|
||||
|
||||
func New(logger *slog.Logger, reportBrokenConnection ReportBrokenConnectionFunc, eventsChan chan *manager.ClientStreamMessage) (Listener, error) {
|
||||
l, err := vsock.Listen(ManagerVsockPort, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &events{
|
||||
lis: l,
|
||||
reportBrokenConnection: reportBrokenConnection,
|
||||
logger: logger,
|
||||
eventsChan: eventsChan,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (e *events) Listen(ctx context.Context) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
e.logger.Info("Listener shutting down")
|
||||
return
|
||||
default:
|
||||
conn, err := e.lis.Accept()
|
||||
if err != nil {
|
||||
e.logger.Warn(err.Error())
|
||||
continue
|
||||
}
|
||||
|
||||
go e.handleConnection(conn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *events) handleConnection(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
|
||||
ackReader := internalvsock.NewAckReader(conn)
|
||||
|
||||
for {
|
||||
var message agentevents.EventsLogs
|
||||
data, err := ackReader.Read()
|
||||
if err != nil {
|
||||
go e.reportBrokenConnection(conn.RemoteAddr().String())
|
||||
e.logger.Warn(err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
if err := proto.Unmarshal(data, &message); err != nil {
|
||||
e.logger.Warn(err.Error())
|
||||
continue
|
||||
}
|
||||
|
||||
var mes manager.ClientStreamMessage
|
||||
|
||||
args := []any{}
|
||||
|
||||
switch message.Message.(type) {
|
||||
case *agentevents.EventsLogs_AgentEvent:
|
||||
args = append(args, slog.Group("agent-event",
|
||||
slog.String("event-type", message.GetAgentEvent().GetEventType()),
|
||||
slog.String("computation-id", message.GetAgentEvent().GetComputationId()),
|
||||
slog.String("status", message.GetAgentEvent().GetStatus()),
|
||||
slog.String("originator", message.GetAgentEvent().GetOriginator()),
|
||||
slog.String("timestamp", message.GetAgentEvent().GetTimestamp().String()),
|
||||
slog.String("details", string(message.GetAgentEvent().GetDetails()))))
|
||||
mes = manager.ClientStreamMessage{
|
||||
Message: &manager.ClientStreamMessage_AgentEvent{
|
||||
AgentEvent: &manager.AgentEvent{
|
||||
EventType: message.GetAgentEvent().GetEventType(),
|
||||
ComputationId: message.GetAgentEvent().GetComputationId(),
|
||||
Status: message.GetAgentEvent().GetStatus(),
|
||||
Originator: message.GetAgentEvent().GetOriginator(),
|
||||
Timestamp: message.GetAgentEvent().GetTimestamp(),
|
||||
Details: message.GetAgentEvent().GetDetails(),
|
||||
},
|
||||
},
|
||||
}
|
||||
case *agentevents.EventsLogs_AgentLog:
|
||||
args = append(args, slog.Group("agent-log",
|
||||
slog.String("computation-id", message.GetAgentLog().GetComputationId()),
|
||||
slog.String("level", message.GetAgentLog().GetLevel()),
|
||||
slog.String("timestamp", message.GetAgentLog().GetTimestamp().String()),
|
||||
slog.String("message", message.GetAgentLog().GetMessage())))
|
||||
mes = manager.ClientStreamMessage{
|
||||
Message: &manager.ClientStreamMessage_AgentLog{
|
||||
AgentLog: &manager.AgentLog{
|
||||
ComputationId: message.GetAgentLog().GetComputationId(),
|
||||
Level: message.GetAgentLog().GetLevel(),
|
||||
Timestamp: message.GetAgentLog().GetTimestamp(),
|
||||
Message: message.GetAgentLog().GetMessage(),
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
e.eventsChan <- &mes
|
||||
|
||||
e.logger.Info("", args...)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package events
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
mglog "github.com/absmach/magistrala/logger"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
)
|
||||
|
||||
type MockVsockListener struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockVsockListener) Accept() (net.Conn, error) {
|
||||
args := m.Called()
|
||||
return args.Get(0).(net.Conn), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockVsockListener) Close() error {
|
||||
args := m.Called()
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockVsockListener) Addr() net.Addr {
|
||||
args := m.Called()
|
||||
return args.Get(0).(net.Addr)
|
||||
}
|
||||
|
||||
var _ net.Conn = (*MockConn)(nil)
|
||||
|
||||
type MockConn struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockConn) Read(b []byte) (n int, err error) {
|
||||
args := m.Called(b)
|
||||
return args.Int(0), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockConn) Write(b []byte) (n int, err error) {
|
||||
args := m.Called(b)
|
||||
return args.Int(0), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockConn) Close() error {
|
||||
args := m.Called()
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockConn) LocalAddr() net.Addr {
|
||||
args := m.Called()
|
||||
return args.Get(0).(net.Addr)
|
||||
}
|
||||
|
||||
func (m *MockConn) RemoteAddr() net.Addr {
|
||||
args := m.Called()
|
||||
return args.Get(0).(net.Addr)
|
||||
}
|
||||
|
||||
func (m *MockConn) SetDeadline(t time.Time) error {
|
||||
args := m.Called(t)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockConn) SetReadDeadline(t time.Time) error {
|
||||
args := m.Called(t)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockConn) SetWriteDeadline(t time.Time) error {
|
||||
args := m.Called(t)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func TestNew(t *testing.T) {
|
||||
logger := &slog.Logger{}
|
||||
reportBrokenConnection := func(address string) {}
|
||||
eventsChan := make(chan *manager.ClientStreamMessage)
|
||||
|
||||
e, err := New(logger, reportBrokenConnection, eventsChan)
|
||||
|
||||
if vsockDeviceExists() {
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.NotNil(t, e)
|
||||
assert.IsType(t, &events{}, e)
|
||||
} else {
|
||||
assert.Error(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListen(t *testing.T) {
|
||||
mockListener := new(MockVsockListener)
|
||||
mockConn := new(MockConn)
|
||||
|
||||
e := &events{
|
||||
lis: mockListener,
|
||||
logger: mglog.NewMock(),
|
||||
}
|
||||
|
||||
mockListener.On("Accept").Return(mockConn, fmt.Errorf("mock error")).Once()
|
||||
mockListener.On("Accept").Return(mockConn, nil)
|
||||
mockConn.On("Close").Return(nil)
|
||||
mockConn.On("Read", mock.Anything).Return(0, nil)
|
||||
|
||||
go e.Listen(context.Background())
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
mockListener.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestListenContextDone(t *testing.T) {
|
||||
mockListener := new(MockVsockListener)
|
||||
mockConn := new(MockConn)
|
||||
|
||||
e := &events{
|
||||
lis: mockListener,
|
||||
logger: mglog.NewMock(),
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
mockListener.On("Accept").Return(mockConn, nil)
|
||||
|
||||
e.Listen(ctx)
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
|
||||
func vsockDeviceExists() bool {
|
||||
fs, err := os.Stat("/dev/vsock")
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
if fs.Mode()&os.ModeDevice == 0 {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
type MockConnWithBuffer struct {
|
||||
mock.Mock
|
||||
readBuf *bytes.Buffer
|
||||
writeBuf *bytes.Buffer
|
||||
}
|
||||
|
||||
func NewMockConnWithBuffer() *MockConnWithBuffer {
|
||||
return &MockConnWithBuffer{
|
||||
readBuf: new(bytes.Buffer),
|
||||
writeBuf: new(bytes.Buffer),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MockConnWithBuffer) Read(b []byte) (n int, err error) {
|
||||
return m.readBuf.Read(b)
|
||||
}
|
||||
|
||||
func (m *MockConnWithBuffer) Write(b []byte) (n int, err error) {
|
||||
return m.writeBuf.Write(b)
|
||||
}
|
||||
|
||||
func (m *MockConnWithBuffer) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *MockConnWithBuffer) LocalAddr() net.Addr {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *MockConnWithBuffer) RemoteAddr() net.Addr {
|
||||
return &net.IPAddr{IP: net.ParseIP("localhost")}
|
||||
}
|
||||
|
||||
func (m *MockConnWithBuffer) SetDeadline(t time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *MockConnWithBuffer) SetReadDeadline(t time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *MockConnWithBuffer) SetWriteDeadline(t time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestHandleConnection(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
message *manager.ClientStreamMessage
|
||||
}{
|
||||
{
|
||||
name: "handle agent event",
|
||||
message: &manager.ClientStreamMessage{
|
||||
Message: &manager.ClientStreamMessage_AgentEvent{
|
||||
AgentEvent: &manager.AgentEvent{
|
||||
EventType: "test_event",
|
||||
ComputationId: "test_computation",
|
||||
Status: "test_status",
|
||||
Originator: "test_originator",
|
||||
Timestamp: timestamppb.Now(),
|
||||
Details: []byte("test_details"),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "handle agent log",
|
||||
message: &manager.ClientStreamMessage{
|
||||
Message: &manager.ClientStreamMessage_AgentLog{
|
||||
AgentLog: &manager.AgentLog{
|
||||
ComputationId: "test_computation",
|
||||
Timestamp: timestamppb.Now(),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockConn := NewMockConnWithBuffer()
|
||||
eventsChan := make(chan *manager.ClientStreamMessage, 1)
|
||||
|
||||
e := &events{
|
||||
logger: mglog.NewMock(),
|
||||
eventsChan: eventsChan,
|
||||
reportBrokenConnection: func(address string) {},
|
||||
}
|
||||
|
||||
data, err := proto.Marshal(tt.message)
|
||||
assert.NoError(t, err)
|
||||
|
||||
messageID := uint32(1)
|
||||
err = binary.Write(mockConn.readBuf, binary.LittleEndian, messageID)
|
||||
assert.NoError(t, err)
|
||||
err = binary.Write(mockConn.readBuf, binary.LittleEndian, uint32(len(data)))
|
||||
assert.NoError(t, err)
|
||||
_, err = mockConn.readBuf.Write(data)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Add EOF to signal end of stream
|
||||
err = binary.Write(mockConn.readBuf, binary.LittleEndian, uint32(0))
|
||||
assert.NoError(t, err)
|
||||
err = binary.Write(mockConn.readBuf, binary.LittleEndian, uint32(0))
|
||||
assert.NoError(t, err)
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
e.handleConnection(mockConn)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
var receivedMessage *manager.ClientStreamMessage
|
||||
select {
|
||||
case receivedMessage = <-eventsChan:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Timeout waiting for message in eventsChan")
|
||||
}
|
||||
|
||||
assert.NotNil(t, receivedMessage)
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
// handleConnection has exited
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Timeout waiting for handleConnection to exit")
|
||||
}
|
||||
|
||||
// Check if ack was written
|
||||
var receivedAck uint32
|
||||
err = binary.Read(mockConn.writeBuf, binary.LittleEndian, &receivedAck)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, messageID, receivedAck)
|
||||
|
||||
// Ensure no unexpected calls were made on the mock
|
||||
mockConn.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
||||
+1521
-252
File diff suppressed because it is too large
Load Diff
+96
-18
@@ -3,42 +3,40 @@
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
import "google/protobuf/empty.proto";
|
||||
import "google/protobuf/timestamp.proto";
|
||||
|
||||
package manager;
|
||||
|
||||
option go_package = "./manager";
|
||||
|
||||
service ManagerService {
|
||||
rpc CreateVm(CreateReq) returns (CreateRes) {}
|
||||
rpc RemoveVm(RemoveReq) returns (google.protobuf.Empty) {}
|
||||
rpc SVMInfo(SVMInfoReq) returns (SVMInfoRes) {}
|
||||
rpc AttestationPolicy(AttestationPolicyReq) returns (AttestationPolicyRes) {}
|
||||
rpc Process(stream ClientStreamMessage) returns (stream ServerStreamMessage) {}
|
||||
}
|
||||
|
||||
message CreateReq{
|
||||
string agent_log_level = 1;
|
||||
bytes agent_cvm_server_ca_cert = 2;
|
||||
bytes agent_cvm_client_key = 3;
|
||||
bytes agent_cvm_client_cert = 4;
|
||||
string agent_cvm_server_url = 5;
|
||||
message Terminate {
|
||||
string message = 1;
|
||||
}
|
||||
|
||||
message CreateRes{
|
||||
string forwarded_port = 1;
|
||||
string svm_id = 2;
|
||||
message StopComputation {
|
||||
string computation_id = 1;
|
||||
}
|
||||
|
||||
message RemoveReq{
|
||||
string svm_id = 1;
|
||||
message StopComputationResponse {
|
||||
string computation_id = 1;
|
||||
string message = 2;
|
||||
}
|
||||
|
||||
message AttestationPolicyRes{
|
||||
message RunResponse{
|
||||
string agent_port = 1;
|
||||
string computation_id = 2;
|
||||
}
|
||||
|
||||
message AttestationPolicy{
|
||||
bytes info = 1;
|
||||
string id = 2;
|
||||
}
|
||||
|
||||
message SVMInfoRes{
|
||||
message SVMInfo{
|
||||
string id = 1;
|
||||
string ovmf_version = 2;
|
||||
int32 cpu_num = 3;
|
||||
@@ -47,6 +45,60 @@ message SVMInfoRes{
|
||||
string eos_version = 6;
|
||||
}
|
||||
|
||||
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;
|
||||
AttestationPolicy attestationPolicy = 4;
|
||||
StopComputationResponse stopComputationRes = 5;
|
||||
SVMInfo svm_info = 6;
|
||||
}
|
||||
}
|
||||
|
||||
message ServerStreamMessage {
|
||||
oneof message {
|
||||
RunReqChunks runReqChunks = 1;
|
||||
ComputationRunReq runReq = 2;
|
||||
Terminate terminateReq = 3;
|
||||
StopComputation stopComputation = 4;
|
||||
AttestationPolicyReq attestationPolicyReq = 5;
|
||||
SVMInfoReq svmInfoReq = 6;
|
||||
}
|
||||
}
|
||||
|
||||
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 AttestationPolicyReq {
|
||||
string id = 1;
|
||||
}
|
||||
@@ -55,3 +107,29 @@ message SVMInfoReq {
|
||||
string id = 1;
|
||||
}
|
||||
|
||||
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 host = 2;
|
||||
string cert_file = 3;
|
||||
string key_file = 4;
|
||||
string client_ca_file = 5;
|
||||
string server_ca_file = 6;
|
||||
string log_level = 7;
|
||||
bool attested_tls = 8;
|
||||
}
|
||||
|
||||
|
||||
+22
-143
@@ -4,7 +4,7 @@
|
||||
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
|
||||
// versions:
|
||||
// - protoc-gen-go-grpc v1.5.1
|
||||
// - protoc v5.29.3
|
||||
// - protoc v5.28.1
|
||||
// source: manager/manager.proto
|
||||
|
||||
package manager
|
||||
@@ -14,7 +14,6 @@ import (
|
||||
grpc "google.golang.org/grpc"
|
||||
codes "google.golang.org/grpc/codes"
|
||||
status "google.golang.org/grpc/status"
|
||||
emptypb "google.golang.org/protobuf/types/known/emptypb"
|
||||
)
|
||||
|
||||
// This is a compile-time assertion to ensure that this generated file
|
||||
@@ -23,20 +22,14 @@ import (
|
||||
const _ = grpc.SupportPackageIsVersion9
|
||||
|
||||
const (
|
||||
ManagerService_CreateVm_FullMethodName = "/manager.ManagerService/CreateVm"
|
||||
ManagerService_RemoveVm_FullMethodName = "/manager.ManagerService/RemoveVm"
|
||||
ManagerService_SVMInfo_FullMethodName = "/manager.ManagerService/SVMInfo"
|
||||
ManagerService_AttestationPolicy_FullMethodName = "/manager.ManagerService/AttestationPolicy"
|
||||
ManagerService_Process_FullMethodName = "/manager.ManagerService/Process"
|
||||
)
|
||||
|
||||
// ManagerServiceClient is the client API for ManagerService 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 ManagerServiceClient interface {
|
||||
CreateVm(ctx context.Context, in *CreateReq, opts ...grpc.CallOption) (*CreateRes, error)
|
||||
RemoveVm(ctx context.Context, in *RemoveReq, opts ...grpc.CallOption) (*emptypb.Empty, error)
|
||||
SVMInfo(ctx context.Context, in *SVMInfoReq, opts ...grpc.CallOption) (*SVMInfoRes, error)
|
||||
AttestationPolicy(ctx context.Context, in *AttestationPolicyReq, opts ...grpc.CallOption) (*AttestationPolicyRes, error)
|
||||
Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage], error)
|
||||
}
|
||||
|
||||
type managerServiceClient struct {
|
||||
@@ -47,54 +40,24 @@ func NewManagerServiceClient(cc grpc.ClientConnInterface) ManagerServiceClient {
|
||||
return &managerServiceClient{cc}
|
||||
}
|
||||
|
||||
func (c *managerServiceClient) CreateVm(ctx context.Context, in *CreateReq, opts ...grpc.CallOption) (*CreateRes, error) {
|
||||
func (c *managerServiceClient) Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage], error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(CreateRes)
|
||||
err := c.cc.Invoke(ctx, ManagerService_CreateVm_FullMethodName, in, out, cOpts...)
|
||||
stream, err := c.cc.NewStream(ctx, &ManagerService_ServiceDesc.Streams[0], ManagerService_Process_FullMethodName, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
x := &grpc.GenericClientStream[ClientStreamMessage, ServerStreamMessage]{ClientStream: stream}
|
||||
return x, nil
|
||||
}
|
||||
|
||||
func (c *managerServiceClient) RemoveVm(ctx context.Context, in *RemoveReq, opts ...grpc.CallOption) (*emptypb.Empty, error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(emptypb.Empty)
|
||||
err := c.cc.Invoke(ctx, ManagerService_RemoveVm_FullMethodName, in, out, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (c *managerServiceClient) SVMInfo(ctx context.Context, in *SVMInfoReq, opts ...grpc.CallOption) (*SVMInfoRes, error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(SVMInfoRes)
|
||||
err := c.cc.Invoke(ctx, ManagerService_SVMInfo_FullMethodName, in, out, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (c *managerServiceClient) AttestationPolicy(ctx context.Context, in *AttestationPolicyReq, opts ...grpc.CallOption) (*AttestationPolicyRes, error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(AttestationPolicyRes)
|
||||
err := c.cc.Invoke(ctx, ManagerService_AttestationPolicy_FullMethodName, in, out, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
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 ManagerService_ProcessClient = grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage]
|
||||
|
||||
// ManagerServiceServer is the server API for ManagerService service.
|
||||
// All implementations must embed UnimplementedManagerServiceServer
|
||||
// for forward compatibility.
|
||||
type ManagerServiceServer interface {
|
||||
CreateVm(context.Context, *CreateReq) (*CreateRes, error)
|
||||
RemoveVm(context.Context, *RemoveReq) (*emptypb.Empty, error)
|
||||
SVMInfo(context.Context, *SVMInfoReq) (*SVMInfoRes, error)
|
||||
AttestationPolicy(context.Context, *AttestationPolicyReq) (*AttestationPolicyRes, error)
|
||||
Process(grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]) error
|
||||
mustEmbedUnimplementedManagerServiceServer()
|
||||
}
|
||||
|
||||
@@ -105,17 +68,8 @@ type ManagerServiceServer interface {
|
||||
// pointer dereference when methods are called.
|
||||
type UnimplementedManagerServiceServer struct{}
|
||||
|
||||
func (UnimplementedManagerServiceServer) CreateVm(context.Context, *CreateReq) (*CreateRes, error) {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method CreateVm not implemented")
|
||||
}
|
||||
func (UnimplementedManagerServiceServer) RemoveVm(context.Context, *RemoveReq) (*emptypb.Empty, error) {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method RemoveVm not implemented")
|
||||
}
|
||||
func (UnimplementedManagerServiceServer) SVMInfo(context.Context, *SVMInfoReq) (*SVMInfoRes, error) {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method SVMInfo not implemented")
|
||||
}
|
||||
func (UnimplementedManagerServiceServer) AttestationPolicy(context.Context, *AttestationPolicyReq) (*AttestationPolicyRes, error) {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method AttestationPolicy not implemented")
|
||||
func (UnimplementedManagerServiceServer) Process(grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]) error {
|
||||
return status.Errorf(codes.Unimplemented, "method Process not implemented")
|
||||
}
|
||||
func (UnimplementedManagerServiceServer) mustEmbedUnimplementedManagerServiceServer() {}
|
||||
func (UnimplementedManagerServiceServer) testEmbeddedByValue() {}
|
||||
@@ -138,77 +92,12 @@ func RegisterManagerServiceServer(s grpc.ServiceRegistrar, srv ManagerServiceSer
|
||||
s.RegisterService(&ManagerService_ServiceDesc, srv)
|
||||
}
|
||||
|
||||
func _ManagerService_CreateVm_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(CreateReq)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(ManagerServiceServer).CreateVm(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: ManagerService_CreateVm_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(ManagerServiceServer).CreateVm(ctx, req.(*CreateReq))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
func _ManagerService_Process_Handler(srv interface{}, stream grpc.ServerStream) error {
|
||||
return srv.(ManagerServiceServer).Process(&grpc.GenericServerStream[ClientStreamMessage, ServerStreamMessage]{ServerStream: stream})
|
||||
}
|
||||
|
||||
func _ManagerService_RemoveVm_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(RemoveReq)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(ManagerServiceServer).RemoveVm(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: ManagerService_RemoveVm_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(ManagerServiceServer).RemoveVm(ctx, req.(*RemoveReq))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
func _ManagerService_SVMInfo_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(SVMInfoReq)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(ManagerServiceServer).SVMInfo(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: ManagerService_SVMInfo_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(ManagerServiceServer).SVMInfo(ctx, req.(*SVMInfoReq))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
func _ManagerService_AttestationPolicy_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(AttestationPolicyReq)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(ManagerServiceServer).AttestationPolicy(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: ManagerService_AttestationPolicy_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(ManagerServiceServer).AttestationPolicy(ctx, req.(*AttestationPolicyReq))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type ManagerService_ProcessServer = grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]
|
||||
|
||||
// ManagerService_ServiceDesc is the grpc.ServiceDesc for ManagerService service.
|
||||
// It's only intended for direct use with grpc.RegisterService,
|
||||
@@ -216,24 +105,14 @@ func _ManagerService_AttestationPolicy_Handler(srv interface{}, ctx context.Cont
|
||||
var ManagerService_ServiceDesc = grpc.ServiceDesc{
|
||||
ServiceName: "manager.ManagerService",
|
||||
HandlerType: (*ManagerServiceServer)(nil),
|
||||
Methods: []grpc.MethodDesc{
|
||||
Methods: []grpc.MethodDesc{},
|
||||
Streams: []grpc.StreamDesc{
|
||||
{
|
||||
MethodName: "CreateVm",
|
||||
Handler: _ManagerService_CreateVm_Handler,
|
||||
},
|
||||
{
|
||||
MethodName: "RemoveVm",
|
||||
Handler: _ManagerService_RemoveVm_Handler,
|
||||
},
|
||||
{
|
||||
MethodName: "SVMInfo",
|
||||
Handler: _ManagerService_SVMInfo_Handler,
|
||||
},
|
||||
{
|
||||
MethodName: "AttestationPolicy",
|
||||
Handler: _ManagerService_AttestationPolicy_Handler,
|
||||
StreamName: "Process",
|
||||
Handler: _ManagerService_Process_Handler,
|
||||
ServerStreams: true,
|
||||
ClientStreams: true,
|
||||
},
|
||||
},
|
||||
Streams: []grpc.StreamDesc{},
|
||||
Metadata: "manager/manager.proto",
|
||||
}
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package manager_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
func TestProcess(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
conn, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to dial bufnet: %v", err)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
client := manager.NewManagerServiceClient(conn)
|
||||
stream, err := client.Process(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("Process failed: %v", err)
|
||||
}
|
||||
|
||||
var data bytes.Buffer
|
||||
for {
|
||||
msg, err := stream.Recv()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to receive ServerStreamMessage: %v", err)
|
||||
}
|
||||
|
||||
switch m := msg.Message.(type) {
|
||||
case *manager.ServerStreamMessage_TerminateReq:
|
||||
if m.TerminateReq.Message != "test terminate" {
|
||||
t.Fatalf("Unexpected terminate message: %v", m.TerminateReq.Message)
|
||||
}
|
||||
case *manager.ServerStreamMessage_RunReqChunks:
|
||||
if len(m.RunReqChunks.Data) == 0 {
|
||||
var runReq manager.ComputationRunReq
|
||||
if err = proto.Unmarshal(data.Bytes(), &runReq); err != nil {
|
||||
t.Fatalf("Failed to create run request: %v", err)
|
||||
}
|
||||
|
||||
runRes := &manager.ClientStreamMessage_AgentLog{
|
||||
AgentLog: &manager.AgentLog{
|
||||
Message: "test log",
|
||||
ComputationId: "comp1",
|
||||
Level: "DEBUG",
|
||||
},
|
||||
}
|
||||
if runReq.Id != "1" || runReq.Name != "sample computation" || runReq.Description != "sample description" {
|
||||
t.Fatalf("Unexpected run request message: %v", &runReq)
|
||||
}
|
||||
if err := stream.Send(&manager.ClientStreamMessage{Message: runRes}); err != nil {
|
||||
t.Fatalf("Failed to send ClientStreamMessage: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
data.Write(m.RunReqChunks.Data)
|
||||
default:
|
||||
t.Fatalf("Unexpected message type: %T", m)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,340 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery v2.43.2. DO NOT EDIT.
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
context "context"
|
||||
|
||||
grpc "google.golang.org/grpc"
|
||||
emptypb "google.golang.org/protobuf/types/known/emptypb"
|
||||
|
||||
manager "github.com/ultravioletrs/cocos/manager"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
// ManagerServiceClient is an autogenerated mock type for the ManagerServiceClient type
|
||||
type ManagerServiceClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type ManagerServiceClient_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *ManagerServiceClient) EXPECT() *ManagerServiceClient_Expecter {
|
||||
return &ManagerServiceClient_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// AttestationPolicy provides a mock function with given fields: ctx, in, opts
|
||||
func (_m *ManagerServiceClient) AttestationPolicy(ctx context.Context, in *manager.AttestationPolicyReq, opts ...grpc.CallOption) (*manager.AttestationPolicyRes, error) {
|
||||
_va := make([]interface{}, len(opts))
|
||||
for _i := range opts {
|
||||
_va[_i] = opts[_i]
|
||||
}
|
||||
var _ca []interface{}
|
||||
_ca = append(_ca, ctx, in)
|
||||
_ca = append(_ca, _va...)
|
||||
ret := _m.Called(_ca...)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for AttestationPolicy")
|
||||
}
|
||||
|
||||
var r0 *manager.AttestationPolicyRes
|
||||
var r1 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) (*manager.AttestationPolicyRes, error)); ok {
|
||||
return rf(ctx, in, opts...)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) *manager.AttestationPolicyRes); ok {
|
||||
r0 = rf(ctx, in, opts...)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*manager.AttestationPolicyRes)
|
||||
}
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) error); ok {
|
||||
r1 = rf(ctx, in, opts...)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// ManagerServiceClient_AttestationPolicy_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AttestationPolicy'
|
||||
type ManagerServiceClient_AttestationPolicy_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// AttestationPolicy is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - in *manager.AttestationPolicyReq
|
||||
// - opts ...grpc.CallOption
|
||||
func (_e *ManagerServiceClient_Expecter) AttestationPolicy(ctx interface{}, in interface{}, opts ...interface{}) *ManagerServiceClient_AttestationPolicy_Call {
|
||||
return &ManagerServiceClient_AttestationPolicy_Call{Call: _e.mock.On("AttestationPolicy",
|
||||
append([]interface{}{ctx, in}, opts...)...)}
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_AttestationPolicy_Call) Run(run func(ctx context.Context, in *manager.AttestationPolicyReq, opts ...grpc.CallOption)) *ManagerServiceClient_AttestationPolicy_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
variadicArgs := make([]grpc.CallOption, len(args)-2)
|
||||
for i, a := range args[2:] {
|
||||
if a != nil {
|
||||
variadicArgs[i] = a.(grpc.CallOption)
|
||||
}
|
||||
}
|
||||
run(args[0].(context.Context), args[1].(*manager.AttestationPolicyReq), variadicArgs...)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_AttestationPolicy_Call) Return(_a0 *manager.AttestationPolicyRes, _a1 error) *ManagerServiceClient_AttestationPolicy_Call {
|
||||
_c.Call.Return(_a0, _a1)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_AttestationPolicy_Call) RunAndReturn(run func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) (*manager.AttestationPolicyRes, error)) *ManagerServiceClient_AttestationPolicy_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// CreateVm provides a mock function with given fields: ctx, in, opts
|
||||
func (_m *ManagerServiceClient) CreateVm(ctx context.Context, in *manager.CreateReq, opts ...grpc.CallOption) (*manager.CreateRes, error) {
|
||||
_va := make([]interface{}, len(opts))
|
||||
for _i := range opts {
|
||||
_va[_i] = opts[_i]
|
||||
}
|
||||
var _ca []interface{}
|
||||
_ca = append(_ca, ctx, in)
|
||||
_ca = append(_ca, _va...)
|
||||
ret := _m.Called(_ca...)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for CreateVm")
|
||||
}
|
||||
|
||||
var r0 *manager.CreateRes
|
||||
var r1 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) (*manager.CreateRes, error)); ok {
|
||||
return rf(ctx, in, opts...)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) *manager.CreateRes); ok {
|
||||
r0 = rf(ctx, in, opts...)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*manager.CreateRes)
|
||||
}
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) error); ok {
|
||||
r1 = rf(ctx, in, opts...)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// ManagerServiceClient_CreateVm_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateVm'
|
||||
type ManagerServiceClient_CreateVm_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// CreateVm is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - in *manager.CreateReq
|
||||
// - opts ...grpc.CallOption
|
||||
func (_e *ManagerServiceClient_Expecter) CreateVm(ctx interface{}, in interface{}, opts ...interface{}) *ManagerServiceClient_CreateVm_Call {
|
||||
return &ManagerServiceClient_CreateVm_Call{Call: _e.mock.On("CreateVm",
|
||||
append([]interface{}{ctx, in}, opts...)...)}
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_CreateVm_Call) Run(run func(ctx context.Context, in *manager.CreateReq, opts ...grpc.CallOption)) *ManagerServiceClient_CreateVm_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
variadicArgs := make([]grpc.CallOption, len(args)-2)
|
||||
for i, a := range args[2:] {
|
||||
if a != nil {
|
||||
variadicArgs[i] = a.(grpc.CallOption)
|
||||
}
|
||||
}
|
||||
run(args[0].(context.Context), args[1].(*manager.CreateReq), variadicArgs...)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_CreateVm_Call) Return(_a0 *manager.CreateRes, _a1 error) *ManagerServiceClient_CreateVm_Call {
|
||||
_c.Call.Return(_a0, _a1)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_CreateVm_Call) RunAndReturn(run func(context.Context, *manager.CreateReq, ...grpc.CallOption) (*manager.CreateRes, error)) *ManagerServiceClient_CreateVm_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// RemoveVm provides a mock function with given fields: ctx, in, opts
|
||||
func (_m *ManagerServiceClient) RemoveVm(ctx context.Context, in *manager.RemoveReq, opts ...grpc.CallOption) (*emptypb.Empty, error) {
|
||||
_va := make([]interface{}, len(opts))
|
||||
for _i := range opts {
|
||||
_va[_i] = opts[_i]
|
||||
}
|
||||
var _ca []interface{}
|
||||
_ca = append(_ca, ctx, in)
|
||||
_ca = append(_ca, _va...)
|
||||
ret := _m.Called(_ca...)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for RemoveVm")
|
||||
}
|
||||
|
||||
var r0 *emptypb.Empty
|
||||
var r1 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) (*emptypb.Empty, error)); ok {
|
||||
return rf(ctx, in, opts...)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) *emptypb.Empty); ok {
|
||||
r0 = rf(ctx, in, opts...)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*emptypb.Empty)
|
||||
}
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) error); ok {
|
||||
r1 = rf(ctx, in, opts...)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// ManagerServiceClient_RemoveVm_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveVm'
|
||||
type ManagerServiceClient_RemoveVm_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// RemoveVm is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - in *manager.RemoveReq
|
||||
// - opts ...grpc.CallOption
|
||||
func (_e *ManagerServiceClient_Expecter) RemoveVm(ctx interface{}, in interface{}, opts ...interface{}) *ManagerServiceClient_RemoveVm_Call {
|
||||
return &ManagerServiceClient_RemoveVm_Call{Call: _e.mock.On("RemoveVm",
|
||||
append([]interface{}{ctx, in}, opts...)...)}
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_RemoveVm_Call) Run(run func(ctx context.Context, in *manager.RemoveReq, opts ...grpc.CallOption)) *ManagerServiceClient_RemoveVm_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
variadicArgs := make([]grpc.CallOption, len(args)-2)
|
||||
for i, a := range args[2:] {
|
||||
if a != nil {
|
||||
variadicArgs[i] = a.(grpc.CallOption)
|
||||
}
|
||||
}
|
||||
run(args[0].(context.Context), args[1].(*manager.RemoveReq), variadicArgs...)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_RemoveVm_Call) Return(_a0 *emptypb.Empty, _a1 error) *ManagerServiceClient_RemoveVm_Call {
|
||||
_c.Call.Return(_a0, _a1)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_RemoveVm_Call) RunAndReturn(run func(context.Context, *manager.RemoveReq, ...grpc.CallOption) (*emptypb.Empty, error)) *ManagerServiceClient_RemoveVm_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SVMInfo provides a mock function with given fields: ctx, in, opts
|
||||
func (_m *ManagerServiceClient) SVMInfo(ctx context.Context, in *manager.SVMInfoReq, opts ...grpc.CallOption) (*manager.SVMInfoRes, error) {
|
||||
_va := make([]interface{}, len(opts))
|
||||
for _i := range opts {
|
||||
_va[_i] = opts[_i]
|
||||
}
|
||||
var _ca []interface{}
|
||||
_ca = append(_ca, ctx, in)
|
||||
_ca = append(_ca, _va...)
|
||||
ret := _m.Called(_ca...)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for SVMInfo")
|
||||
}
|
||||
|
||||
var r0 *manager.SVMInfoRes
|
||||
var r1 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.SVMInfoReq, ...grpc.CallOption) (*manager.SVMInfoRes, error)); ok {
|
||||
return rf(ctx, in, opts...)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.SVMInfoReq, ...grpc.CallOption) *manager.SVMInfoRes); ok {
|
||||
r0 = rf(ctx, in, opts...)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*manager.SVMInfoRes)
|
||||
}
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, *manager.SVMInfoReq, ...grpc.CallOption) error); ok {
|
||||
r1 = rf(ctx, in, opts...)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// ManagerServiceClient_SVMInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SVMInfo'
|
||||
type ManagerServiceClient_SVMInfo_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// SVMInfo is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - in *manager.SVMInfoReq
|
||||
// - opts ...grpc.CallOption
|
||||
func (_e *ManagerServiceClient_Expecter) SVMInfo(ctx interface{}, in interface{}, opts ...interface{}) *ManagerServiceClient_SVMInfo_Call {
|
||||
return &ManagerServiceClient_SVMInfo_Call{Call: _e.mock.On("SVMInfo",
|
||||
append([]interface{}{ctx, in}, opts...)...)}
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_SVMInfo_Call) Run(run func(ctx context.Context, in *manager.SVMInfoReq, opts ...grpc.CallOption)) *ManagerServiceClient_SVMInfo_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
variadicArgs := make([]grpc.CallOption, len(args)-2)
|
||||
for i, a := range args[2:] {
|
||||
if a != nil {
|
||||
variadicArgs[i] = a.(grpc.CallOption)
|
||||
}
|
||||
}
|
||||
run(args[0].(context.Context), args[1].(*manager.SVMInfoReq), variadicArgs...)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_SVMInfo_Call) Return(_a0 *manager.SVMInfoRes, _a1 error) *ManagerServiceClient_SVMInfo_Call {
|
||||
_c.Call.Return(_a0, _a1)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *ManagerServiceClient_SVMInfo_Call) RunAndReturn(run func(context.Context, *manager.SVMInfoReq, ...grpc.CallOption) (*manager.SVMInfoRes, error)) *ManagerServiceClient_SVMInfo_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// NewManagerServiceClient creates a new instance of ManagerServiceClient. 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 NewManagerServiceClient(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *ManagerServiceClient {
|
||||
mock := &ManagerServiceClient{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
+118
-92
@@ -25,70 +25,6 @@ func (_m *Service) EXPECT() *Service_Expecter {
|
||||
return &Service_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// CreateVM provides a mock function with given fields: ctx, req
|
||||
func (_m *Service) CreateVM(ctx context.Context, req *manager.CreateReq) (string, string, error) {
|
||||
ret := _m.Called(ctx, req)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for CreateVM")
|
||||
}
|
||||
|
||||
var r0 string
|
||||
var r1 string
|
||||
var r2 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq) (string, string, error)); ok {
|
||||
return rf(ctx, req)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq) string); ok {
|
||||
r0 = rf(ctx, req)
|
||||
} else {
|
||||
r0 = ret.Get(0).(string)
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, *manager.CreateReq) string); ok {
|
||||
r1 = rf(ctx, req)
|
||||
} else {
|
||||
r1 = ret.Get(1).(string)
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(2).(func(context.Context, *manager.CreateReq) error); ok {
|
||||
r2 = rf(ctx, req)
|
||||
} else {
|
||||
r2 = ret.Error(2)
|
||||
}
|
||||
|
||||
return r0, r1, r2
|
||||
}
|
||||
|
||||
// Service_CreateVM_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateVM'
|
||||
type Service_CreateVM_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// CreateVM is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - req *manager.CreateReq
|
||||
func (_e *Service_Expecter) CreateVM(ctx interface{}, req interface{}) *Service_CreateVM_Call {
|
||||
return &Service_CreateVM_Call{Call: _e.mock.On("CreateVM", ctx, req)}
|
||||
}
|
||||
|
||||
func (_c *Service_CreateVM_Call) Run(run func(ctx context.Context, req *manager.CreateReq)) *Service_CreateVM_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run(args[0].(context.Context), args[1].(*manager.CreateReq))
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_CreateVM_Call) Return(_a0 string, _a1 string, _a2 error) *Service_CreateVM_Call {
|
||||
_c.Call.Return(_a0, _a1, _a2)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_CreateVM_Call) RunAndReturn(run func(context.Context, *manager.CreateReq) (string, string, error)) *Service_CreateVM_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// FetchAttestationPolicy provides a mock function with given fields: ctx, computationID
|
||||
func (_m *Service) FetchAttestationPolicy(ctx context.Context, computationID string) ([]byte, error) {
|
||||
ret := _m.Called(ctx, computationID)
|
||||
@@ -148,49 +84,35 @@ func (_c *Service_FetchAttestationPolicy_Call) RunAndReturn(run func(context.Con
|
||||
return _c
|
||||
}
|
||||
|
||||
// RemoveVM provides a mock function with given fields: ctx, computationID
|
||||
func (_m *Service) RemoveVM(ctx context.Context, computationID string) error {
|
||||
ret := _m.Called(ctx, computationID)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for RemoveVM")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
|
||||
r0 = rf(ctx, computationID)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
|
||||
return r0
|
||||
// ReportBrokenConnection provides a mock function with given fields: addr
|
||||
func (_m *Service) ReportBrokenConnection(addr string) {
|
||||
_m.Called(addr)
|
||||
}
|
||||
|
||||
// Service_RemoveVM_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveVM'
|
||||
type Service_RemoveVM_Call struct {
|
||||
// Service_ReportBrokenConnection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ReportBrokenConnection'
|
||||
type Service_ReportBrokenConnection_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// RemoveVM is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - computationID string
|
||||
func (_e *Service_Expecter) RemoveVM(ctx interface{}, computationID interface{}) *Service_RemoveVM_Call {
|
||||
return &Service_RemoveVM_Call{Call: _e.mock.On("RemoveVM", ctx, computationID)}
|
||||
// ReportBrokenConnection is a helper method to define mock.On call
|
||||
// - addr string
|
||||
func (_e *Service_Expecter) ReportBrokenConnection(addr interface{}) *Service_ReportBrokenConnection_Call {
|
||||
return &Service_ReportBrokenConnection_Call{Call: _e.mock.On("ReportBrokenConnection", addr)}
|
||||
}
|
||||
|
||||
func (_c *Service_RemoveVM_Call) Run(run func(ctx context.Context, computationID string)) *Service_RemoveVM_Call {
|
||||
func (_c *Service_ReportBrokenConnection_Call) Run(run func(addr string)) *Service_ReportBrokenConnection_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run(args[0].(context.Context), args[1].(string))
|
||||
run(args[0].(string))
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_RemoveVM_Call) Return(_a0 error) *Service_RemoveVM_Call {
|
||||
_c.Call.Return(_a0)
|
||||
func (_c *Service_ReportBrokenConnection_Call) Return() *Service_ReportBrokenConnection_Call {
|
||||
_c.Call.Return()
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_RemoveVM_Call) RunAndReturn(run func(context.Context, string) error) *Service_RemoveVM_Call {
|
||||
func (_c *Service_ReportBrokenConnection_Call) RunAndReturn(run func(string)) *Service_ReportBrokenConnection_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -265,6 +187,110 @@ func (_c *Service_ReturnSVMInfo_Call) RunAndReturn(run func(context.Context) (st
|
||||
return _c
|
||||
}
|
||||
|
||||
// Run provides a mock function with given fields: ctx, c
|
||||
func (_m *Service) Run(ctx context.Context, c *manager.ComputationRunReq) (string, error) {
|
||||
ret := _m.Called(ctx, c)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Run")
|
||||
}
|
||||
|
||||
var r0 string
|
||||
var r1 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.ComputationRunReq) (string, error)); ok {
|
||||
return rf(ctx, c)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, *manager.ComputationRunReq) string); ok {
|
||||
r0 = rf(ctx, c)
|
||||
} else {
|
||||
r0 = ret.Get(0).(string)
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, *manager.ComputationRunReq) error); ok {
|
||||
r1 = rf(ctx, c)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// Service_Run_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Run'
|
||||
type Service_Run_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Run is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - c *manager.ComputationRunReq
|
||||
func (_e *Service_Expecter) Run(ctx interface{}, c interface{}) *Service_Run_Call {
|
||||
return &Service_Run_Call{Call: _e.mock.On("Run", ctx, c)}
|
||||
}
|
||||
|
||||
func (_c *Service_Run_Call) Run(run func(ctx context.Context, c *manager.ComputationRunReq)) *Service_Run_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run(args[0].(context.Context), args[1].(*manager.ComputationRunReq))
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Run_Call) Return(_a0 string, _a1 error) *Service_Run_Call {
|
||||
_c.Call.Return(_a0, _a1)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Run_Call) RunAndReturn(run func(context.Context, *manager.ComputationRunReq) (string, error)) *Service_Run_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Stop provides a mock function with given fields: ctx, computationID
|
||||
func (_m *Service) Stop(ctx context.Context, computationID string) error {
|
||||
ret := _m.Called(ctx, computationID)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Stop")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
|
||||
r0 = rf(ctx, computationID)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
|
||||
return r0
|
||||
}
|
||||
|
||||
// Service_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
|
||||
type Service_Stop_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Stop is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - computationID string
|
||||
func (_e *Service_Expecter) Stop(ctx interface{}, computationID interface{}) *Service_Stop_Call {
|
||||
return &Service_Stop_Call{Call: _e.mock.On("Stop", ctx, computationID)}
|
||||
}
|
||||
|
||||
func (_c *Service_Stop_Call) Run(run func(ctx context.Context, computationID string)) *Service_Stop_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run(args[0].(context.Context), args[1].(string))
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Stop_Call) Return(_a0 error) *Service_Stop_Call {
|
||||
_c.Call.Return(_a0)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Stop_Call) RunAndReturn(run func(context.Context, string) error) *Service_Stop_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 {
|
||||
|
||||
+15
-34
@@ -56,13 +56,7 @@ type SevConfig struct {
|
||||
ID string `env:"SEV_ID" envDefault:"sev0"`
|
||||
CBitPos int `env:"SEV_CBITPOS" envDefault:"51"`
|
||||
ReducedPhysBits int `env:"SEV_REDUCED_PHYS_BITS" envDefault:"1"`
|
||||
EnableHostData bool `env:"ENABLE_HOST_DATA" envDefault:"false"`
|
||||
HostData string `env:"HOST_DATA" envDefault:""`
|
||||
}
|
||||
|
||||
type IGVMConfig struct {
|
||||
ID string `env:"IGVM_ID" envDefault:"igvm0"`
|
||||
File string `env:"IGVM_FILE" envDefault:"/root/coconut-qemu.igvm"`
|
||||
HostData string `env:"HOST_DATA" envDefault:""`
|
||||
}
|
||||
|
||||
type VSockConfig struct {
|
||||
@@ -86,6 +80,9 @@ type Config struct {
|
||||
MemID string `env:"MEM_ID" envDefault:"ram1"`
|
||||
MemoryConfig
|
||||
|
||||
// Kernel hash
|
||||
KernelHash bool `env:"KERNEL_HASH" envDefault:"false"`
|
||||
|
||||
// OVMF
|
||||
OVMFCodeConfig
|
||||
OVMFVarsConfig
|
||||
@@ -103,19 +100,12 @@ type Config struct {
|
||||
// SEV
|
||||
SevConfig
|
||||
|
||||
// vTPM
|
||||
IGVMConfig
|
||||
|
||||
// display
|
||||
NoGraphic bool `env:"NO_GRAPHIC" envDefault:"true"`
|
||||
Monitor string `env:"MONITOR" envDefault:"pty"`
|
||||
|
||||
// ports
|
||||
HostFwdRange string `env:"HOST_FWD_RANGE" envDefault:"6100-6200"`
|
||||
|
||||
// mounts
|
||||
CertsMount string `env:"CERTS_MOUNT" envDefault:""`
|
||||
EnvMount string `env:"ENV_MOUNT" envDefault:""`
|
||||
}
|
||||
|
||||
func (config Config) ConstructQemuArgs() []string {
|
||||
@@ -179,39 +169,40 @@ func (config Config) ConstructQemuArgs() []string {
|
||||
// SEV
|
||||
if config.EnableSEV || config.EnableSEVSNP {
|
||||
sevType := "sev-guest"
|
||||
kernelHash := ""
|
||||
hostData := ""
|
||||
|
||||
args = append(args, "-machine",
|
||||
fmt.Sprintf("confidential-guest-support=%s,memory-backend=%s,igvm-cfg=%s",
|
||||
fmt.Sprintf("confidential-guest-support=%s,memory-backend=%s",
|
||||
config.SevConfig.ID,
|
||||
config.MemID,
|
||||
config.IGVMConfig.ID))
|
||||
config.MemID))
|
||||
|
||||
if config.EnableSEVSNP {
|
||||
args = append(args, "-bios", config.OVMFCodeConfig.File)
|
||||
sevType = "sev-snp-guest"
|
||||
|
||||
if config.SevConfig.EnableHostData {
|
||||
if config.SevConfig.HostData != "" {
|
||||
hostData = fmt.Sprintf(",host-data=%s", config.SevConfig.HostData)
|
||||
}
|
||||
}
|
||||
|
||||
if config.KernelHash {
|
||||
kernelHash = ",kernel-hashes=on"
|
||||
}
|
||||
|
||||
args = append(args, "-object",
|
||||
fmt.Sprintf("memory-backend-memfd,id=%s,size=%s,share=true,prealloc=false",
|
||||
config.MemID,
|
||||
config.MemoryConfig.Size))
|
||||
|
||||
args = append(args, "-object",
|
||||
fmt.Sprintf("%s,id=%s,cbitpos=%d,reduced-phys-bits=%d%s",
|
||||
fmt.Sprintf("%s,id=%s,cbitpos=%d,reduced-phys-bits=%d%s%s",
|
||||
sevType,
|
||||
config.SevConfig.ID,
|
||||
config.SevConfig.CBitPos,
|
||||
config.SevConfig.ReducedPhysBits,
|
||||
kernelHash,
|
||||
hostData))
|
||||
|
||||
args = append(args, "-object",
|
||||
fmt.Sprintf("igvm-cfg,id=%s,file=%s",
|
||||
config.IGVMConfig.ID,
|
||||
config.IGVMConfig.File))
|
||||
}
|
||||
|
||||
args = append(args, "-kernel", config.DiskImgConfig.KernelFile)
|
||||
@@ -225,15 +216,5 @@ func (config Config) ConstructQemuArgs() []string {
|
||||
|
||||
args = append(args, "-monitor", config.Monitor)
|
||||
|
||||
if config.CertsMount != "" {
|
||||
args = append(args, "-fsdev", fmt.Sprintf("local,id=cert_fs,path=%s,security_model=mapped", config.CertsMount))
|
||||
args = append(args, "-device", "virtio-9p-pci,fsdev=cert_fs,mount_tag=certs_share")
|
||||
}
|
||||
|
||||
if config.EnvMount != "" {
|
||||
args = append(args, "-fsdev", fmt.Sprintf("local,id=env_fs,path=%s,security_model=mapped", config.EnvMount))
|
||||
args = append(args, "-device", "virtio-9p-pci,fsdev=env_fs,mount_tag=env_share")
|
||||
}
|
||||
|
||||
return args
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user