mirror of
https://github.com/ultravioletrs/cocos.git
synced 2026-08-07 23:31:55 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Drasko Draskovic
|
f2e55aead2 |
@@ -0,0 +1,33 @@
|
||||
## Jaeger
|
||||
COCOS_JAEGER_PORT=6831
|
||||
COCOS_JAEGER_FRONTEND=16686
|
||||
COCOS_JAEGER_COLLECTOR=14268
|
||||
COCOS_JAEGER_CONFIGS=5778
|
||||
COCOS_JAEGER_URL=http://jaeger:14268/api/traces
|
||||
COCOS_JAEGER_TRACE_RATIO=1.0
|
||||
|
||||
## Core Services
|
||||
|
||||
### Manager
|
||||
MANAGER_HTTP_HOST="cocos-manager"
|
||||
MANAGER_HTTP_PORT=9021
|
||||
MANAGER_HTTP_SERVER_CERT=""
|
||||
MANAGER_HTTP_SERVER_KEY=""
|
||||
MANAGER_GRPC_HOST="cocos-manager"
|
||||
MANAGER_GRPC_PORT=7003
|
||||
MANAGER_GRPC_SERVER_CERT=""
|
||||
MANAGER_GRPC_SERVER_KEY=""
|
||||
AGENT_GRPC_URL="192.168.100.4:7002"
|
||||
AGENT_GRPC_TIMEOUT=""
|
||||
AGENT_GRPC_CA_CERTS=""
|
||||
AGENT_GRPC_CLIENT_TLS=""
|
||||
MANAGER_INSTANCE_ID=""
|
||||
MANAGER_LOG_LEVEL=debug
|
||||
MANAGER_QEMU_USE_SUDO=false
|
||||
MANAGER_QEMU_ENABLE_SEV=false
|
||||
MANAGER_QEMU_SEV_CBITPOS=51
|
||||
MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/OVMF/OVMF_CODE.fd
|
||||
MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/OVMF/OVMF_VARS.fd
|
||||
|
||||
# Docker image tag
|
||||
COCOS_RELEASE_TAG=latest
|
||||
@@ -1,8 +0,0 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "gomod"
|
||||
directories:
|
||||
- "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
day: "monday"
|
||||
@@ -9,7 +9,6 @@ on:
|
||||
- "pkg/manager/*.pb.go"
|
||||
- "agent/agent.proto"
|
||||
- "agent/*.pb.go"
|
||||
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
@@ -30,13 +29,13 @@ jobs:
|
||||
- name: Install Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.26.x
|
||||
go-version: 1.21.x
|
||||
|
||||
- name: Set up protoc
|
||||
run: |
|
||||
PROTOC_VERSION=35.0
|
||||
PROTOC_GEN_VERSION=v1.36.11
|
||||
PROTOC_GRPC_VERSION=v1.6.0
|
||||
PROTOC_VERSION=25.2
|
||||
PROTOC_GEN_VERSION=v1.31.0
|
||||
PROTOC_GRPC_VERSION=v1.3.0
|
||||
|
||||
# Download and install protoc
|
||||
PROTOC_ZIP=protoc-$PROTOC_VERSION-linux-x86_64.zip
|
||||
@@ -56,7 +55,7 @@ jobs:
|
||||
- name: Set up Cocos-AI
|
||||
run: |
|
||||
# Rename .pb.go files to .pb.go.tmp to prevent conflicts
|
||||
for p in $(ls manager/*.pb.go); do
|
||||
for p in $(ls pkg/manager/*.pb.go); do
|
||||
mv $p $p.tmp
|
||||
done
|
||||
|
||||
@@ -68,7 +67,7 @@ jobs:
|
||||
make protoc
|
||||
|
||||
# Compare generated Go files with the original ones
|
||||
for p in $(ls manager/*.pb.go); do
|
||||
for p in $(ls pkg/manager/*.pb.go); do
|
||||
if ! cmp -s $p $p.tmp; then
|
||||
echo "Proto file and generated Go file $p are out of sync!"
|
||||
exit 1
|
||||
|
||||
+58
-73
@@ -1,88 +1,73 @@
|
||||
name: Build and Release Hal
|
||||
name: Build and Release
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- "*"
|
||||
- '*'
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 120
|
||||
|
||||
steps:
|
||||
- name: Free Disk Space
|
||||
run: |
|
||||
sudo apt-get clean
|
||||
sudo rm -rf /usr/share/dotnet
|
||||
sudo rm -rf /opt/ghc
|
||||
sudo rm -rf "/usr/local/share/boost"
|
||||
sudo rm -rf "$AGENT_TOOLSDIRECTORY"
|
||||
mkdir -p "$RUNNER_TEMP/hostedtoolcache"
|
||||
echo "RUNNER_TOOL_CACHE=$RUNNER_TEMP/hostedtoolcache" >> "$GITHUB_ENV"
|
||||
echo "AGENT_TOOLSDIRECTORY=$RUNNER_TEMP/hostedtoolcache" >> "$GITHUB_ENV"
|
||||
df -h
|
||||
- name: Update Ubuntu
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get upgrade -y
|
||||
|
||||
- name: Install Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.21.x
|
||||
cache-dependency-path: "go.sum"
|
||||
|
||||
- name: Update Ubuntu
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get upgrade -y
|
||||
sudo apt-get clean
|
||||
sudo apt-get autoremove -y
|
||||
df -h
|
||||
- name: Checkout cocos
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: 'ultravioletrs/cocos'
|
||||
path: cocos
|
||||
|
||||
- name: Maximize build space
|
||||
uses: easimon/maximize-build-space@master
|
||||
with:
|
||||
root-reserve-mb: 35000
|
||||
swap-size-mb: 1024
|
||||
remove-dotnet: "true"
|
||||
remove-android: "true"
|
||||
- name: Check free space
|
||||
run: |
|
||||
echo "Free space:"
|
||||
df -h
|
||||
- name: Checkout buildroot
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: 'buildroot/buildroot'
|
||||
path: buildroot
|
||||
|
||||
- name: Checkout cocos
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
path: cocos
|
||||
- name: Build
|
||||
run: |
|
||||
cd buildroot
|
||||
make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
|
||||
make
|
||||
|
||||
- name: Install Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.26.x
|
||||
cache-dependency-path: "cocos/go.sum"
|
||||
- name: Create Release
|
||||
id: create_release
|
||||
uses: actions/create-release@latest
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
with:
|
||||
tag_name: ${{ github.ref }}
|
||||
release_name: Release ${{ github.ref }}
|
||||
draft: false
|
||||
prerelease: false
|
||||
|
||||
- name: Checkout buildroot
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: "buildroot/buildroot"
|
||||
path: buildroot
|
||||
ref: 2025.08-rc3
|
||||
- name: Upload Release Asset
|
||||
id: upload-release-kernel
|
||||
uses: actions/upload-release-asset@latest
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
with:
|
||||
upload_url: ${{ steps.create_release.outputs.upload_url }}
|
||||
asset_path: ./buildroot/output/images/bzImage
|
||||
asset_name: bzImage
|
||||
asset_content_type: application/octet-stream
|
||||
|
||||
- name: Build hal
|
||||
run: |
|
||||
cd buildroot
|
||||
make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
|
||||
make
|
||||
|
||||
- name: Build cocos
|
||||
run: |
|
||||
cd cocos
|
||||
make
|
||||
|
||||
- name: Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
files: |
|
||||
buildroot/output/images/bzImage
|
||||
buildroot/output/images/rootfs.cpio.gz
|
||||
cocos/build/cocos-agent
|
||||
cocos/build/cocos-cli
|
||||
cocos/build/cocos-manager
|
||||
cocos/build/cocos-attestation-service
|
||||
cocos/build/cocos-log-forwarder
|
||||
cocos/build/cocos-computation-runner
|
||||
cocos/build/cocos-egress-proxy
|
||||
cocos/build/cocos-ingress-proxy
|
||||
cocos/build/igvmmeasure
|
||||
- name: Upload Release Asset
|
||||
id: upload-release-rootfs
|
||||
uses: actions/upload-release-asset@latest
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
with:
|
||||
upload_url: ${{ steps.create_release.outputs.upload_url }}
|
||||
asset_path: ./buildroot/output/images/rootfs.cpio.gz
|
||||
asset_name: rootfs.cpio.gz
|
||||
asset_content_type: application/gzip
|
||||
|
||||
@@ -9,8 +9,9 @@ on:
|
||||
- main
|
||||
|
||||
jobs:
|
||||
lint:
|
||||
ci:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
@@ -18,71 +19,16 @@ jobs:
|
||||
- name: Install Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.26.x
|
||||
go-version: 1.21.x
|
||||
|
||||
- name: golangci-lint
|
||||
uses: golangci/golangci-lint-action@v8
|
||||
uses: golangci/golangci-lint-action@v3
|
||||
with:
|
||||
version: v2.11.1
|
||||
version: latest
|
||||
|
||||
- name: Build
|
||||
run: make
|
||||
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
needs: lint
|
||||
strategy:
|
||||
matrix:
|
||||
module: [agent, cli, cmd, internal, pkg, manager]
|
||||
include:
|
||||
- module: manager
|
||||
sudo: true
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Install Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.26.x
|
||||
|
||||
- name: Create coverage directory
|
||||
run: mkdir -p coverage
|
||||
|
||||
- name: Run tests for ${{ matrix.module }}
|
||||
run: |
|
||||
if [[ "${{ matrix.module }}" == "manager" ]]; then
|
||||
sudo GOTOOLCHAIN=go1.26.0+auto go test -v --race -covermode=atomic -coverprofile coverage/${{ matrix.module }}.out ./${{ matrix.module }}/...
|
||||
else
|
||||
GOTOOLCHAIN=go1.26.0+auto go test -v --race -covermode=atomic -coverprofile coverage/${{ matrix.module }}.out ./${{ matrix.module }}/...
|
||||
fi
|
||||
make
|
||||
|
||||
- name: Upload coverage artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: coverage-${{ matrix.module }}
|
||||
path: coverage/${{ matrix.module }}.out
|
||||
retention-days: 1
|
||||
|
||||
upload-coverage:
|
||||
runs-on: ubuntu-latest
|
||||
needs: test
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Download all coverage artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
pattern: coverage-*
|
||||
path: coverage/
|
||||
merge-multiple: true
|
||||
|
||||
- name: Upload results to Codecov
|
||||
uses: codecov/codecov-action@v4
|
||||
with:
|
||||
token: ${{ secrets.CODECOV_TOKEN }}
|
||||
files: ./coverage/*.out
|
||||
codecov_yml_path: codecov.yml
|
||||
verbose: true
|
||||
- name: Run tests
|
||||
run: go test -v --race -covermode=atomic -coverprofile cover.out ./...
|
||||
|
||||
+3
-28
@@ -1,33 +1,8 @@
|
||||
build
|
||||
|
||||
build
|
||||
cmd/manager/img
|
||||
cmd/manager/iso
|
||||
cmd/manager/tmp
|
||||
|
||||
.cov
|
||||
|
||||
*.pem
|
||||
|
||||
dist/
|
||||
*.zip
|
||||
*.spec
|
||||
*.tar
|
||||
|
||||
# Generated by Cargo
|
||||
# will have compiled files and executables
|
||||
debug/
|
||||
target/
|
||||
|
||||
# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
|
||||
# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
|
||||
Cargo.lock
|
||||
!tools/nvidia-attestation-helper/Cargo.lock
|
||||
|
||||
# These are backup files generated by rustfmt
|
||||
**/*.rs.bk
|
||||
|
||||
# MSVC Windows builds of rustc generate these, which store debugging information
|
||||
*.pdb
|
||||
|
||||
*.enc
|
||||
*.key
|
||||
*.pub
|
||||
.codex
|
||||
|
||||
+69
-81
@@ -1,98 +1,86 @@
|
||||
version: "2"
|
||||
run:
|
||||
timeout: 3m
|
||||
|
||||
issues:
|
||||
max-issues-per-linter: 10
|
||||
max-same-issues: 10
|
||||
|
||||
linters-settings:
|
||||
importas:
|
||||
no-unaliased: true
|
||||
no-extra-aliases: false
|
||||
alias:
|
||||
- pkg: github.com/absmach/magistrala/logger
|
||||
alias: mglog
|
||||
|
||||
gocritic:
|
||||
enabled-checks:
|
||||
- captLocal
|
||||
- dupImport
|
||||
- singleCaseSwitch
|
||||
- switchTrue
|
||||
- importShadow
|
||||
- httpNoBody
|
||||
- paramTypeCombine
|
||||
- emptyStringTest
|
||||
- builtinShadow
|
||||
- exposedSyncMutex
|
||||
disabled-checks:
|
||||
- appendAssign
|
||||
enabled-tags:
|
||||
- diagnostic
|
||||
disabled-tags:
|
||||
- performance
|
||||
- style
|
||||
- experimental
|
||||
- opinionated
|
||||
misspell:
|
||||
ignore-words:
|
||||
- "mosquitto"
|
||||
stylecheck:
|
||||
checks: ["-ST1000", "-ST1003", "-ST1020", "-ST1021", "-ST1022"]
|
||||
goheader:
|
||||
template: |-
|
||||
Copyright (c) Ultraviolet
|
||||
SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
linters:
|
||||
default: none
|
||||
disable-all: true
|
||||
enable:
|
||||
- importas
|
||||
- gocritic
|
||||
- gosimple
|
||||
- errcheck
|
||||
- govet
|
||||
- unused
|
||||
- goconst
|
||||
- godot
|
||||
- godox
|
||||
- ineffassign
|
||||
- misspell
|
||||
- stylecheck
|
||||
- whitespace
|
||||
- gci
|
||||
- gofmt
|
||||
- goimports
|
||||
- loggercheck
|
||||
- goheader
|
||||
- asasalint
|
||||
- asciicheck
|
||||
- bidichk
|
||||
- contextcheck
|
||||
- copyloopvar
|
||||
- decorder
|
||||
- dogsled
|
||||
- dupword
|
||||
- errcheck
|
||||
- errchkjson
|
||||
- errname
|
||||
- execinquery
|
||||
- exportloopref
|
||||
- ginkgolinter
|
||||
- gocheckcompilerdirectives
|
||||
- goconst
|
||||
- gocritic
|
||||
- godot
|
||||
- godox
|
||||
- goheader
|
||||
- gofumpt
|
||||
- goprintffuncname
|
||||
- govet
|
||||
- importas
|
||||
- ineffassign
|
||||
- loggercheck
|
||||
- makezero
|
||||
- mirror
|
||||
- misspell
|
||||
- nakedret
|
||||
- staticcheck
|
||||
- unused
|
||||
- whitespace
|
||||
settings:
|
||||
gocritic:
|
||||
enabled-checks:
|
||||
- dupImport
|
||||
- importShadow
|
||||
- httpNoBody
|
||||
- paramTypeCombine
|
||||
- emptyStringTest
|
||||
- builtinShadow
|
||||
- exposedSyncMutex
|
||||
disabled-checks:
|
||||
- appendAssign
|
||||
enabled-tags:
|
||||
- diagnostic
|
||||
disabled-tags:
|
||||
- performance
|
||||
- style
|
||||
- experimental
|
||||
- opinionated
|
||||
goheader:
|
||||
template: |-
|
||||
Copyright (c) Ultraviolet
|
||||
SPDX-License-Identifier: Apache-2.0
|
||||
staticcheck:
|
||||
checks:
|
||||
- -ST1000
|
||||
- -ST1003
|
||||
- -ST1020
|
||||
- -ST1021
|
||||
- -ST1022
|
||||
exclusions:
|
||||
generated: lax
|
||||
presets:
|
||||
- comments
|
||||
- common-false-positives
|
||||
- legacy
|
||||
- std-error-handling
|
||||
rules:
|
||||
- linters:
|
||||
- errcheck
|
||||
path: build/
|
||||
- linters:
|
||||
- makezero
|
||||
text: with non-zero initialized length
|
||||
paths:
|
||||
- build
|
||||
- third_party$
|
||||
- builtin$
|
||||
- examples$
|
||||
issues:
|
||||
max-issues-per-linter: 10
|
||||
max-same-issues: 10
|
||||
formatters:
|
||||
enable:
|
||||
- gci
|
||||
- gofmt
|
||||
- gofumpt
|
||||
- goimports
|
||||
exclusions:
|
||||
generated: lax
|
||||
paths:
|
||||
- third_party$
|
||||
- builtin$
|
||||
- examples$
|
||||
- dupword
|
||||
-162
@@ -1,162 +0,0 @@
|
||||
pkgname: mocks
|
||||
template: testify
|
||||
template-data:
|
||||
boilerplate-file: ./boilerplate.txt
|
||||
unroll-variadic: true
|
||||
packages:
|
||||
github.com/ultravioletrs/cocos/agent:
|
||||
interfaces:
|
||||
AgentService_AlgoClient:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
AgentService_DataClient:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
AgentService_IMAMeasurementsClient:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
Service:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/agent/algorithm:
|
||||
interfaces:
|
||||
Algorithm:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/agent/auth:
|
||||
interfaces:
|
||||
Authenticator:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage:
|
||||
interfaces:
|
||||
Storage:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/agent/cvms/server:
|
||||
interfaces:
|
||||
AgentServer:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/agent/events:
|
||||
interfaces:
|
||||
Service:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/agent/statemachine:
|
||||
interfaces:
|
||||
StateMachine:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/pkg/server:
|
||||
interfaces:
|
||||
Server:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/manager:
|
||||
interfaces:
|
||||
ManagerServiceClient:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
Service:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/manager/qemu:
|
||||
interfaces:
|
||||
Persistence:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/manager/vm:
|
||||
interfaces:
|
||||
StateMachine:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
VM:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/pkg/attestation:
|
||||
interfaces:
|
||||
Provider:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
Verifier:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig:
|
||||
interfaces:
|
||||
MeasurementProvider:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/pkg/clients/grpc:
|
||||
interfaces:
|
||||
Client:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/pkg/sdk:
|
||||
interfaces:
|
||||
SDK:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/pkg/atls:
|
||||
interfaces:
|
||||
CertificateProvider:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/pkg/clients/grpc/runner:
|
||||
interfaces:
|
||||
Client:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
github.com/ultravioletrs/cocos/internal/proto/attestation-agent:
|
||||
interfaces:
|
||||
AttestationAgentServiceClient:
|
||||
config:
|
||||
dir: '{{.InterfaceDir}}/mocks'
|
||||
structname: '{{.InterfaceName}}'
|
||||
filename: "{{.InterfaceName | lower}}.go"
|
||||
@@ -1,30 +1,12 @@
|
||||
BUILD_DIR = build
|
||||
SERVICES = manager agent cli attestation-service log-forwarder computation-runner egress-proxy ingress-proxy
|
||||
NVIDIA_ATTESTATION_HELPER = nvidia-attestation-helper
|
||||
NVIDIA_ATTESTATION_HELPER_DIR = tools/$(NVIDIA_ATTESTATION_HELPER)
|
||||
NVIDIA_ATTESTATION_HELPER_MANIFEST = $(NVIDIA_ATTESTATION_HELPER_DIR)/Cargo.toml
|
||||
NVIDIA_ATTESTATION_HELPER_BINARY = $(BUILD_DIR)/$(NVIDIA_ATTESTATION_HELPER)
|
||||
NVIDIA_ATTESTATION_HELPER_LIB_DIR = $(BUILD_DIR)/lib
|
||||
NVAT_SDK_CPP_DIR ?= $(firstword $(wildcard $(HOME)/.cargo/git/checkouts/attestation-sdk-*/*/nv-attestation-sdk-cpp))
|
||||
NVAT_SDK_CPP_BUILD_DIR ?= $(NVAT_SDK_CPP_DIR)/build
|
||||
NVAT_SDK_HEADER ?= $(NVAT_SDK_CPP_BUILD_DIR)/include/nvat.h
|
||||
NVAT_SDK_SHARED_LIB ?= $(NVAT_SDK_CPP_BUILD_DIR)/libnvat.so.1
|
||||
NVAT_SYSTEM_HEADER ?= /usr/include/nvat.h
|
||||
CARGO ?= cargo
|
||||
CMAKE ?= cmake
|
||||
SERVICES = manager agent cli
|
||||
CGO_ENABLED ?= 0
|
||||
GOARCH ?= amd64
|
||||
VERSION ?= $(shell git describe --abbrev=0 --tags --always)
|
||||
COMMIT ?= $(shell git rev-parse HEAD)
|
||||
TIME ?= $(shell date +%F_%T)
|
||||
EMBED_ENABLED ?= 0
|
||||
NVAT_USE_SYSTEM_LIB ?=
|
||||
INSTALL_DIR ?= /usr/local/bin
|
||||
CONFIG_DIR ?= /etc/cocos
|
||||
SERVICE_NAME ?= cocos-manager
|
||||
SERVICE_DIR ?= /etc/systemd/system
|
||||
SERVICE_FILE = init/systemd/$(SERVICE_NAME).service
|
||||
IGVM_BUILD_SCRIPT := ./scripts/igvmmeasure/igvm.sh
|
||||
empty:=
|
||||
space:= $(empty) $(empty)
|
||||
|
||||
define compile_service
|
||||
CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) \
|
||||
@@ -32,83 +14,16 @@ define compile_service
|
||||
-X 'github.com/absmach/magistrala.BuildTime=$(TIME)' \
|
||||
-X 'github.com/absmach/magistrala.Version=$(VERSION)' \
|
||||
-X 'github.com/absmach/magistrala.Commit=$(COMMIT)'" \
|
||||
$(if $(filter 1,$(EMBED_ENABLED)),-tags "embed",) \
|
||||
-o ${BUILD_DIR}/cocos-$(1) ./cmd/$(1)
|
||||
-o ${BUILD_DIR}/cocos-$(1) cmd/$(1)/main.go
|
||||
endef
|
||||
|
||||
NVIDIA_ATTESTATION_HELPER_CARGO_ENV = $(if $(filter 1,$(NVAT_USE_SYSTEM_LIB)),NVAT_USE_SYSTEM_LIB=1,)
|
||||
NVIDIA_ATTESTATION_HELPER_RUSTFLAGS = $(strip $(RUSTFLAGS) $(if $(filter 1,$(NVAT_USE_SYSTEM_LIB)),,-C link-arg=-Wl,-rpath,$$ORIGIN/lib))
|
||||
|
||||
.PHONY: all $(SERVICES) $(NVIDIA_ATTESTATION_HELPER) nvidia-attestation-helper-prereqs install clean
|
||||
.PHONY: all $(SERVICES)
|
||||
|
||||
all: $(SERVICES)
|
||||
|
||||
$(BUILD_DIR):
|
||||
mkdir -p $(BUILD_DIR)
|
||||
|
||||
$(SERVICES): | $(BUILD_DIR)
|
||||
$(call compile_service,$@)
|
||||
@if [ "$@" = "cli" ] || [ "$@" = "manager" ]; then $(MAKE) build-igvm; fi
|
||||
|
||||
nvidia-attestation-helper-prereqs:
|
||||
ifeq ($(filter 1,$(NVAT_USE_SYSTEM_LIB)),1)
|
||||
@test -f $(NVAT_SYSTEM_HEADER) || \
|
||||
( echo "Missing $(NVAT_SYSTEM_HEADER). Install the NVAT development package or run without NVAT_USE_SYSTEM_LIB=1."; exit 1 )
|
||||
@ldconfig -p | grep -q libnvat.so.1 || \
|
||||
( echo "libnvat.so.1 not found in the dynamic linker cache. Install the NVAT runtime package or run without NVAT_USE_SYSTEM_LIB=1."; exit 1 )
|
||||
else
|
||||
@if [ -z "$(NVAT_SDK_CPP_DIR)" ]; then \
|
||||
echo "Unable to locate nv-attestation-sdk-cpp under $$HOME/.cargo/git/checkouts."; \
|
||||
echo "Run 'cargo fetch --manifest-path $(NVIDIA_ATTESTATION_HELPER_MANIFEST)' first, or install NVAT and use 'make NVAT_USE_SYSTEM_LIB=1 $(NVIDIA_ATTESTATION_HELPER)'."; \
|
||||
exit 1; \
|
||||
fi
|
||||
@if [ ! -f "$(NVAT_SDK_HEADER)" ] || [ ! -f "$(NVAT_SDK_SHARED_LIB)" ]; then \
|
||||
$(CMAKE) -S $(NVAT_SDK_CPP_DIR) -B $(NVAT_SDK_CPP_BUILD_DIR) && \
|
||||
$(CMAKE) --build $(NVAT_SDK_CPP_BUILD_DIR); \
|
||||
fi
|
||||
endif
|
||||
|
||||
$(NVIDIA_ATTESTATION_HELPER): nvidia-attestation-helper-prereqs | $(BUILD_DIR)
|
||||
RUSTFLAGS='$(NVIDIA_ATTESTATION_HELPER_RUSTFLAGS)' $(NVIDIA_ATTESTATION_HELPER_CARGO_ENV) $(CARGO) build --manifest-path $(NVIDIA_ATTESTATION_HELPER_MANIFEST) --release
|
||||
install -m 755 $(NVIDIA_ATTESTATION_HELPER_DIR)/target/release/$(NVIDIA_ATTESTATION_HELPER) $(NVIDIA_ATTESTATION_HELPER_BINARY)
|
||||
@if [ "$(filter 1,$(NVAT_USE_SYSTEM_LIB))" != "1" ]; then \
|
||||
install -d $(NVIDIA_ATTESTATION_HELPER_LIB_DIR); \
|
||||
install -m 755 $(NVAT_SDK_SHARED_LIB) $(NVIDIA_ATTESTATION_HELPER_LIB_DIR)/libnvat.so.1; \
|
||||
fi
|
||||
$(SERVICES):
|
||||
$(call compile_service,$(@))
|
||||
|
||||
protoc:
|
||||
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/agent.proto
|
||||
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative manager/manager.proto
|
||||
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/events/events.proto
|
||||
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/cvms/cvms.proto
|
||||
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative internal/proto/attestation/v1/attestation.proto
|
||||
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative internal/proto/attestation-agent/attestation-agent.proto
|
||||
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/log/log.proto
|
||||
protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/runner/runner.proto
|
||||
|
||||
mocks:
|
||||
mockery --config ./.mockery.yml
|
||||
|
||||
install: $(SERVICES)
|
||||
install -d $(INSTALL_DIR)
|
||||
install $(BUILD_DIR)/cocos-cli $(INSTALL_DIR)/cocos-cli
|
||||
install $(BUILD_DIR)/cocos-manager $(INSTALL_DIR)/cocos-manager
|
||||
install -d $(CONFIG_DIR)
|
||||
install cocos-manager.env $(CONFIG_DIR)/cocos-manager.env
|
||||
|
||||
clean:
|
||||
rm -rf $(BUILD_DIR)
|
||||
|
||||
run: install_service
|
||||
sudo systemctl start $(SERVICE_NAME).service
|
||||
|
||||
stop:
|
||||
sudo systemctl stop $(SERVICE_NAME).service
|
||||
|
||||
install_service:
|
||||
sudo install -m 644 $(SERVICE_FILE) $(SERVICE_DIR)/$(SERVICE_NAME).service
|
||||
sudo systemctl daemon-reload
|
||||
|
||||
build-igvm:
|
||||
@echo "Running build script for igvmmeasure..."
|
||||
@$(IGVM_BUILD_SCRIPT)
|
||||
protoc -I. --go_out=./pkg --go_opt=paths=source_relative --go-grpc_out=./pkg --go-grpc_opt=paths=source_relative manager/manager.proto
|
||||
|
||||
@@ -1,80 +1,40 @@
|
||||
<div align="center">
|
||||
|
||||
# Cocos AI 🥥
|
||||
|
||||
**Confidential Computing System for AI**
|
||||
|
||||
**Made with ❤️ by [Ultraviolet](https://ultraviolet.rs/)**
|
||||
|
||||
[](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
|
||||
[Cocos AI (Confdential Computing System for AI/ML)][cocos] is a platform for secure multiparty computation (SMPC)
|
||||
based on the [Confidential Computing][cc] and [Trusted Execution Environments (TEEs)][tee].
|
||||
|
||||
<p align="center">
|
||||
<img src="https://cocos.ai/images/Collaborative%20AI.drawio.svg" alt="Cocos AI Illustration" width="400" height="400">
|
||||
<img src="https://cocos.ai/images/Collaborative%20AI.drawio.svg" width="500" height="500">
|
||||
</p>
|
||||
|
||||
## Features 🛠️
|
||||
With Cocos AI it becomes possible to run AI/ML workloads on combined datasets from multiple organizations
|
||||
while guaranteeing the privacy and security of the data and the algorithms.
|
||||
Data is always encrypted, protected by hardware secure enclaves (Trusted Execution Environments),
|
||||
attested via secure remote attestation protocols, and invisible to cloud processors or any other
|
||||
3rd party to which computation is offloaded.
|
||||
|
||||
Cocos AI provides essential features for secure and efficient collaborative AI/ML:
|
||||
## Usage
|
||||
|
||||
- 🖥️ **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.
|
||||
The quickest way to start using Cocos is via the CLI. The latest version can be downloaded from the [official releases page][rel].
|
||||
|
||||
## 🚀 Quick Start
|
||||
It can also be built and used from the project's root directory:
|
||||
|
||||
### Clone the Repository and Build Binaries
|
||||
```bash
|
||||
git clone git@github.com:ultravioletrs/cocos.git
|
||||
make
|
||||
make cli
|
||||
./build/cocos-cli version
|
||||
```
|
||||
|
||||
This will generate three binaries:
|
||||
```bash
|
||||
ls build/
|
||||
# cocos-agent cocos-cli cocos-manager
|
||||
```
|
||||
Additional details on using the CLI can be found in the [CLI documentation](https://docs.cocos.ai/cli).
|
||||
|
||||
### 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.
|
||||
## Documentation
|
||||
|
||||
## 📚 Documentation
|
||||
Official documentation is hosted at [Cocos official docs page][docs]. Documentation is auto-generated, checkout the instructions on [official docs repository](https://github.com/ultravioletrs/docs).
|
||||
|
||||
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).
|
||||
## License
|
||||
Cocos AI is a proprietary product created by Ultraviolet company.
|
||||
|
||||
Documentation is automatically generated from the [docs repository](https://github.com/ultravioletrs/docs). Contributions to documentation are welcome!
|
||||
[cc]: https://confidentialcomputing.io/white-papers-reports/
|
||||
[cocos]: https://cocos.ai/
|
||||
[rel]: https://github.com/ultraviolet/cocos/releases
|
||||
[tee]: https://en.wikipedia.org/wiki/Trusted_execution_environment
|
||||
[docs]: https://docs.cocos.ultraviolet.rs
|
||||
|
||||
## 🛡️ License
|
||||
|
||||
Cocos AI is published under the permissive open-source [Apache-2.0](LICENSE) license. Contributions are encouraged and appreciated!
|
||||
|
||||
## 🌐 Links and Resources
|
||||
|
||||
- [Cocos AI Website](https://cocos.ai/)
|
||||
- [Official Releases](https://github.com/ultravioletrs/cocos/releases)
|
||||
- [Confidential Computing Overview](https://confidentialcomputing.io/white-papers-reports/)
|
||||
- [Trusted Execution Environments (TEEs)](https://en.wikipedia.org/wiki/Trusted_execution_environment)
|
||||
|
||||
>This work has been partially supported by the [ELASTIC](https://elasticproject.eu/) and [CONFIDENTIAL6G](https://confidential6g.eu/), which received funding from the Smart Networks and Services Joint Undertaking (SNS JU) under the European Union’s Horizon Europe research and innovation programme under [Grant Agreement No. 101139067](https://cordis.europa.eu/project/id/101139067) and [Grant Agreement No. 101096435](https://cordis.europa.eu/project/id/101096435). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union. Neither the European Union nor the granting authority can be held responsible for them.
|
||||
|
||||
+10
-40
@@ -6,47 +6,17 @@ Agent service provides a barebones HTTP and gRPC API and Service interface imple
|
||||
|
||||
The service is configured using the environment variables from the following table. Note that any unset variables will be replaced with their default values.
|
||||
|
||||
| Variable | Description | Default |
|
||||
| ------------------------------ | ------------------------------------------------------------------------------------------------------------- | ----------------------------------------------- |
|
||||
| AGENT_LOG_LEVEL | Log level for agent service (debug, info, warn, error) | debug |
|
||||
| AGENT_VMPL | VMPL (Virtual Machine Privilege Level) for AMD SEV-SNP attestation (0-3) | 2 |
|
||||
| AGENT_GRPC_HOST | Agent service gRPC host address | 0.0.0.0 |
|
||||
| AGENT_CVM_GRPC_HOST | Agent service gRPC host | "" |
|
||||
| AGENT_CVM_GRPC_PORT | Agent service gRPC port | 7001 |
|
||||
| AGENT_CVM_GRPC_SERVER_CERT | Path to gRPC server certificate in pem format | "" |
|
||||
| AGENT_CVM_GRPC_SERVER_KEY | Path to gRPC server key in pem format | "" |
|
||||
| AGENT_CVM_GRPC_SERVER_CA_CERTS | Path to gRPC server CA certificate | "" |
|
||||
| AGENT_CVM_GRPC_CLIENT_CA_CERTS | Path to gRPC client CA certificate | "" |
|
||||
| AGENT_CVM_CA_URL | URL for CA service, if provided it will be used for certificate generation, used only with aTLS at the moment | "" |
|
||||
| AGENT_CVM_ID | Unique identifier for the CVM (Confidential Virtual Machine) | "" |
|
||||
| AGENT_CERTS_TOKEN | Authentication token for certificate service access | "" |
|
||||
| AGENT_MAA_URL | Microsoft Azure Attestation service URL for Azure attestation | https://sharedeus2.eus2.attest.azure.net |
|
||||
| AZURE_TDX_IMDS_URL | Azure TDX quote endpoint used by direct Azure TDX attestation | http://169.254.169.254/acc/tdquote |
|
||||
| AZURE_HCL_REFRESH_WAIT | Wait after writing TDX report data to Azure HCL vTPM storage before reading the refreshed HCL report | 3s |
|
||||
| AGENT_OS_BUILD | Operating system build information for attestation | UVC |
|
||||
| AGENT_OS_DISTRO | Operating system distribution information for attestation | UVC |
|
||||
| AGENT_OS_TYPE | Operating system type information for attestation | UVC |
|
||||
| ATTESTATION_SERVICE_SOCKET | Unix socket path for attestation service communication | /run/cocos/attestation.sock |
|
||||
| AGENT_ENABLE_ATLS | Enable Attestation TLS for secure communication | true |
|
||||
| Variable | Description | Default |
|
||||
| ----------------------------- | ------------------------------------------------------ | ------------------------------ |
|
||||
| AGENT_LOG_LEVEL | Log level for agent service (debug, info, warn, error) | info |
|
||||
| AGENT_GRPC_HOST | Agent service gRPC host | "" |
|
||||
| AGENT_GRPC_PORT | Agent service gRPC port | 7002 |
|
||||
| AGENT_GRPC_SERVER_CERT | Path to gRPC server certificate in pem format | "" |
|
||||
| AGENT_GRPC_SERVER_KEY | Path to gRPC server key in pem format | "" |
|
||||
| AGENT_GRPC_SERVER_CA_CERTS | Path to gRPC server CA certificate | "" |
|
||||
| AGENT_GRPC_CLIENT_CA_CERTS | Path to gRPC client CA certificate | "" |
|
||||
| COCOS_NOTIFICATION_SERVER_URL | Server to receive notification events from agent. | http:/localhost:9000 |
|
||||
|
||||
### Azure TDX Attestation
|
||||
|
||||
When the agent runs on an Azure TDX CVM, Azure attestation uses the direct Azure TDX flow. The agent writes TDX report data to Azure HCL vTPM storage, reads the refreshed HCL report, requests a TD quote from Azure IMDS, and submits the quote plus HCL runtime data to Microsoft Azure Attestation. This path does not depend on Confidential Containers attestation-agent `GetEvidence` or KBS token retrieval.
|
||||
|
||||
`AGENT_MAA_URL` selects the Microsoft Azure Attestation endpoint. `AZURE_TDX_IMDS_URL` can override the Azure IMDS TDX quote endpoint, and `AZURE_HCL_REFRESH_WAIT` controls the wait used to avoid reading a stale HCL report after report-data is written.
|
||||
|
||||
### Remote Resource Download (Optional)
|
||||
|
||||
The agent supports downloading encrypted algorithms and datasets from remote registries (S3, HTTP/HTTPS) and retrieving decryption keys from a Key Broker Service (KBS) via attestation.
|
||||
|
||||
| Variable | Description | Default |
|
||||
| ------------------------------ | ------------------------------------------------------------------------------------------------------------- | ----------------------------------------------- |
|
||||
| AWS_REGION | AWS region for S3 access (required for S3 downloads) | \"\" |
|
||||
| AWS_ACCESS_KEY_ID | AWS access key ID for S3 authentication | \"\" |
|
||||
| AWS_SECRET_ACCESS_KEY | AWS secret access key for S3 authentication | \"\" |
|
||||
| AWS_ENDPOINT_URL | Custom S3 endpoint URL (for S3-compatible services like MinIO) | \"\" |
|
||||
|
||||
**Note**: KBS URL is specified in the computation manifest, not as an environment variable. See [TESTING_REMOTE_RESOURCES.md](./TESTING_REMOTE_RESOURCES.md) for details on using remote resources.
|
||||
|
||||
## Deployment
|
||||
|
||||
|
||||
@@ -1,468 +0,0 @@
|
||||
# Testing Remote Resources with CoCo Key Provider
|
||||
|
||||
This guide explains how to test Cocos with encrypted remote resources using the Confidential Containers Key Provider ecosystem.
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
```
|
||||
┌────────────────────────────────────────────────────────────┐
|
||||
│ CVM (Agent) │
|
||||
│ │
|
||||
│ ┌──────────┐ ┌────────────────┐ ┌─────────────────┐ │
|
||||
│ │ Agent │──▶│ Skopeo │──▶│ CoCo Keyprovider│ │
|
||||
│ │ │ │ (ocicrypt) │ │ (gRPC:50011) │ │
|
||||
│ │ │ └───────┬────────┘ └────────┬────────┘ │
|
||||
│ │ │ │ │ │
|
||||
│ │ │ ┌───────▼────────┐ ┌────────▼────────┐ │
|
||||
│ │ │──▶│ S3/HTTP │ │ Attestation │ │
|
||||
│ │ │ │ Downloader │ │ Agent (50002) │ │
|
||||
│ └────┬─────┘ └───────┬────────┘ └────────┬────────┘ │
|
||||
│ │ │ │ │
|
||||
│ └──────────────────┼──────────────────────┘ │
|
||||
└────────┬─────────────────┼──────────────────────┬──────────┘
|
||||
│ (Resource) │ (Resource) │ (Attest)
|
||||
▼ ▼ ▼
|
||||
OCI Registry S3 / HTTP / GCS KBS
|
||||
(Key Broker)
|
||||
```
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### 1. Install Skopeo (Host Machine)
|
||||
|
||||
```bash
|
||||
# Ubuntu/Debian
|
||||
sudo apt-get install skopeo
|
||||
|
||||
# macOS
|
||||
brew install skopeo
|
||||
|
||||
# Or build from source
|
||||
git clone https://github.com/containers/skopeo
|
||||
cd skopeo
|
||||
make bin/skopeo
|
||||
sudo make install
|
||||
```
|
||||
|
||||
### 2. Start KBS Server (Host Machine)
|
||||
|
||||
```bash
|
||||
# Clone and build KBS
|
||||
git clone https://github.com/confidential-containers/trustee
|
||||
cd trustee/kbs
|
||||
# Patch Cargo.toml to disable SGX requirement (for testing only)
|
||||
sed -i 's/"all-verifier",//g' Cargo.toml
|
||||
|
||||
make
|
||||
make cli
|
||||
|
||||
# Generate admin keys
|
||||
openssl genpkey -algorithm ed25519 -out kbs-admin.key
|
||||
openssl pkey -in kbs-admin.key -pubout -out kbs-admin.pub
|
||||
|
||||
# Create KBS configuration file
|
||||
cat > kbs-config.toml << 'EOF'
|
||||
[http_server]
|
||||
sockets = ["0.0.0.0:8080"]
|
||||
insecure_http = true
|
||||
|
||||
[admin]
|
||||
type = "Simple"
|
||||
[[admin.personas]]
|
||||
id = "admin"
|
||||
public_key_path = "kbs-admin.pub"
|
||||
|
||||
[attestation_service]
|
||||
type = "coco_as_builtin"
|
||||
work_dir = "kbs-data/as"
|
||||
|
||||
[attestation_service.rvps_config]
|
||||
type = "BuiltIn"
|
||||
|
||||
[attestation_service.rvps_config.storage]
|
||||
type = "LocalFs"
|
||||
file_path = "kbs-data/rvps-values"
|
||||
|
||||
[[plugins]]
|
||||
name = "resource"
|
||||
type = "LocalFs"
|
||||
dir_path = "kbs-data/repository"
|
||||
EOF
|
||||
|
||||
# Create configuration directories
|
||||
mkdir -p kbs-data/as kbs-data/rvps kbs-data/repository
|
||||
|
||||
# Start KBS
|
||||
sudo ../target/release/kbs --config-file kbs-config.toml
|
||||
```
|
||||
|
||||
KBS will listen on `http://localhost:8080`
|
||||
|
||||
### 3. Setup Local OCI Registry (Optional)
|
||||
|
||||
For testing, you can use a local registry:
|
||||
|
||||
```bash
|
||||
docker run -d -p 5000:5000 --name registry registry:2
|
||||
```
|
||||
|
||||
## Creating Encrypted Resources
|
||||
|
||||
### Encrypt an Algorithm (Python Script)
|
||||
|
||||
```bash
|
||||
# 1. Create a simple algorithm
|
||||
cat > lin_reg.py << 'EOF'
|
||||
import pandas as pd
|
||||
from sklearn.linear_model import LinearRegression
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Load dataset
|
||||
data = pd.read_csv(sys.argv[1])
|
||||
X = data[['feature1', 'feature2']]
|
||||
y = data['target']
|
||||
|
||||
# Train model
|
||||
model = LinearRegression()
|
||||
model.fit(X, y)
|
||||
|
||||
# Save results
|
||||
os.makedirs("results", exist_ok=True)
|
||||
with open("results/output.txt", "w") as f:
|
||||
f.write(f"Coefficients: {model.coef_}\n")
|
||||
f.write(f"Intercept: {model.intercept_}\n")
|
||||
|
||||
print(f"Coefficients: {model.coef_}")
|
||||
print(f"Intercept: {model.intercept_}")
|
||||
EOF
|
||||
|
||||
# 2. Create requirements.txt
|
||||
cat > requirements.txt << 'EOF'
|
||||
pandas
|
||||
scikit-learn
|
||||
EOF
|
||||
|
||||
# 3. Create a Dockerfile
|
||||
cat > Dockerfile << 'EOF'
|
||||
FROM python:3.9-slim
|
||||
RUN pip install pandas scikit-learn
|
||||
COPY lin_reg.py /app/algorithm.py
|
||||
COPY requirements.txt /app/requirements.txt
|
||||
WORKDIR /app
|
||||
ENTRYPOINT ["python", "algorithm.py"]
|
||||
EOF
|
||||
|
||||
# 4. Build the image
|
||||
docker build -t localhost:5000/lin-reg-algo:v1.0 .
|
||||
docker push localhost:5000/lin-reg-algo:v1.0
|
||||
|
||||
# 5. Generate and store key
|
||||
openssl rand -out algo.key 32
|
||||
|
||||
# 6. Store key in KBS using kbs-client
|
||||
../target/release/kbs-client --url http://localhost:8080 config \
|
||||
--auth-private-key kbs-admin.key \
|
||||
set-resource \
|
||||
--path default/key/algo-key \
|
||||
--resource-file algo.key
|
||||
|
||||
# 7. Encrypt the image using Host Skopeo + Docker Keyprovider
|
||||
# Start Keyprovider in background
|
||||
docker run -d --rm --name keyprovider --network host \
|
||||
-v "$PWD:/work" -w /work \
|
||||
ghcr.io/confidential-containers/staged-images/coco-keyprovider:latest \
|
||||
coco_keyprovider --socket 127.0.0.1:50000
|
||||
|
||||
# Configure Ocicrypt to use local Keyprovider
|
||||
cat <<EOF > ocicrypt.conf
|
||||
{
|
||||
"key-providers": {
|
||||
"attestation-agent": {
|
||||
"grpc": "127.0.0.1:50000"
|
||||
}
|
||||
}
|
||||
}
|
||||
EOF
|
||||
export OCICRYPT_KEYPROVIDER_CONFIG=$(pwd)/ocicrypt.conf
|
||||
|
||||
# Encrypt Algo
|
||||
skopeo copy \
|
||||
--src-tls-verify=false \
|
||||
--dest-tls-verify=false \
|
||||
--encryption-key "provider:attestation-agent:keypath=/work/algo.key::keyid=kbs:///default/key/algo-key::algorithm=A256GCM" \
|
||||
docker://localhost:5000/lin-reg-algo:v1.0 \
|
||||
docker://localhost:5000/encrypted-lin-reg:v1.0
|
||||
|
||||
# Stop Keyprovider
|
||||
docker stop keyprovider
|
||||
```
|
||||
|
||||
### Encrypt a Dataset (CSV in OCI Image)
|
||||
|
||||
```bash
|
||||
# 1. Create dataset
|
||||
cat > iris.csv << 'EOF'
|
||||
feature1,feature2,target
|
||||
5.1,3.5,0
|
||||
4.9,3.0,0
|
||||
6.2,3.4,1
|
||||
5.9,3.0,1
|
||||
EOF
|
||||
|
||||
# 2. Create Dockerfile for dataset
|
||||
cat > Dockerfile.dataset << 'EOF'
|
||||
FROM scratch
|
||||
COPY iris.csv /data/iris.csv
|
||||
EOF
|
||||
|
||||
# 3. Build and push
|
||||
docker build -f Dockerfile.dataset -t localhost:5000/iris-dataset:v1.0 .
|
||||
docker push localhost:5000/iris-dataset:v1.0
|
||||
|
||||
# 4. Generate and store key
|
||||
# 4. Generate and store key
|
||||
openssl rand -out dataset.key 32
|
||||
../target/release/kbs-client --url http://localhost:8080 config \
|
||||
--auth-private-key kbs-admin.key \
|
||||
set-resource \
|
||||
--path default/key/dataset-key \
|
||||
--resource-file dataset.key
|
||||
|
||||
# 5. Encrypt dataset image using Host Skopeo + Docker Keyprovider
|
||||
# Start Keyprovider in background
|
||||
docker run -d --rm --name keyprovider --network host \
|
||||
-v "$PWD:/work" -w /work \
|
||||
ghcr.io/confidential-containers/staged-images/coco-keyprovider:latest \
|
||||
coco_keyprovider --socket 127.0.0.1:50000
|
||||
|
||||
# Configure Ocicrypt (if not already done)
|
||||
export OCICRYPT_KEYPROVIDER_CONFIG=$(pwd)/ocicrypt.conf
|
||||
|
||||
# Encrypt Dataset
|
||||
skopeo copy \
|
||||
--src-tls-verify=false \
|
||||
--dest-tls-verify=false \
|
||||
--encryption-key "provider:attestation-agent:keypath=/work/dataset.key::keyid=kbs:///default/key/dataset-key::algorithm=A256GCM" \
|
||||
docker://localhost:5000/iris-dataset:v1.0 \
|
||||
docker://localhost:5000/encrypted-iris:v1.0
|
||||
|
||||
# Stop Keyprovider
|
||||
docker stop keyprovider
|
||||
```
|
||||
|
||||
## Running a Computation
|
||||
|
||||
### 1. Start Manager (Host)
|
||||
|
||||
```bash
|
||||
cd /path/to/cocos-ai
|
||||
./build/cocos-manager
|
||||
```
|
||||
|
||||
### 2. Start CVMS Test Server (Host)
|
||||
|
||||
Get your host IP:
|
||||
```bash
|
||||
HOST_IP=$(ip -4 addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v 127.0.0.1 | head -n1)
|
||||
```
|
||||
|
||||
Start CVMS server:
|
||||
```bash
|
||||
# Calculate SHA3-256 of decrypted files using cocos-cli or cvms-test
|
||||
# NOTE: We use the hash of the original plaintext files, as the Agent validates the decrypted content.
|
||||
# For single files, use the file hash. For directories, use the hash of the directory (which the tools zip deterministically).
|
||||
|
||||
ALGO_HASH=$(./build/cocos-cli checksum lin_reg.py 2>&1 | awk '{print $NF}')
|
||||
|
||||
DATASET_HASH=$(./build/cocos-cli checksum iris.csv 2>&1 | awk '{print $NF}')
|
||||
|
||||
go build -o build/cvms-test ./test/cvms/main.go
|
||||
HOST=$HOST_IP PORT=7001 ./build/cvms-test \
|
||||
-public-key-path ./public.pem \
|
||||
-attested-tls-bool false \
|
||||
-algo-type python \
|
||||
-algo-source-url docker://$HOST_IP:5000/encrypted-lin-reg:v1.0 \
|
||||
-algo-kbs-path default/key/algo-key \
|
||||
-algo-kbs-url http://$HOST_IP:8080 \
|
||||
-algo-hash $ALGO_HASH \
|
||||
-algo-args datasets/dataset_0.csv \
|
||||
-dataset-source-urls docker://$HOST_IP:5000/encrypted-iris:v1.0 \
|
||||
-dataset-kbs-paths default/key/dataset-key \
|
||||
-dataset-kbs-urls http://$HOST_IP:8080 \
|
||||
-dataset-hash $DATASET_HASH
|
||||
```
|
||||
|
||||
> [!NOTE]
|
||||
> You must specify the KBS URL for each encrypted resource using `-algo-kbs-url` and `-dataset-kbs-urls`. A global KBS is no longer supported.
|
||||
|
||||
|
||||
### 3. Create VM via CLI (Host)
|
||||
|
||||
```bash
|
||||
export MANAGER_GRPC_URL=localhost:7002
|
||||
./build/cocos-cli create-vm \
|
||||
--server-url $HOST_IP:7001 \
|
||||
--log-level debug
|
||||
```
|
||||
|
||||
The agent will:
|
||||
1. Receive computation manifest from CVMS
|
||||
2. Use Skopeo to download encrypted OCI images
|
||||
3. Skopeo invokes CoCo Keyprovider via ocicrypt
|
||||
4. CoCo Keyprovider requests decryption key from KBS
|
||||
5. Attestation Agent generates TEE evidence for KBS
|
||||
6. KBS validates evidence and returns decryption key
|
||||
7. Image layers are decrypted and extracted
|
||||
8. Computation executes with decrypted algorithm and dataset
|
||||
|
||||
## Verifying the Setup
|
||||
|
||||
### Check CoCo Keyprovider Status (Inside CVM)
|
||||
|
||||
```bash
|
||||
# SSH into CVM or use console
|
||||
systemctl status coco-keyprovider
|
||||
journalctl -u coco-keyprovider -f
|
||||
```
|
||||
|
||||
### Check Attestation Agent Status
|
||||
|
||||
```bash
|
||||
systemctl status attestation-agent
|
||||
journalctl -u attestation-agent -f
|
||||
```
|
||||
|
||||
### Test Skopeo Decryption Manually
|
||||
|
||||
```bash
|
||||
# Inside CVM
|
||||
export OCICRYPT_KEYPROVIDER_CONFIG=/etc/ocicrypt_keyprovider.conf
|
||||
|
||||
skopeo copy \
|
||||
--src-tls-verify=false \
|
||||
--dest-tls-verify=false \
|
||||
--decryption-key provider:attestation-agent:cc_kbc::null \
|
||||
docker://localhost:5000/encrypted-lin-reg:v1.0 \
|
||||
oci:/tmp/decrypted-algo
|
||||
|
||||
# Verify decryption
|
||||
skopeo inspect oci:/tmp/decrypted-algo | jq -r '.LayersData[].MIMEType'
|
||||
# Should show: application/vnd.oci.image.layer.v1.tar+gzip
|
||||
```
|
||||
|
||||
## Computation Manifest Format
|
||||
|
||||
The CVMS server sends this manifest to the agent:
|
||||
|
||||
```json
|
||||
{
|
||||
"computation_id": "1",
|
||||
"algorithm": {
|
||||
"type": "oci-image",
|
||||
"uri": "docker://localhost:5000/encrypted-lin-reg:v1.0",
|
||||
"encrypted": true,
|
||||
"kbs_resource_path": "default/key/algo-key",
|
||||
"kbs": {
|
||||
"url": "http://192.168.100.15:8080",
|
||||
"enabled": true
|
||||
}
|
||||
},
|
||||
"datasets": [
|
||||
{
|
||||
"filename": "iris.csv",
|
||||
"source": {
|
||||
"type": "oci-image",
|
||||
"url": "docker://localhost:5000/encrypted-iris:v1.0",
|
||||
"encrypted": true,
|
||||
"kbs_resource_path": "default/key/dataset-key"
|
||||
},
|
||||
"kbs": {
|
||||
"url": "http://192.168.100.20:8080",
|
||||
"enabled": true
|
||||
}
|
||||
}
|
||||
],
|
||||
"kbs": {
|
||||
"url": "http://192.168.100.15:8080",
|
||||
"enabled": true
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### CoCo Keyprovider Not Starting
|
||||
|
||||
```bash
|
||||
# Check logs
|
||||
journalctl -u coco-keyprovider -n 50
|
||||
|
||||
# Verify socket is listening
|
||||
ss -tlnp | grep 50011
|
||||
|
||||
# Check environment
|
||||
cat /etc/default/coco-keyprovider
|
||||
```
|
||||
|
||||
### Skopeo Decryption Fails
|
||||
|
||||
```bash
|
||||
# Verify ocicrypt config
|
||||
cat /etc/ocicrypt_keyprovider.conf
|
||||
|
||||
# Test keyprovider connection
|
||||
grpcurl -plaintext 127.0.0.1:50011 list
|
||||
|
||||
# Check KBS connectivity from CVM
|
||||
curl http://HOST_IP:8080/kbs/v0/auth
|
||||
```
|
||||
|
||||
### KBS Returns 401
|
||||
|
||||
```bash
|
||||
# Check KBS logs on host
|
||||
# Verify attestation evidence format
|
||||
# Ensure KBS is configured for sample attestation
|
||||
```
|
||||
|
||||
## 4. Testing with Non-OCI Sources (S3, HTTP, GCS)
|
||||
|
||||
The `cvms` test utility also supports testing remote encrypted resources hosted in more traditional environments like S3-compatible storage or simple web servers, bypassing the need for container registries and OCI images.
|
||||
|
||||
### Supported Flags
|
||||
|
||||
The following flags define how resources should be fetched:
|
||||
|
||||
- `--algo-source-url`: The URL of the algorithm (e.g. `s3://bucket/algo.bin`, `https://server/algo.bin`)
|
||||
- `--algo-source-type`: The type of remote endpoint (`s3`, `gcs`, `https`, `http`). If omitted, it will automatically be inferred from the URL scheme.
|
||||
- `--algo-kbs-path`: The KBS path to retrieve the AES-256-GCM key from. If present, the agent will attempt decryption.
|
||||
- `--dataset-source-urls` and `--dataset-source-type`: Defines the locations and protocols for datasets.
|
||||
|
||||
### Encryption Format for Non-OCI Sources
|
||||
|
||||
Unlike OCI images where `ocicrypt` wraps the dataset, resources hosted on HTTP/S3 must be straightforwardly encrypted using **AES-256-GCM**.
|
||||
|
||||
The expected format is exactly as produced by standard Go AES-GCM:
|
||||
`nonce (12 bytes) || ciphertext || tag`
|
||||
|
||||
### Test Example
|
||||
|
||||
If you had a Python script encrypted using a key hosted at KBS path `default/my-keys/python-script` and uploaded to `s3://my-secure-bucket/script.enc`, you could run:
|
||||
|
||||
```bash
|
||||
cd test
|
||||
go run cvms/main.go --algo-source-url="s3://my-secure-bucket/script.enc" \
|
||||
--algo-source-type="s3" \
|
||||
--algo-kbs-path="default/my-keys/python-script" \
|
||||
--algo-type="python" \
|
||||
--public-key-path=./test-data/public-key.pem
|
||||
```
|
||||
|
||||
The system will:
|
||||
1. Connect via `attestation-agent` to the KBS to retrieve the symmetric key
|
||||
2. Use Google Cloud Storage client library methods (support for generic S3 via environment variables is standard) to fetch the resource
|
||||
3. Decrypt using AES-256-GCM
|
||||
4. Run the code normally
|
||||
|
||||
---
|
||||
+281
-338
@@ -3,8 +3,8 @@
|
||||
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.11
|
||||
// protoc v7.35.0
|
||||
// protoc-gen-go v1.31.0
|
||||
// protoc v4.25.2
|
||||
// source: agent/agent.proto
|
||||
|
||||
package agent
|
||||
@@ -14,7 +14,6 @@ import (
|
||||
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
|
||||
reflect "reflect"
|
||||
sync "sync"
|
||||
unsafe "unsafe"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -25,18 +24,22 @@ const (
|
||||
)
|
||||
|
||||
type AlgoRequest struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
|
||||
Requirements []byte `protobuf:"bytes,2,opt,name=requirements,proto3" json:"requirements,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
|
||||
Provider string `protobuf:"bytes,2,opt,name=provider,proto3" json:"provider,omitempty"`
|
||||
Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"`
|
||||
}
|
||||
|
||||
func (x *AlgoRequest) Reset() {
|
||||
*x = AlgoRequest{}
|
||||
mi := &file_agent_agent_proto_msgTypes[0]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_agent_proto_msgTypes[0]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *AlgoRequest) String() string {
|
||||
@@ -47,7 +50,7 @@ func (*AlgoRequest) ProtoMessage() {}
|
||||
|
||||
func (x *AlgoRequest) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[0]
|
||||
if x != nil {
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
@@ -69,24 +72,33 @@ func (x *AlgoRequest) GetAlgorithm() []byte {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *AlgoRequest) GetRequirements() []byte {
|
||||
func (x *AlgoRequest) GetProvider() string {
|
||||
if x != nil {
|
||||
return x.Requirements
|
||||
return x.Provider
|
||||
}
|
||||
return nil
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *AlgoRequest) GetId() string {
|
||||
if x != nil {
|
||||
return x.Id
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type AlgoResponse struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
}
|
||||
|
||||
func (x *AlgoResponse) Reset() {
|
||||
*x = AlgoResponse{}
|
||||
mi := &file_agent_agent_proto_msgTypes[1]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_agent_proto_msgTypes[1]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *AlgoResponse) String() string {
|
||||
@@ -97,7 +109,7 @@ func (*AlgoResponse) ProtoMessage() {}
|
||||
|
||||
func (x *AlgoResponse) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[1]
|
||||
if x != nil {
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
@@ -113,18 +125,22 @@ func (*AlgoResponse) Descriptor() ([]byte, []int) {
|
||||
}
|
||||
|
||||
type DataRequest struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Dataset []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
|
||||
Filename string `protobuf:"bytes,2,opt,name=filename,proto3" json:"filename,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
Dataset []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
|
||||
Provider string `protobuf:"bytes,2,opt,name=provider,proto3" json:"provider,omitempty"`
|
||||
Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"`
|
||||
}
|
||||
|
||||
func (x *DataRequest) Reset() {
|
||||
*x = DataRequest{}
|
||||
mi := &file_agent_agent_proto_msgTypes[2]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_agent_proto_msgTypes[2]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *DataRequest) String() string {
|
||||
@@ -135,7 +151,7 @@ func (*DataRequest) ProtoMessage() {}
|
||||
|
||||
func (x *DataRequest) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[2]
|
||||
if x != nil {
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
@@ -157,24 +173,33 @@ func (x *DataRequest) GetDataset() []byte {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *DataRequest) GetFilename() string {
|
||||
func (x *DataRequest) GetProvider() string {
|
||||
if x != nil {
|
||||
return x.Filename
|
||||
return x.Provider
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *DataRequest) GetId() string {
|
||||
if x != nil {
|
||||
return x.Id
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type DataResponse struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
}
|
||||
|
||||
func (x *DataResponse) Reset() {
|
||||
*x = DataResponse{}
|
||||
mi := &file_agent_agent_proto_msgTypes[3]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_agent_proto_msgTypes[3]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *DataResponse) String() string {
|
||||
@@ -185,7 +210,7 @@ func (*DataResponse) ProtoMessage() {}
|
||||
|
||||
func (x *DataResponse) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[3]
|
||||
if x != nil {
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
@@ -201,16 +226,20 @@ func (*DataResponse) Descriptor() ([]byte, []int) {
|
||||
}
|
||||
|
||||
type ResultRequest struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
Consumer string `protobuf:"bytes,1,opt,name=consumer,proto3" json:"consumer,omitempty"`
|
||||
}
|
||||
|
||||
func (x *ResultRequest) Reset() {
|
||||
*x = ResultRequest{}
|
||||
mi := &file_agent_agent_proto_msgTypes[4]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_agent_proto_msgTypes[4]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *ResultRequest) String() string {
|
||||
@@ -221,7 +250,7 @@ func (*ResultRequest) ProtoMessage() {}
|
||||
|
||||
func (x *ResultRequest) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[4]
|
||||
if x != nil {
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
@@ -236,18 +265,28 @@ func (*ResultRequest) Descriptor() ([]byte, []int) {
|
||||
return file_agent_agent_proto_rawDescGZIP(), []int{4}
|
||||
}
|
||||
|
||||
func (x *ResultRequest) GetConsumer() string {
|
||||
if x != nil {
|
||||
return x.Consumer
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type ResultResponse struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
|
||||
}
|
||||
|
||||
func (x *ResultResponse) Reset() {
|
||||
*x = ResultResponse{}
|
||||
mi := &file_agent_agent_proto_msgTypes[5]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_agent_proto_msgTypes[5]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *ResultResponse) String() string {
|
||||
@@ -258,7 +297,7 @@ func (*ResultResponse) ProtoMessage() {}
|
||||
|
||||
func (x *ResultResponse) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[5]
|
||||
if x != nil {
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
@@ -281,19 +320,20 @@ func (x *ResultResponse) GetFile() []byte {
|
||||
}
|
||||
|
||||
type AttestationRequest struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
TeeNonce []byte `protobuf:"bytes,1,opt,name=teeNonce,proto3" json:"teeNonce,omitempty"` // Should be less or equal 64 bytes.
|
||||
VtpmNonce []byte `protobuf:"bytes,2,opt,name=vtpmNonce,proto3" json:"vtpmNonce,omitempty"` // Should be less or equal 32 bytes.
|
||||
Type int32 `protobuf:"varint,3,opt,name=type,proto3" json:"type,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
ReportData []byte `protobuf:"bytes,1,opt,name=report_data,json=reportData,proto3" json:"report_data,omitempty"` // Should be of length 64.
|
||||
}
|
||||
|
||||
func (x *AttestationRequest) Reset() {
|
||||
*x = AttestationRequest{}
|
||||
mi := &file_agent_agent_proto_msgTypes[6]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_agent_proto_msgTypes[6]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *AttestationRequest) String() string {
|
||||
@@ -304,7 +344,7 @@ func (*AttestationRequest) ProtoMessage() {}
|
||||
|
||||
func (x *AttestationRequest) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[6]
|
||||
if x != nil {
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
@@ -319,39 +359,28 @@ func (*AttestationRequest) Descriptor() ([]byte, []int) {
|
||||
return file_agent_agent_proto_rawDescGZIP(), []int{6}
|
||||
}
|
||||
|
||||
func (x *AttestationRequest) GetTeeNonce() []byte {
|
||||
func (x *AttestationRequest) GetReportData() []byte {
|
||||
if x != nil {
|
||||
return x.TeeNonce
|
||||
return x.ReportData
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *AttestationRequest) GetVtpmNonce() []byte {
|
||||
if x != nil {
|
||||
return x.VtpmNonce
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *AttestationRequest) GetType() int32 {
|
||||
if x != nil {
|
||||
return x.Type
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
type AttestationResponse struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
|
||||
}
|
||||
|
||||
func (x *AttestationResponse) Reset() {
|
||||
*x = AttestationResponse{}
|
||||
mi := &file_agent_agent_proto_msgTypes[7]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_agent_agent_proto_msgTypes[7]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *AttestationResponse) String() string {
|
||||
@@ -362,7 +391,7 @@ func (*AttestationResponse) ProtoMessage() {}
|
||||
|
||||
func (x *AttestationResponse) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[7]
|
||||
if x != nil {
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
@@ -384,276 +413,91 @@ func (x *AttestationResponse) GetFile() []byte {
|
||||
return nil
|
||||
}
|
||||
|
||||
type IMAMeasurementsRequest struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *IMAMeasurementsRequest) Reset() {
|
||||
*x = IMAMeasurementsRequest{}
|
||||
mi := &file_agent_agent_proto_msgTypes[8]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *IMAMeasurementsRequest) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*IMAMeasurementsRequest) ProtoMessage() {}
|
||||
|
||||
func (x *IMAMeasurementsRequest) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[8]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use IMAMeasurementsRequest.ProtoReflect.Descriptor instead.
|
||||
func (*IMAMeasurementsRequest) Descriptor() ([]byte, []int) {
|
||||
return file_agent_agent_proto_rawDescGZIP(), []int{8}
|
||||
}
|
||||
|
||||
type IMAMeasurementsResponse struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
|
||||
Pcr10 []byte `protobuf:"bytes,2,opt,name=pcr10,proto3" json:"pcr10,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *IMAMeasurementsResponse) Reset() {
|
||||
*x = IMAMeasurementsResponse{}
|
||||
mi := &file_agent_agent_proto_msgTypes[9]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *IMAMeasurementsResponse) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*IMAMeasurementsResponse) ProtoMessage() {}
|
||||
|
||||
func (x *IMAMeasurementsResponse) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[9]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use IMAMeasurementsResponse.ProtoReflect.Descriptor instead.
|
||||
func (*IMAMeasurementsResponse) Descriptor() ([]byte, []int) {
|
||||
return file_agent_agent_proto_rawDescGZIP(), []int{9}
|
||||
}
|
||||
|
||||
func (x *IMAMeasurementsResponse) GetFile() []byte {
|
||||
if x != nil {
|
||||
return x.File
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *IMAMeasurementsResponse) GetPcr10() []byte {
|
||||
if x != nil {
|
||||
return x.Pcr10
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type AttestationTokenRequest struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
TokenNonce []byte `protobuf:"bytes,1,opt,name=tokenNonce,proto3" json:"tokenNonce,omitempty"` // Should be less or equal 32 bytes
|
||||
Type int32 `protobuf:"varint,3,opt,name=type,proto3" json:"type,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *AttestationTokenRequest) Reset() {
|
||||
*x = AttestationTokenRequest{}
|
||||
mi := &file_agent_agent_proto_msgTypes[10]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *AttestationTokenRequest) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*AttestationTokenRequest) ProtoMessage() {}
|
||||
|
||||
func (x *AttestationTokenRequest) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[10]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use AttestationTokenRequest.ProtoReflect.Descriptor instead.
|
||||
func (*AttestationTokenRequest) Descriptor() ([]byte, []int) {
|
||||
return file_agent_agent_proto_rawDescGZIP(), []int{10}
|
||||
}
|
||||
|
||||
func (x *AttestationTokenRequest) GetTokenNonce() []byte {
|
||||
if x != nil {
|
||||
return x.TokenNonce
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *AttestationTokenRequest) GetType() int32 {
|
||||
if x != nil {
|
||||
return x.Type
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
type AttestationTokenResponse struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *AttestationTokenResponse) Reset() {
|
||||
*x = AttestationTokenResponse{}
|
||||
mi := &file_agent_agent_proto_msgTypes[11]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *AttestationTokenResponse) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*AttestationTokenResponse) ProtoMessage() {}
|
||||
|
||||
func (x *AttestationTokenResponse) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_agent_proto_msgTypes[11]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use AttestationTokenResponse.ProtoReflect.Descriptor instead.
|
||||
func (*AttestationTokenResponse) Descriptor() ([]byte, []int) {
|
||||
return file_agent_agent_proto_rawDescGZIP(), []int{11}
|
||||
}
|
||||
|
||||
func (x *AttestationTokenResponse) GetFile() []byte {
|
||||
if x != nil {
|
||||
return x.File
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var File_agent_agent_proto protoreflect.FileDescriptor
|
||||
|
||||
const file_agent_agent_proto_rawDesc = "" +
|
||||
"\n" +
|
||||
"\x11agent/agent.proto\x12\x05agent\"O\n" +
|
||||
"\vAlgoRequest\x12\x1c\n" +
|
||||
"\talgorithm\x18\x01 \x01(\fR\talgorithm\x12\"\n" +
|
||||
"\frequirements\x18\x02 \x01(\fR\frequirements\"\x0e\n" +
|
||||
"\fAlgoResponse\"C\n" +
|
||||
"\vDataRequest\x12\x18\n" +
|
||||
"\adataset\x18\x01 \x01(\fR\adataset\x12\x1a\n" +
|
||||
"\bfilename\x18\x02 \x01(\tR\bfilename\"\x0e\n" +
|
||||
"\fDataResponse\"\x0f\n" +
|
||||
"\rResultRequest\"$\n" +
|
||||
"\x0eResultResponse\x12\x12\n" +
|
||||
"\x04file\x18\x01 \x01(\fR\x04file\"b\n" +
|
||||
"\x12AttestationRequest\x12\x1a\n" +
|
||||
"\bteeNonce\x18\x01 \x01(\fR\bteeNonce\x12\x1c\n" +
|
||||
"\tvtpmNonce\x18\x02 \x01(\fR\tvtpmNonce\x12\x12\n" +
|
||||
"\x04type\x18\x03 \x01(\x05R\x04type\")\n" +
|
||||
"\x13AttestationResponse\x12\x12\n" +
|
||||
"\x04file\x18\x01 \x01(\fR\x04file\"\x18\n" +
|
||||
"\x16IMAMeasurementsRequest\"C\n" +
|
||||
"\x17IMAMeasurementsResponse\x12\x12\n" +
|
||||
"\x04file\x18\x01 \x01(\fR\x04file\x12\x14\n" +
|
||||
"\x05pcr10\x18\x02 \x01(\fR\x05pcr10\"M\n" +
|
||||
"\x17AttestationTokenRequest\x12\x1e\n" +
|
||||
"\n" +
|
||||
"tokenNonce\x18\x01 \x01(\fR\n" +
|
||||
"tokenNonce\x12\x12\n" +
|
||||
"\x04type\x18\x03 \x01(\x05R\x04type\".\n" +
|
||||
"\x18AttestationTokenResponse\x12\x12\n" +
|
||||
"\x04file\x18\x01 \x01(\fR\x04file2\xaf\x03\n" +
|
||||
"\fAgentService\x123\n" +
|
||||
"\x04Algo\x12\x12.agent.AlgoRequest\x1a\x13.agent.AlgoResponse\"\x00(\x01\x123\n" +
|
||||
"\x04Data\x12\x12.agent.DataRequest\x1a\x13.agent.DataResponse\"\x00(\x01\x129\n" +
|
||||
"\x06Result\x12\x14.agent.ResultRequest\x1a\x15.agent.ResultResponse\"\x000\x01\x12H\n" +
|
||||
"\vAttestation\x12\x19.agent.AttestationRequest\x1a\x1a.agent.AttestationResponse\"\x000\x01\x12T\n" +
|
||||
"\x0fIMAMeasurements\x12\x1d.agent.IMAMeasurementsRequest\x1a\x1e.agent.IMAMeasurementsResponse\"\x000\x01\x12Z\n" +
|
||||
"\x15AzureAttestationToken\x12\x1e.agent.AttestationTokenRequest\x1a\x1f.agent.AttestationTokenResponse\"\x00B\tZ\a./agentb\x06proto3"
|
||||
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, 0x57, 0x0a, 0x0b, 0x41, 0x6c,
|
||||
0x67, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1c, 0x0a, 0x09, 0x61, 0x6c, 0x67,
|
||||
0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x09, 0x61, 0x6c,
|
||||
0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x12, 0x1a, 0x0a, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69,
|
||||
0x64, 0x65, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69,
|
||||
0x64, 0x65, 0x72, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52,
|
||||
0x02, 0x69, 0x64, 0x22, 0x0e, 0x0a, 0x0c, 0x41, 0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f,
|
||||
0x6e, 0x73, 0x65, 0x22, 0x53, 0x0a, 0x0b, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65,
|
||||
0x73, 0x74, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x18, 0x01, 0x20,
|
||||
0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x12, 0x1a, 0x0a, 0x08,
|
||||
0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08,
|
||||
0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x03,
|
||||
0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x22, 0x0e, 0x0a, 0x0c, 0x44, 0x61, 0x74, 0x61,
|
||||
0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x2b, 0x0a, 0x0d, 0x52, 0x65, 0x73, 0x75,
|
||||
0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x63, 0x6f, 0x6e,
|
||||
0x73, 0x75, 0x6d, 0x65, 0x72, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x63, 0x6f, 0x6e,
|
||||
0x73, 0x75, 0x6d, 0x65, 0x72, 0x22, 0x24, 0x0a, 0x0e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52,
|
||||
0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18,
|
||||
0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x22, 0x35, 0x0a, 0x12, 0x41,
|
||||
0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
|
||||
0x74, 0x12, 0x1f, 0x0a, 0x0b, 0x72, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x5f, 0x64, 0x61, 0x74, 0x61,
|
||||
0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0a, 0x72, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x44, 0x61,
|
||||
0x74, 0x61, 0x22, 0x29, 0x0a, 0x13, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f,
|
||||
0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c,
|
||||
0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x32, 0xf5, 0x01,
|
||||
0x0a, 0x0c, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x31,
|
||||
0x0a, 0x04, 0x41, 0x6c, 0x67, 0x6f, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41,
|
||||
0x6c, 0x67, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65,
|
||||
0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22,
|
||||
0x00, 0x12, 0x31, 0x0a, 0x04, 0x44, 0x61, 0x74, 0x61, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e,
|
||||
0x74, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e,
|
||||
0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
|
||||
0x73, 0x65, 0x22, 0x00, 0x12, 0x37, 0x0a, 0x06, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x12, 0x14,
|
||||
0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71,
|
||||
0x75, 0x65, 0x73, 0x74, 0x1a, 0x15, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73,
|
||||
0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x46, 0x0a,
|
||||
0x0b, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x19, 0x2e, 0x61,
|
||||
0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e,
|
||||
0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e,
|
||||
0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f,
|
||||
0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61, 0x67, 0x65, 0x6e, 0x74,
|
||||
0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
|
||||
}
|
||||
|
||||
var (
|
||||
file_agent_agent_proto_rawDescOnce sync.Once
|
||||
file_agent_agent_proto_rawDescData []byte
|
||||
file_agent_agent_proto_rawDescData = file_agent_agent_proto_rawDesc
|
||||
)
|
||||
|
||||
func file_agent_agent_proto_rawDescGZIP() []byte {
|
||||
file_agent_agent_proto_rawDescOnce.Do(func() {
|
||||
file_agent_agent_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_agent_proto_rawDesc), len(file_agent_agent_proto_rawDesc)))
|
||||
file_agent_agent_proto_rawDescData = protoimpl.X.CompressGZIP(file_agent_agent_proto_rawDescData)
|
||||
})
|
||||
return file_agent_agent_proto_rawDescData
|
||||
}
|
||||
|
||||
var file_agent_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 12)
|
||||
var file_agent_agent_proto_goTypes = []any{
|
||||
(*AlgoRequest)(nil), // 0: agent.AlgoRequest
|
||||
(*AlgoResponse)(nil), // 1: agent.AlgoResponse
|
||||
(*DataRequest)(nil), // 2: agent.DataRequest
|
||||
(*DataResponse)(nil), // 3: agent.DataResponse
|
||||
(*ResultRequest)(nil), // 4: agent.ResultRequest
|
||||
(*ResultResponse)(nil), // 5: agent.ResultResponse
|
||||
(*AttestationRequest)(nil), // 6: agent.AttestationRequest
|
||||
(*AttestationResponse)(nil), // 7: agent.AttestationResponse
|
||||
(*IMAMeasurementsRequest)(nil), // 8: agent.IMAMeasurementsRequest
|
||||
(*IMAMeasurementsResponse)(nil), // 9: agent.IMAMeasurementsResponse
|
||||
(*AttestationTokenRequest)(nil), // 10: agent.AttestationTokenRequest
|
||||
(*AttestationTokenResponse)(nil), // 11: agent.AttestationTokenResponse
|
||||
var file_agent_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 8)
|
||||
var file_agent_agent_proto_goTypes = []interface{}{
|
||||
(*AlgoRequest)(nil), // 0: agent.AlgoRequest
|
||||
(*AlgoResponse)(nil), // 1: agent.AlgoResponse
|
||||
(*DataRequest)(nil), // 2: agent.DataRequest
|
||||
(*DataResponse)(nil), // 3: agent.DataResponse
|
||||
(*ResultRequest)(nil), // 4: agent.ResultRequest
|
||||
(*ResultResponse)(nil), // 5: agent.ResultResponse
|
||||
(*AttestationRequest)(nil), // 6: agent.AttestationRequest
|
||||
(*AttestationResponse)(nil), // 7: agent.AttestationResponse
|
||||
}
|
||||
var file_agent_agent_proto_depIdxs = []int32{
|
||||
0, // 0: agent.AgentService.Algo:input_type -> agent.AlgoRequest
|
||||
2, // 1: agent.AgentService.Data:input_type -> agent.DataRequest
|
||||
4, // 2: agent.AgentService.Result:input_type -> agent.ResultRequest
|
||||
6, // 3: agent.AgentService.Attestation:input_type -> agent.AttestationRequest
|
||||
8, // 4: agent.AgentService.IMAMeasurements:input_type -> agent.IMAMeasurementsRequest
|
||||
10, // 5: agent.AgentService.AzureAttestationToken:input_type -> agent.AttestationTokenRequest
|
||||
1, // 6: agent.AgentService.Algo:output_type -> agent.AlgoResponse
|
||||
3, // 7: agent.AgentService.Data:output_type -> agent.DataResponse
|
||||
5, // 8: agent.AgentService.Result:output_type -> agent.ResultResponse
|
||||
7, // 9: agent.AgentService.Attestation:output_type -> agent.AttestationResponse
|
||||
9, // 10: agent.AgentService.IMAMeasurements:output_type -> agent.IMAMeasurementsResponse
|
||||
11, // 11: agent.AgentService.AzureAttestationToken:output_type -> agent.AttestationTokenResponse
|
||||
6, // [6:12] is the sub-list for method output_type
|
||||
0, // [0:6] is the sub-list for method input_type
|
||||
0, // [0:0] is the sub-list for extension type_name
|
||||
0, // [0:0] is the sub-list for extension extendee
|
||||
0, // [0:0] is the sub-list for field type_name
|
||||
0, // 0: agent.AgentService.Algo:input_type -> agent.AlgoRequest
|
||||
2, // 1: agent.AgentService.Data:input_type -> agent.DataRequest
|
||||
4, // 2: agent.AgentService.Result:input_type -> agent.ResultRequest
|
||||
6, // 3: agent.AgentService.Attestation:input_type -> agent.AttestationRequest
|
||||
1, // 4: agent.AgentService.Algo:output_type -> agent.AlgoResponse
|
||||
3, // 5: agent.AgentService.Data:output_type -> agent.DataResponse
|
||||
5, // 6: agent.AgentService.Result:output_type -> agent.ResultResponse
|
||||
7, // 7: agent.AgentService.Attestation:output_type -> agent.AttestationResponse
|
||||
4, // [4:8] is the sub-list for method output_type
|
||||
0, // [0:4] is the sub-list for method input_type
|
||||
0, // [0:0] is the sub-list for extension type_name
|
||||
0, // [0:0] is the sub-list for extension extendee
|
||||
0, // [0:0] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_agent_agent_proto_init() }
|
||||
@@ -661,13 +505,111 @@ func file_agent_agent_proto_init() {
|
||||
if File_agent_agent_proto != nil {
|
||||
return
|
||||
}
|
||||
if !protoimpl.UnsafeEnabled {
|
||||
file_agent_agent_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*AlgoRequest); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_agent_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*AlgoResponse); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_agent_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*DataRequest); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_agent_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*DataResponse); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_agent_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*ResultRequest); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_agent_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*ResultResponse); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_agent_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*AttestationRequest); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
file_agent_agent_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*AttestationResponse); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
type x struct{}
|
||||
out := protoimpl.TypeBuilder{
|
||||
File: protoimpl.DescBuilder{
|
||||
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
|
||||
RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_agent_proto_rawDesc), len(file_agent_agent_proto_rawDesc)),
|
||||
RawDescriptor: file_agent_agent_proto_rawDesc,
|
||||
NumEnums: 0,
|
||||
NumMessages: 12,
|
||||
NumMessages: 8,
|
||||
NumExtensions: 0,
|
||||
NumServices: 1,
|
||||
},
|
||||
@@ -676,6 +618,7 @@ func file_agent_agent_proto_init() {
|
||||
MessageInfos: file_agent_agent_proto_msgTypes,
|
||||
}.Build()
|
||||
File_agent_agent_proto = out.File
|
||||
file_agent_agent_proto_rawDesc = nil
|
||||
file_agent_agent_proto_goTypes = nil
|
||||
file_agent_agent_proto_depIdxs = nil
|
||||
}
|
||||
|
||||
+10
-27
@@ -8,29 +8,30 @@ package agent;
|
||||
option go_package = "./agent";
|
||||
|
||||
service AgentService {
|
||||
rpc Algo(stream AlgoRequest) returns (AlgoResponse) {}
|
||||
rpc Data(stream DataRequest) returns (DataResponse) {}
|
||||
rpc Result(ResultRequest) returns (stream ResultResponse) {}
|
||||
rpc Attestation(AttestationRequest) returns (stream AttestationResponse) {}
|
||||
rpc IMAMeasurements(IMAMeasurementsRequest) returns (stream IMAMeasurementsResponse) {}
|
||||
rpc AzureAttestationToken(AttestationTokenRequest) returns (AttestationTokenResponse) {}
|
||||
rpc Algo(AlgoRequest) returns (AlgoResponse) {}
|
||||
rpc Data(DataRequest) returns (DataResponse) {}
|
||||
rpc Result(ResultRequest) returns (ResultResponse) {}
|
||||
rpc Attestation(AttestationRequest) returns (AttestationResponse) {}
|
||||
}
|
||||
|
||||
message AlgoRequest {
|
||||
bytes algorithm = 1;
|
||||
bytes requirements = 2;
|
||||
string provider = 2;
|
||||
string id = 3;
|
||||
}
|
||||
|
||||
message AlgoResponse {}
|
||||
|
||||
message DataRequest {
|
||||
bytes dataset = 1;
|
||||
string filename = 2;
|
||||
string provider = 2;
|
||||
string id = 3;
|
||||
}
|
||||
|
||||
message DataResponse {}
|
||||
|
||||
message ResultRequest {
|
||||
string consumer = 1;
|
||||
}
|
||||
|
||||
message ResultResponse {
|
||||
@@ -38,27 +39,9 @@ message ResultResponse {
|
||||
}
|
||||
|
||||
message AttestationRequest {
|
||||
bytes teeNonce = 1; // Should be less or equal 64 bytes.
|
||||
bytes vtpmNonce = 2; // Should be less or equal 32 bytes.
|
||||
int32 type = 3;
|
||||
bytes report_data = 1; // Should be of length 64.
|
||||
}
|
||||
|
||||
message AttestationResponse {
|
||||
bytes file = 1;
|
||||
}
|
||||
|
||||
message IMAMeasurementsRequest {
|
||||
}
|
||||
|
||||
message IMAMeasurementsResponse {
|
||||
bytes file = 1;
|
||||
bytes pcr10 = 2;
|
||||
}
|
||||
|
||||
message AttestationTokenRequest{
|
||||
bytes tokenNonce = 1; // Should be less or equal 32 bytes
|
||||
int32 type = 3;
|
||||
}
|
||||
message AttestationTokenResponse{
|
||||
bytes file = 1;
|
||||
}
|
||||
|
||||
+117
-204
@@ -3,8 +3,8 @@
|
||||
|
||||
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
|
||||
// versions:
|
||||
// - protoc-gen-go-grpc v1.6.0
|
||||
// - protoc v7.35.0
|
||||
// - protoc-gen-go-grpc v1.3.0
|
||||
// - protoc v4.25.2
|
||||
// source: agent/agent.proto
|
||||
|
||||
package agent
|
||||
@@ -18,28 +18,24 @@ import (
|
||||
|
||||
// This is a compile-time assertion to ensure that this generated file
|
||||
// is compatible with the grpc package it is being compiled against.
|
||||
// Requires gRPC-Go v1.64.0 or later.
|
||||
const _ = grpc.SupportPackageIsVersion9
|
||||
// Requires gRPC-Go v1.32.0 or later.
|
||||
const _ = grpc.SupportPackageIsVersion7
|
||||
|
||||
const (
|
||||
AgentService_Algo_FullMethodName = "/agent.AgentService/Algo"
|
||||
AgentService_Data_FullMethodName = "/agent.AgentService/Data"
|
||||
AgentService_Result_FullMethodName = "/agent.AgentService/Result"
|
||||
AgentService_Attestation_FullMethodName = "/agent.AgentService/Attestation"
|
||||
AgentService_IMAMeasurements_FullMethodName = "/agent.AgentService/IMAMeasurements"
|
||||
AgentService_AzureAttestationToken_FullMethodName = "/agent.AgentService/AzureAttestationToken"
|
||||
AgentService_Algo_FullMethodName = "/agent.AgentService/Algo"
|
||||
AgentService_Data_FullMethodName = "/agent.AgentService/Data"
|
||||
AgentService_Result_FullMethodName = "/agent.AgentService/Result"
|
||||
AgentService_Attestation_FullMethodName = "/agent.AgentService/Attestation"
|
||||
)
|
||||
|
||||
// AgentServiceClient is the client API for AgentService service.
|
||||
//
|
||||
// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
|
||||
type AgentServiceClient interface {
|
||||
Algo(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[AlgoRequest, AlgoResponse], error)
|
||||
Data(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[DataRequest, DataResponse], error)
|
||||
Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ResultResponse], error)
|
||||
Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AttestationResponse], error)
|
||||
IMAMeasurements(ctx context.Context, in *IMAMeasurementsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[IMAMeasurementsResponse], error)
|
||||
AzureAttestationToken(ctx context.Context, in *AttestationTokenRequest, opts ...grpc.CallOption) (*AttestationTokenResponse, error)
|
||||
Algo(ctx context.Context, in *AlgoRequest, opts ...grpc.CallOption) (*AlgoResponse, error)
|
||||
Data(ctx context.Context, in *DataRequest, opts ...grpc.CallOption) (*DataResponse, error)
|
||||
Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (*ResultResponse, error)
|
||||
Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error)
|
||||
}
|
||||
|
||||
type agentServiceClient struct {
|
||||
@@ -50,93 +46,36 @@ func NewAgentServiceClient(cc grpc.ClientConnInterface) AgentServiceClient {
|
||||
return &agentServiceClient{cc}
|
||||
}
|
||||
|
||||
func (c *agentServiceClient) Algo(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[AlgoRequest, AlgoResponse], error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[0], AgentService_Algo_FullMethodName, cOpts...)
|
||||
func (c *agentServiceClient) Algo(ctx context.Context, in *AlgoRequest, opts ...grpc.CallOption) (*AlgoResponse, error) {
|
||||
out := new(AlgoResponse)
|
||||
err := c.cc.Invoke(ctx, AgentService_Algo_FullMethodName, in, out, opts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
x := &grpc.GenericClientStream[AlgoRequest, AlgoResponse]{ClientStream: stream}
|
||||
return x, nil
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_AlgoClient = grpc.ClientStreamingClient[AlgoRequest, AlgoResponse]
|
||||
|
||||
func (c *agentServiceClient) Data(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[DataRequest, DataResponse], error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[1], AgentService_Data_FullMethodName, cOpts...)
|
||||
func (c *agentServiceClient) Data(ctx context.Context, in *DataRequest, opts ...grpc.CallOption) (*DataResponse, error) {
|
||||
out := new(DataResponse)
|
||||
err := c.cc.Invoke(ctx, AgentService_Data_FullMethodName, in, out, opts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
x := &grpc.GenericClientStream[DataRequest, DataResponse]{ClientStream: stream}
|
||||
return x, nil
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_DataClient = grpc.ClientStreamingClient[DataRequest, DataResponse]
|
||||
|
||||
func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ResultResponse], error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[2], AgentService_Result_FullMethodName, cOpts...)
|
||||
func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (*ResultResponse, error) {
|
||||
out := new(ResultResponse)
|
||||
err := c.cc.Invoke(ctx, AgentService_Result_FullMethodName, in, out, opts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
x := &grpc.GenericClientStream[ResultRequest, ResultResponse]{ClientStream: stream}
|
||||
if err := x.ClientStream.SendMsg(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := x.ClientStream.CloseSend(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return x, nil
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_ResultClient = grpc.ServerStreamingClient[ResultResponse]
|
||||
|
||||
func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AttestationResponse], error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[3], AgentService_Attestation_FullMethodName, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
x := &grpc.GenericClientStream[AttestationRequest, AttestationResponse]{ClientStream: stream}
|
||||
if err := x.ClientStream.SendMsg(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := x.ClientStream.CloseSend(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return x, nil
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_AttestationClient = grpc.ServerStreamingClient[AttestationResponse]
|
||||
|
||||
func (c *agentServiceClient) IMAMeasurements(ctx context.Context, in *IMAMeasurementsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[IMAMeasurementsResponse], error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[4], AgentService_IMAMeasurements_FullMethodName, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
x := &grpc.GenericClientStream[IMAMeasurementsRequest, IMAMeasurementsResponse]{ClientStream: stream}
|
||||
if err := x.ClientStream.SendMsg(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := x.ClientStream.CloseSend(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return x, nil
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_IMAMeasurementsClient = grpc.ServerStreamingClient[IMAMeasurementsResponse]
|
||||
|
||||
func (c *agentServiceClient) AzureAttestationToken(ctx context.Context, in *AttestationTokenRequest, opts ...grpc.CallOption) (*AttestationTokenResponse, error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(AttestationTokenResponse)
|
||||
err := c.cc.Invoke(ctx, AgentService_AzureAttestationToken_FullMethodName, in, out, cOpts...)
|
||||
func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error) {
|
||||
out := new(AttestationResponse)
|
||||
err := c.cc.Invoke(ctx, AgentService_Attestation_FullMethodName, in, out, opts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -145,44 +84,32 @@ func (c *agentServiceClient) AzureAttestationToken(ctx context.Context, in *Atte
|
||||
|
||||
// AgentServiceServer is the server API for AgentService service.
|
||||
// All implementations must embed UnimplementedAgentServiceServer
|
||||
// for forward compatibility.
|
||||
// for forward compatibility
|
||||
type AgentServiceServer interface {
|
||||
Algo(grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]) error
|
||||
Data(grpc.ClientStreamingServer[DataRequest, DataResponse]) error
|
||||
Result(*ResultRequest, grpc.ServerStreamingServer[ResultResponse]) error
|
||||
Attestation(*AttestationRequest, grpc.ServerStreamingServer[AttestationResponse]) error
|
||||
IMAMeasurements(*IMAMeasurementsRequest, grpc.ServerStreamingServer[IMAMeasurementsResponse]) error
|
||||
AzureAttestationToken(context.Context, *AttestationTokenRequest) (*AttestationTokenResponse, error)
|
||||
Algo(context.Context, *AlgoRequest) (*AlgoResponse, error)
|
||||
Data(context.Context, *DataRequest) (*DataResponse, error)
|
||||
Result(context.Context, *ResultRequest) (*ResultResponse, error)
|
||||
Attestation(context.Context, *AttestationRequest) (*AttestationResponse, error)
|
||||
mustEmbedUnimplementedAgentServiceServer()
|
||||
}
|
||||
|
||||
// UnimplementedAgentServiceServer must be embedded to have
|
||||
// forward compatible implementations.
|
||||
//
|
||||
// NOTE: this should be embedded by value instead of pointer to avoid a nil
|
||||
// pointer dereference when methods are called.
|
||||
type UnimplementedAgentServiceServer struct{}
|
||||
// UnimplementedAgentServiceServer must be embedded to have forward compatible implementations.
|
||||
type UnimplementedAgentServiceServer struct {
|
||||
}
|
||||
|
||||
func (UnimplementedAgentServiceServer) Algo(grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]) error {
|
||||
return status.Error(codes.Unimplemented, "method Algo not implemented")
|
||||
func (UnimplementedAgentServiceServer) Algo(context.Context, *AlgoRequest) (*AlgoResponse, error) {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method Algo not implemented")
|
||||
}
|
||||
func (UnimplementedAgentServiceServer) Data(grpc.ClientStreamingServer[DataRequest, DataResponse]) error {
|
||||
return status.Error(codes.Unimplemented, "method Data not implemented")
|
||||
func (UnimplementedAgentServiceServer) Data(context.Context, *DataRequest) (*DataResponse, error) {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method Data not implemented")
|
||||
}
|
||||
func (UnimplementedAgentServiceServer) Result(*ResultRequest, grpc.ServerStreamingServer[ResultResponse]) error {
|
||||
return status.Error(codes.Unimplemented, "method Result not implemented")
|
||||
func (UnimplementedAgentServiceServer) Result(context.Context, *ResultRequest) (*ResultResponse, error) {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method Result not implemented")
|
||||
}
|
||||
func (UnimplementedAgentServiceServer) Attestation(*AttestationRequest, grpc.ServerStreamingServer[AttestationResponse]) error {
|
||||
return status.Error(codes.Unimplemented, "method Attestation not implemented")
|
||||
}
|
||||
func (UnimplementedAgentServiceServer) IMAMeasurements(*IMAMeasurementsRequest, grpc.ServerStreamingServer[IMAMeasurementsResponse]) error {
|
||||
return status.Error(codes.Unimplemented, "method IMAMeasurements not implemented")
|
||||
}
|
||||
func (UnimplementedAgentServiceServer) AzureAttestationToken(context.Context, *AttestationTokenRequest) (*AttestationTokenResponse, error) {
|
||||
return nil, status.Error(codes.Unimplemented, "method AzureAttestationToken not implemented")
|
||||
func (UnimplementedAgentServiceServer) Attestation(context.Context, *AttestationRequest) (*AttestationResponse, error) {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method Attestation not implemented")
|
||||
}
|
||||
func (UnimplementedAgentServiceServer) mustEmbedUnimplementedAgentServiceServer() {}
|
||||
func (UnimplementedAgentServiceServer) testEmbeddedByValue() {}
|
||||
|
||||
// UnsafeAgentServiceServer may be embedded to opt out of forward compatibility for this service.
|
||||
// Use of this interface is not recommended, as added methods to AgentServiceServer will
|
||||
@@ -192,77 +119,77 @@ type UnsafeAgentServiceServer interface {
|
||||
}
|
||||
|
||||
func RegisterAgentServiceServer(s grpc.ServiceRegistrar, srv AgentServiceServer) {
|
||||
// If the following call panics, it indicates UnimplementedAgentServiceServer was
|
||||
// embedded by pointer and is nil. This will cause panics if an
|
||||
// unimplemented method is ever invoked, so we test this at initialization
|
||||
// time to prevent it from happening at runtime later due to I/O.
|
||||
if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
|
||||
t.testEmbeddedByValue()
|
||||
}
|
||||
s.RegisterService(&AgentService_ServiceDesc, srv)
|
||||
}
|
||||
|
||||
func _AgentService_Algo_Handler(srv interface{}, stream grpc.ServerStream) error {
|
||||
return srv.(AgentServiceServer).Algo(&grpc.GenericServerStream[AlgoRequest, AlgoResponse]{ServerStream: stream})
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_AlgoServer = grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]
|
||||
|
||||
func _AgentService_Data_Handler(srv interface{}, stream grpc.ServerStream) error {
|
||||
return srv.(AgentServiceServer).Data(&grpc.GenericServerStream[DataRequest, DataResponse]{ServerStream: stream})
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_DataServer = grpc.ClientStreamingServer[DataRequest, DataResponse]
|
||||
|
||||
func _AgentService_Result_Handler(srv interface{}, stream grpc.ServerStream) error {
|
||||
m := new(ResultRequest)
|
||||
if err := stream.RecvMsg(m); err != nil {
|
||||
return err
|
||||
}
|
||||
return srv.(AgentServiceServer).Result(m, &grpc.GenericServerStream[ResultRequest, ResultResponse]{ServerStream: stream})
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_ResultServer = grpc.ServerStreamingServer[ResultResponse]
|
||||
|
||||
func _AgentService_Attestation_Handler(srv interface{}, stream grpc.ServerStream) error {
|
||||
m := new(AttestationRequest)
|
||||
if err := stream.RecvMsg(m); err != nil {
|
||||
return err
|
||||
}
|
||||
return srv.(AgentServiceServer).Attestation(m, &grpc.GenericServerStream[AttestationRequest, AttestationResponse]{ServerStream: stream})
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_AttestationServer = grpc.ServerStreamingServer[AttestationResponse]
|
||||
|
||||
func _AgentService_IMAMeasurements_Handler(srv interface{}, stream grpc.ServerStream) error {
|
||||
m := new(IMAMeasurementsRequest)
|
||||
if err := stream.RecvMsg(m); err != nil {
|
||||
return err
|
||||
}
|
||||
return srv.(AgentServiceServer).IMAMeasurements(m, &grpc.GenericServerStream[IMAMeasurementsRequest, IMAMeasurementsResponse]{ServerStream: stream})
|
||||
}
|
||||
|
||||
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
|
||||
type AgentService_IMAMeasurementsServer = grpc.ServerStreamingServer[IMAMeasurementsResponse]
|
||||
|
||||
func _AgentService_AzureAttestationToken_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(AttestationTokenRequest)
|
||||
func _AgentService_Algo_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(AlgoRequest)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(AgentServiceServer).AzureAttestationToken(ctx, in)
|
||||
return srv.(AgentServiceServer).Algo(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: AgentService_AzureAttestationToken_FullMethodName,
|
||||
FullMethod: AgentService_Algo_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(AgentServiceServer).AzureAttestationToken(ctx, req.(*AttestationTokenRequest))
|
||||
return srv.(AgentServiceServer).Algo(ctx, req.(*AlgoRequest))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
func _AgentService_Data_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(DataRequest)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(AgentServiceServer).Data(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: AgentService_Data_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(AgentServiceServer).Data(ctx, req.(*DataRequest))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
func _AgentService_Result_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(ResultRequest)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(AgentServiceServer).Result(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: AgentService_Result_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(AgentServiceServer).Result(ctx, req.(*ResultRequest))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
func _AgentService_Attestation_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(AttestationRequest)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(AgentServiceServer).Attestation(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: AgentService_Attestation_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(AgentServiceServer).Attestation(ctx, req.(*AttestationRequest))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
@@ -275,36 +202,22 @@ var AgentService_ServiceDesc = grpc.ServiceDesc{
|
||||
HandlerType: (*AgentServiceServer)(nil),
|
||||
Methods: []grpc.MethodDesc{
|
||||
{
|
||||
MethodName: "AzureAttestationToken",
|
||||
Handler: _AgentService_AzureAttestationToken_Handler,
|
||||
},
|
||||
},
|
||||
Streams: []grpc.StreamDesc{
|
||||
{
|
||||
StreamName: "Algo",
|
||||
Handler: _AgentService_Algo_Handler,
|
||||
ClientStreams: true,
|
||||
},
|
||||
{
|
||||
StreamName: "Data",
|
||||
Handler: _AgentService_Data_Handler,
|
||||
ClientStreams: true,
|
||||
},
|
||||
{
|
||||
StreamName: "Result",
|
||||
Handler: _AgentService_Result_Handler,
|
||||
ServerStreams: true,
|
||||
},
|
||||
{
|
||||
StreamName: "Attestation",
|
||||
Handler: _AgentService_Attestation_Handler,
|
||||
ServerStreams: true,
|
||||
},
|
||||
{
|
||||
StreamName: "IMAMeasurements",
|
||||
Handler: _AgentService_IMAMeasurements_Handler,
|
||||
ServerStreams: true,
|
||||
MethodName: "Algo",
|
||||
Handler: _AgentService_Algo_Handler,
|
||||
},
|
||||
{
|
||||
MethodName: "Data",
|
||||
Handler: _AgentService_Data_Handler,
|
||||
},
|
||||
{
|
||||
MethodName: "Result",
|
||||
Handler: _AgentService_Result_Handler,
|
||||
},
|
||||
{
|
||||
MethodName: "Attestation",
|
||||
Handler: _AgentService_Attestation_Handler,
|
||||
},
|
||||
},
|
||||
Streams: []grpc.StreamDesc{},
|
||||
Metadata: "agent/agent.proto",
|
||||
}
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
// Code generated by "stringer -type=AgentEvent"; DO NOT EDIT.
|
||||
|
||||
package agent
|
||||
|
||||
import "strconv"
|
||||
|
||||
func _() {
|
||||
// An "invalid array index" compiler error signifies that the constant values have changed.
|
||||
// Re-run the stringer command to generate them again.
|
||||
var x [1]struct{}
|
||||
_ = x[Start-0]
|
||||
_ = x[ManifestReceived-1]
|
||||
_ = x[AlgorithmReceived-2]
|
||||
_ = x[DataReceived-3]
|
||||
_ = x[RunComplete-4]
|
||||
_ = x[ResultsConsumed-5]
|
||||
_ = x[RunFailed-6]
|
||||
}
|
||||
|
||||
const _AgentEvent_name = "StartManifestReceivedAlgorithmReceivedDataReceivedRunCompleteResultsConsumedRunFailed"
|
||||
|
||||
var _AgentEvent_index = [...]uint8{0, 5, 21, 38, 50, 61, 76, 85}
|
||||
|
||||
func (i AgentEvent) String() string {
|
||||
if i < 0 || i >= AgentEvent(len(_AgentEvent_index)-1) {
|
||||
return "AgentEvent(" + strconv.FormatInt(int64(i), 10) + ")"
|
||||
}
|
||||
return _AgentEvent_name[_AgentEvent_index[i]:_AgentEvent_index[i+1]]
|
||||
}
|
||||
@@ -1,30 +0,0 @@
|
||||
// Code generated by "stringer -type=AgentState"; DO NOT EDIT.
|
||||
|
||||
package agent
|
||||
|
||||
import "strconv"
|
||||
|
||||
func _() {
|
||||
// An "invalid array index" compiler error signifies that the constant values have changed.
|
||||
// Re-run the stringer command to generate them again.
|
||||
var x [1]struct{}
|
||||
_ = x[Idle-0]
|
||||
_ = x[ReceivingManifest-1]
|
||||
_ = x[ReceivingAlgorithm-2]
|
||||
_ = x[ReceivingData-3]
|
||||
_ = x[Running-4]
|
||||
_ = x[ConsumingResults-5]
|
||||
_ = x[Complete-6]
|
||||
_ = x[Failed-7]
|
||||
}
|
||||
|
||||
const _AgentState_name = "IdleReceivingManifestReceivingAlgorithmReceivingDataRunningConsumingResultsCompleteFailed"
|
||||
|
||||
var _AgentState_index = [...]uint8{0, 4, 21, 39, 52, 59, 75, 83, 89}
|
||||
|
||||
func (i AgentState) String() string {
|
||||
if i < 0 || i >= AgentState(len(_AgentState_index)-1) {
|
||||
return "AgentState(" + strconv.FormatInt(int64(i), 10) + ")"
|
||||
}
|
||||
return _AgentState_name[_AgentState_index[i]:_AgentState_index[i+1]]
|
||||
}
|
||||
@@ -1,52 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package algorithm
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
type AlgorithType string
|
||||
|
||||
const (
|
||||
AlgoTypeBin AlgorithType = "bin"
|
||||
AlgoTypePython AlgorithType = "python"
|
||||
AlgoTypeWasm AlgorithType = "wasm"
|
||||
AlgoTypeDocker AlgorithType = "docker"
|
||||
AlgoTypeKey = "algo_type"
|
||||
AlgoArgsKey = "algo_args"
|
||||
|
||||
ResultsDir = "results"
|
||||
DatasetsDir = "datasets"
|
||||
AlgoWorkingDir = "/cocos"
|
||||
)
|
||||
|
||||
func AlgorithmTypeToContext(ctx context.Context, algoType string) context.Context {
|
||||
return metadata.AppendToOutgoingContext(ctx, AlgoTypeKey, algoType)
|
||||
}
|
||||
|
||||
func AlgorithmTypeFromContext(ctx context.Context) string {
|
||||
return metadata.ValueFromIncomingContext(ctx, AlgoTypeKey)[0]
|
||||
}
|
||||
|
||||
func AlgorithmArgsToContext(ctx context.Context, algoArgs []string) context.Context {
|
||||
for _, arg := range algoArgs {
|
||||
ctx = metadata.AppendToOutgoingContext(ctx, AlgoArgsKey, arg)
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
|
||||
func AlgorithmArgsFromContext(ctx context.Context) []string {
|
||||
return metadata.ValueFromIncomingContext(ctx, AlgoArgsKey)
|
||||
}
|
||||
|
||||
// Algorithm is an interface that specifies the API for an algorithm.
|
||||
type Algorithm interface {
|
||||
// Run executes the algorithm and returns the result.
|
||||
Run() error
|
||||
|
||||
// Stop stops the algorithm.
|
||||
Stop() error
|
||||
}
|
||||
@@ -1,77 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package binary
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"sync"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
)
|
||||
|
||||
var execCommand = exec.Command
|
||||
|
||||
var _ algorithm.Algorithm = (*binary)(nil)
|
||||
|
||||
type binary struct {
|
||||
algoFile string
|
||||
stderr io.Writer
|
||||
stdout io.Writer
|
||||
args []string
|
||||
cmd *exec.Cmd
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string, cmpID string) algorithm.Algorithm {
|
||||
return &binary{
|
||||
algoFile: algoFile,
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
|
||||
stdout: &logging.Stdout{Logger: logger},
|
||||
args: args,
|
||||
}
|
||||
}
|
||||
|
||||
func (b *binary) Run() error {
|
||||
b.mu.Lock()
|
||||
b.cmd = execCommand(b.algoFile, b.args...)
|
||||
b.cmd.Stderr = b.stderr
|
||||
b.cmd.Stdout = b.stdout
|
||||
|
||||
if err := b.cmd.Start(); err != nil {
|
||||
b.mu.Unlock()
|
||||
return fmt.Errorf("error starting algorithm: %v", err)
|
||||
}
|
||||
b.mu.Unlock()
|
||||
|
||||
if err := b.cmd.Wait(); err != nil {
|
||||
return fmt.Errorf("algorithm execution error: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *binary) Stop() error {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
|
||||
if b.cmd == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if b.cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := b.cmd.Process.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) {
|
||||
return fmt.Errorf("error stopping algorithm: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,170 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package binary
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
|
||||
"github.com/ultravioletrs/cocos/agent/events/mocks"
|
||||
)
|
||||
|
||||
func TestNewAlgorithm(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventsSvc := new(mocks.Service)
|
||||
algoFile := "/path/to/algo"
|
||||
args := []string{"arg1", "arg2"}
|
||||
|
||||
algo := NewAlgorithm(logger, eventsSvc, algoFile, args, "")
|
||||
|
||||
b, ok := algo.(*binary)
|
||||
if !ok {
|
||||
t.Fatalf("NewAlgorithm did not return a *binary")
|
||||
}
|
||||
|
||||
if b.algoFile != algoFile {
|
||||
t.Errorf("Expected algoFile to be %s, got %s", algoFile, b.algoFile)
|
||||
}
|
||||
|
||||
if len(b.args) != len(args) {
|
||||
t.Errorf("Expected %d args, got %d", len(args), len(b.args))
|
||||
}
|
||||
|
||||
for i, arg := range args {
|
||||
if b.args[i] != arg {
|
||||
t.Errorf("Expected arg %d to be %s, got %s", i, arg, b.args[i])
|
||||
}
|
||||
}
|
||||
|
||||
if _, ok := b.stderr.(*logging.Stderr); !ok {
|
||||
t.Errorf("Expected stderr to be *algorithm.Stderr")
|
||||
}
|
||||
|
||||
if _, ok := b.stdout.(*logging.Stdout); !ok {
|
||||
t.Errorf("Expected stdout to be *algorithm.Stdout")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBinaryRun(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
algoFile string
|
||||
args []string
|
||||
expectedError bool
|
||||
}{
|
||||
{
|
||||
name: "Successful execution",
|
||||
algoFile: "echo",
|
||||
args: []string{"Hello, World!"},
|
||||
expectedError: false,
|
||||
},
|
||||
{
|
||||
name: "Non-existent binary",
|
||||
algoFile: "non_existent_binary",
|
||||
args: []string{},
|
||||
expectedError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventsSvc := new(mocks.Service)
|
||||
eventsSvc.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return().Maybe()
|
||||
|
||||
b := NewAlgorithm(logger, eventsSvc, tt.algoFile, tt.args, "").(*binary)
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
b.stdout = &stdout
|
||||
b.stderr = &stderr
|
||||
|
||||
err := b.Run()
|
||||
|
||||
if tt.expectedError && err == nil {
|
||||
t.Errorf("Expected an error, but got none")
|
||||
}
|
||||
|
||||
if !tt.expectedError && err != nil {
|
||||
t.Errorf("Unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if !tt.expectedError {
|
||||
if stdout.Len() == 0 {
|
||||
t.Errorf("Expected non-empty stdout")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStop(t *testing.T) {
|
||||
t.Run("stop nil cmd", func(t *testing.T) {
|
||||
b := &binary{}
|
||||
err := b.Stop()
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
t.Run("stop with running process", func(t *testing.T) {
|
||||
b := &binary{
|
||||
algoFile: "sleep",
|
||||
args: []string{"10"},
|
||||
}
|
||||
if err := b.Run(); err != nil {
|
||||
t.Fatalf("Failed to start command: %v", err)
|
||||
}
|
||||
|
||||
err := b.Stop()
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify it actually stopped
|
||||
_ = b.cmd.Wait()
|
||||
})
|
||||
|
||||
t.Run("stop already exited", func(t *testing.T) {
|
||||
b := &binary{
|
||||
algoFile: "echo",
|
||||
args: []string{"test"},
|
||||
stdout: io.Discard,
|
||||
stderr: io.Discard,
|
||||
}
|
||||
if err := b.Run(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
err := b.Stop()
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestRunError(t *testing.T) {
|
||||
// Mock execCommand to return an error on Start
|
||||
oldExecCommand := execCommand
|
||||
execCommand = mockExecCommandError
|
||||
defer func() { execCommand = oldExecCommand }()
|
||||
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventsSvc := new(mocks.Service)
|
||||
eventsSvc.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return().Maybe()
|
||||
b := NewAlgorithm(logger, eventsSvc, "test", nil, "").(*binary)
|
||||
|
||||
err := b.Run()
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func mockExecCommandError(command string, args ...string) *exec.Cmd {
|
||||
// This will make Start() fail if we use a non-existent binary
|
||||
return exec.Command("non_existent_binary_for_sure_12345")
|
||||
}
|
||||
|
||||
func TestHelperProcess(t *testing.T) {
|
||||
if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
|
||||
return
|
||||
}
|
||||
os.Exit(0)
|
||||
}
|
||||
@@ -1,184 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package docker
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path"
|
||||
|
||||
"github.com/docker/docker/api/types/container"
|
||||
"github.com/docker/docker/api/types/image"
|
||||
"github.com/docker/docker/api/types/mount"
|
||||
"github.com/docker/docker/client"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
)
|
||||
|
||||
const (
|
||||
containerName = "agent_container"
|
||||
datasetsMountPath = "/cocos/datasets"
|
||||
resultsMountPath = "/cocos/results"
|
||||
)
|
||||
|
||||
var _ algorithm.Algorithm = (*docker)(nil)
|
||||
|
||||
type docker struct {
|
||||
algoFile string
|
||||
logger *slog.Logger
|
||||
stderr io.Writer
|
||||
stdout io.Writer
|
||||
cmpID string
|
||||
}
|
||||
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile, cmpID string) algorithm.Algorithm {
|
||||
d := &docker{
|
||||
algoFile: algoFile,
|
||||
logger: logger,
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
|
||||
stdout: &logging.Stdout{Logger: logger},
|
||||
cmpID: cmpID,
|
||||
}
|
||||
|
||||
return d
|
||||
}
|
||||
|
||||
func (d *docker) Run() error {
|
||||
// Create a new Docker client.
|
||||
cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not create a new Docker client: %v", err)
|
||||
}
|
||||
|
||||
// Open the Docker image tar file.
|
||||
imageFile, err := os.Open(d.algoFile)
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not open Docker image: %v", err)
|
||||
}
|
||||
defer imageFile.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
// Load the Docker image from the tar file.
|
||||
resp, err := cli.ImageLoad(ctx, imageFile, client.ImageLoadWithQuiet(true))
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not load Docker image from file: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// List the loaded images to get the image ID.
|
||||
images, err := cli.ImageList(ctx, image.ListOptions{})
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not get the Docker image list: %v", err)
|
||||
}
|
||||
|
||||
var imageID string = ""
|
||||
var dockerImageName string = ""
|
||||
for _, image := range images {
|
||||
for _, tag := range image.RepoTags {
|
||||
imageID = image.ID
|
||||
dockerImageName = tag
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if imageID == "" {
|
||||
return fmt.Errorf("could not find image ID")
|
||||
}
|
||||
|
||||
// Create and start the container.
|
||||
respContainer, err := cli.ContainerCreate(ctx, &container.Config{
|
||||
Image: dockerImageName,
|
||||
Tty: true,
|
||||
AttachStdout: true,
|
||||
AttachStderr: true,
|
||||
}, &container.HostConfig{
|
||||
Mounts: []mount.Mount{
|
||||
{
|
||||
Type: mount.TypeBind,
|
||||
Source: path.Join(algorithm.AlgoWorkingDir, algorithm.DatasetsDir),
|
||||
Target: datasetsMountPath,
|
||||
},
|
||||
{
|
||||
Type: mount.TypeBind,
|
||||
Source: path.Join(algorithm.AlgoWorkingDir, algorithm.ResultsDir),
|
||||
Target: resultsMountPath,
|
||||
},
|
||||
},
|
||||
}, nil, nil, fmt.Sprintf("%s-%s", containerName, d.cmpID))
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not create a Docker container: %v", err)
|
||||
}
|
||||
|
||||
if err := cli.ContainerStart(ctx, respContainer.ID, container.StartOptions{}); err != nil {
|
||||
return fmt.Errorf("could not start a Docker container: %v", err)
|
||||
}
|
||||
|
||||
stdout, err := cli.ContainerLogs(ctx, respContainer.ID, container.LogsOptions{ShowStdout: true, Follow: true})
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not read stdout from the container: %v", err)
|
||||
}
|
||||
defer stdout.Close()
|
||||
|
||||
go func() {
|
||||
if err := writeToOut(stdout, d.stdout); err != nil {
|
||||
d.logger.Warn(fmt.Sprintf("could not write to stdout: %v", err))
|
||||
}
|
||||
}()
|
||||
|
||||
stderr, err := cli.ContainerLogs(ctx, respContainer.ID, container.LogsOptions{ShowStderr: true, Follow: true})
|
||||
if err != nil {
|
||||
d.logger.Warn(fmt.Sprintf("could not read stderr from the container: %v", err))
|
||||
}
|
||||
defer stderr.Close()
|
||||
|
||||
go func() {
|
||||
if err := writeToOut(stderr, d.stderr); err != nil {
|
||||
d.logger.Warn(fmt.Sprintf("could not write to stderr: %v", err))
|
||||
}
|
||||
}()
|
||||
|
||||
statusCh, errCh := cli.ContainerWait(ctx, respContainer.ID, container.WaitConditionNotRunning)
|
||||
select {
|
||||
case err := <-errCh:
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not wait for a Docker container: %v", err)
|
||||
}
|
||||
case <-statusCh:
|
||||
}
|
||||
|
||||
defer func() {
|
||||
if err = cli.ContainerRemove(ctx, respContainer.ID, container.RemoveOptions{Force: true}); err != nil {
|
||||
d.logger.Warn(fmt.Sprintf("error could not remove container: %v", err))
|
||||
}
|
||||
|
||||
if _, err := cli.ImageRemove(ctx, imageID, image.RemoveOptions{Force: true}); err != nil {
|
||||
d.logger.Warn(fmt.Sprintf("error could not remove image: %v", err))
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func writeToOut(readCloser io.ReadCloser, ioWriter io.Writer) error {
|
||||
scanner := bufio.NewScanner(readCloser)
|
||||
for scanner.Scan() {
|
||||
if _, err := ioWriter.Write(scanner.Bytes()); err != nil {
|
||||
return fmt.Errorf("error writing to output: %v", err)
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return fmt.Errorf("error reading container logs error: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *docker) Stop() error {
|
||||
// To be supported later.
|
||||
return nil
|
||||
}
|
||||
@@ -1,29 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package docker
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
|
||||
"github.com/ultravioletrs/cocos/agent/events/mocks"
|
||||
)
|
||||
|
||||
// TestNewAlgorithm tests the NewAlgorithm function.
|
||||
func TestNewAlgorithm(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventsSvc := new(mocks.Service)
|
||||
algoFile := "/path/to/algo.tar"
|
||||
|
||||
algo := NewAlgorithm(logger, eventsSvc, algoFile, "")
|
||||
|
||||
d, ok := algo.(*docker)
|
||||
assert.True(t, ok, "NewAlgorithm should return a *docker")
|
||||
assert.Equal(t, algoFile, d.algoFile, "algoFile should be set correctly")
|
||||
assert.NotNil(t, d.logger, "logger should be set")
|
||||
assert.IsType(t, &logging.Stderr{}, d.stderr, "stderr should be of type *algorithm.Stderr")
|
||||
assert.IsType(t, &logging.Stdout{}, d.stdout, "stdout should be of type *algorithm.Stdout")
|
||||
}
|
||||
@@ -1,77 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package logging
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
)
|
||||
|
||||
var (
|
||||
_ io.Writer = &Stdout{}
|
||||
_ io.Writer = &Stderr{}
|
||||
)
|
||||
|
||||
const (
|
||||
bufSize = 1024
|
||||
algorithmRun = "AlgorithmRun"
|
||||
warningStatus = "Warning"
|
||||
)
|
||||
|
||||
type Stdout struct {
|
||||
Logger *slog.Logger
|
||||
}
|
||||
|
||||
// Write implements io.Writer.
|
||||
func (s *Stdout) Write(p []byte) (n int, err error) {
|
||||
inBuf := bytes.NewBuffer(p)
|
||||
|
||||
buf := make([]byte, bufSize)
|
||||
|
||||
for {
|
||||
n, err := inBuf.Read(buf)
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
return len(p) - inBuf.Len(), err
|
||||
}
|
||||
|
||||
s.Logger.Debug(string(buf[:n]))
|
||||
}
|
||||
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
type Stderr struct {
|
||||
Logger *slog.Logger
|
||||
EventSvc events.Service
|
||||
CmpID string
|
||||
}
|
||||
|
||||
// Write implements io.Writer.
|
||||
func (s *Stderr) Write(p []byte) (n int, err error) {
|
||||
inBuf := bytes.NewBuffer(p)
|
||||
|
||||
buf := make([]byte, bufSize)
|
||||
|
||||
for {
|
||||
n, err := inBuf.Read(buf)
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
return len(p) - inBuf.Len(), err
|
||||
}
|
||||
|
||||
s.Logger.Error(string(buf[:n]))
|
||||
}
|
||||
|
||||
s.EventSvc.SendEvent(s.CmpID, algorithmRun, warningStatus, json.RawMessage{})
|
||||
|
||||
return len(p), nil
|
||||
}
|
||||
@@ -1,86 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package logging
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
mglog "github.com/absmach/magistrala/logger"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent/events/mocks"
|
||||
"github.com/ultravioletrs/cocos/pkg/manager"
|
||||
)
|
||||
|
||||
func TestStdoutWrite(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected []string
|
||||
}{
|
||||
{
|
||||
name: "Single line",
|
||||
input: "Hello, World!",
|
||||
expected: []string{"Hello, World!"},
|
||||
},
|
||||
{
|
||||
name: "Multiple lines",
|
||||
input: "Line 1\nLine 2\nLine 3",
|
||||
expected: []string{"Line 1\nLine 2\nLine 3"},
|
||||
},
|
||||
{
|
||||
name: "Long input",
|
||||
input: strings.Repeat("a", bufSize+100),
|
||||
expected: []string{strings.Repeat("a", bufSize), strings.Repeat("a", 100)},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
stdout := &Stdout{Logger: mglog.NewMock()}
|
||||
n, err := stdout.Write([]byte(tt.input))
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, len(tt.input), n)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStderrWrite(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected []string
|
||||
}{
|
||||
{
|
||||
name: "Single line",
|
||||
input: "Error: Something went wrong",
|
||||
expected: []string{"Error: Something went wrong"},
|
||||
},
|
||||
{
|
||||
name: "Multiple lines",
|
||||
input: "Error 1\nError 2\nError 3",
|
||||
expected: []string{"Error 1\nError 2\nError 3"},
|
||||
},
|
||||
{
|
||||
name: "Long input",
|
||||
input: strings.Repeat("e", bufSize+100),
|
||||
expected: []string{strings.Repeat("e", bufSize), strings.Repeat("e", 100)},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockEventService := mocks.NewService(t)
|
||||
mockEventService.On("SendEvent", mock.Anything, "AlgorithmRun", manager.Warning.String(), mock.Anything).Return(nil)
|
||||
|
||||
stderr := &Stderr{Logger: mglog.NewMock(), EventSvc: mockEventService}
|
||||
n, err := stderr.Write([]byte(tt.input))
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, len(tt.input), n)
|
||||
mockEventService.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,127 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
// NewAlgorithm creates a new instance of Algorithm. 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 NewAlgorithm(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *Algorithm {
|
||||
mock := &Algorithm{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
|
||||
// Algorithm is an autogenerated mock type for the Algorithm type
|
||||
type Algorithm struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type Algorithm_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *Algorithm) EXPECT() *Algorithm_Expecter {
|
||||
return &Algorithm_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// Run provides a mock function for the type Algorithm
|
||||
func (_mock *Algorithm) Run() error {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Run")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func() error); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// Algorithm_Run_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Run'
|
||||
type Algorithm_Run_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Run is a helper method to define mock.On call
|
||||
func (_e *Algorithm_Expecter) Run() *Algorithm_Run_Call {
|
||||
return &Algorithm_Run_Call{Call: _e.mock.On("Run")}
|
||||
}
|
||||
|
||||
func (_c *Algorithm_Run_Call) Run(run func()) *Algorithm_Run_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Algorithm_Run_Call) Return(err error) *Algorithm_Run_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Algorithm_Run_Call) RunAndReturn(run func() error) *Algorithm_Run_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Stop provides a mock function for the type Algorithm
|
||||
func (_mock *Algorithm) Stop() error {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Stop")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func() error); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// Algorithm_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
|
||||
type Algorithm_Stop_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Stop is a helper method to define mock.On call
|
||||
func (_e *Algorithm_Expecter) Stop() *Algorithm_Stop_Call {
|
||||
return &Algorithm_Stop_Call{Call: _e.mock.On("Stop")}
|
||||
}
|
||||
|
||||
func (_c *Algorithm_Stop_Call) Run(run func()) *Algorithm_Stop_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Algorithm_Stop_Call) Return(err error) *Algorithm_Stop_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Algorithm_Stop_Call) RunAndReturn(run func() error) *Algorithm_Stop_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -1,133 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package python
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
const (
|
||||
PyRuntime = "python3"
|
||||
PyRuntimeKey = "python_runtime"
|
||||
)
|
||||
|
||||
func PythonRunTimeToContext(ctx context.Context, runtime string) context.Context {
|
||||
return metadata.AppendToOutgoingContext(ctx, PyRuntimeKey, runtime)
|
||||
}
|
||||
|
||||
func PythonRunTimeFromContext(ctx context.Context) string {
|
||||
return metadata.ValueFromIncomingContext(ctx, PyRuntimeKey)[0]
|
||||
}
|
||||
|
||||
var _ algorithm.Algorithm = (*python)(nil)
|
||||
|
||||
type python struct {
|
||||
algoFile string
|
||||
stderr io.Writer
|
||||
stdout io.Writer
|
||||
runtime string
|
||||
requirementsFile string
|
||||
args []string
|
||||
cmd *exec.Cmd
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, runtime, requirementsFile, algoFile string, args []string, cmpID string) algorithm.Algorithm {
|
||||
p := &python{
|
||||
algoFile: algoFile,
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
|
||||
stdout: &logging.Stdout{Logger: logger},
|
||||
requirementsFile: requirementsFile,
|
||||
args: args,
|
||||
}
|
||||
if runtime != "" {
|
||||
p.runtime = runtime
|
||||
} else {
|
||||
p.runtime = PyRuntime
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *python) Run() error {
|
||||
venvPath := "venv"
|
||||
defer func() {
|
||||
if err := os.RemoveAll(venvPath); err != nil {
|
||||
_, _ = p.stderr.Write([]byte(fmt.Sprintf("error removing virtual environment: %v\n", err)))
|
||||
}
|
||||
}()
|
||||
|
||||
createVenvCmd := exec.Command(p.runtime, "-m", "venv", venvPath)
|
||||
createVenvCmd.Stderr = p.stderr
|
||||
createVenvCmd.Stdout = p.stdout
|
||||
if err := createVenvCmd.Run(); err != nil {
|
||||
return fmt.Errorf("error creating virtual environment: %v", err)
|
||||
}
|
||||
|
||||
pythonPath := filepath.Join(venvPath, "bin", "python")
|
||||
|
||||
updatePipCmd := exec.Command(pythonPath, "-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel")
|
||||
updatePipCmd.Stderr = p.stderr
|
||||
updatePipCmd.Stdout = p.stdout
|
||||
if err := updatePipCmd.Run(); err != nil {
|
||||
return fmt.Errorf("error updating pip, setuptools and wheel: %v", err)
|
||||
}
|
||||
|
||||
if p.requirementsFile != "" {
|
||||
rcmd := exec.Command(pythonPath, "-m", "pip", "install", "-r", p.requirementsFile)
|
||||
rcmd.Stderr = p.stderr
|
||||
rcmd.Stdout = p.stdout
|
||||
if err := rcmd.Run(); err != nil {
|
||||
return fmt.Errorf("error installing requirements: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
args := append([]string{p.algoFile}, p.args...)
|
||||
p.mu.Lock()
|
||||
p.cmd = exec.Command(pythonPath, args...)
|
||||
p.cmd.Stderr = p.stderr
|
||||
p.cmd.Stdout = p.stdout
|
||||
|
||||
if err := p.cmd.Start(); err != nil {
|
||||
p.mu.Unlock()
|
||||
return fmt.Errorf("error starting algorithm: %v", err)
|
||||
}
|
||||
p.mu.Unlock()
|
||||
|
||||
if err := p.cmd.Wait(); err != nil {
|
||||
return fmt.Errorf("algorithm execution error: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *python) Stop() error {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
if p.cmd == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if p.cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := p.cmd.Process.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) {
|
||||
return fmt.Errorf("error stopping algorithm: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,242 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package python
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
|
||||
"github.com/ultravioletrs/cocos/agent/events/mocks"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
const runtime = "python3"
|
||||
|
||||
func TestPythonRunTimeToContext(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
newCtx := PythonRunTimeToContext(ctx, runtime)
|
||||
|
||||
md, ok := metadata.FromOutgoingContext(newCtx)
|
||||
if !ok {
|
||||
t.Fatal("Expected metadata in context")
|
||||
}
|
||||
|
||||
values := md.Get(PyRuntimeKey)
|
||||
if len(values) != 1 || values[0] != runtime {
|
||||
t.Errorf("Expected runtime %s, got %v", runtime, values)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPythonRunTimeFromContext(t *testing.T) {
|
||||
ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(PyRuntimeKey, runtime))
|
||||
|
||||
got := PythonRunTimeFromContext(ctx)
|
||||
if got != runtime {
|
||||
t.Errorf("Expected runtime %s, got %s", runtime, got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAlgorithm(t *testing.T) {
|
||||
logger := &slog.Logger{}
|
||||
eventsSvc := new(mocks.Service)
|
||||
requirementsFile := "requirements.txt"
|
||||
algoFile := "algorithm.py"
|
||||
args := []string{"--arg1", "value1"}
|
||||
|
||||
algo := NewAlgorithm(logger, eventsSvc, runtime, requirementsFile, algoFile, args, "")
|
||||
|
||||
p, ok := algo.(*python)
|
||||
if !ok {
|
||||
t.Fatal("Expected *python type")
|
||||
}
|
||||
|
||||
if p.runtime != runtime {
|
||||
t.Errorf("Expected runtime %s, got %s", runtime, p.runtime)
|
||||
}
|
||||
if p.requirementsFile != requirementsFile {
|
||||
t.Errorf("Expected requirementsFile %s, got %s", requirementsFile, p.requirementsFile)
|
||||
}
|
||||
if p.algoFile != algoFile {
|
||||
t.Errorf("Expected algoFile %s, got %s", algoFile, p.algoFile)
|
||||
}
|
||||
if len(p.args) != len(args) {
|
||||
t.Errorf("Expected %d args, got %d", len(args), len(p.args))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "python-test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
scriptContent := []byte("print('Hello, World!')")
|
||||
scriptPath := filepath.Join(tmpDir, "test_script.py")
|
||||
if err := os.WriteFile(scriptPath, scriptContent, 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
eventsSvc := new(mocks.Service)
|
||||
eventsSvc.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return().Maybe()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
|
||||
algo := &python{
|
||||
algoFile: scriptPath,
|
||||
stderr: io.MultiWriter(&stderr, &logging.Stderr{Logger: slog.Default(), EventSvc: eventsSvc}),
|
||||
stdout: io.MultiWriter(&stdout, &logging.Stdout{Logger: slog.Default()}),
|
||||
runtime: "python3",
|
||||
}
|
||||
|
||||
err = algo.Run()
|
||||
if err != nil {
|
||||
t.Fatalf("Unexpected error: %v", err)
|
||||
}
|
||||
|
||||
expectedOutput := "Hello, World!\n"
|
||||
if !strings.Contains(stdout.String(), expectedOutput) {
|
||||
t.Errorf("Expected output to contain %q, got %q", expectedOutput, stdout.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunWithRequirements(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "python-test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
scriptContent := []byte("import requests\nprint(requests.__version__)")
|
||||
scriptPath := filepath.Join(tmpDir, "test_script.py")
|
||||
if err := os.WriteFile(scriptPath, scriptContent, 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
requirementsContent := []byte("requests==2.26.0")
|
||||
requirementsPath := filepath.Join(tmpDir, "requirements.txt")
|
||||
if err := os.WriteFile(requirementsPath, requirementsContent, 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
eventsSvc := new(mocks.Service)
|
||||
eventsSvc.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return().Maybe()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
|
||||
algo := &python{
|
||||
algoFile: scriptPath,
|
||||
requirementsFile: requirementsPath,
|
||||
stderr: io.MultiWriter(&stderr, &logging.Stderr{Logger: slog.Default(), EventSvc: eventsSvc}),
|
||||
stdout: io.MultiWriter(&stdout, &logging.Stdout{Logger: slog.Default()}),
|
||||
runtime: "python3",
|
||||
}
|
||||
|
||||
err = algo.Run()
|
||||
if err != nil {
|
||||
t.Fatalf("Unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if !strings.Contains(stdout.String(), "2.26.0") {
|
||||
t.Errorf("Expected output to contain requests version 2.26.0, got %q", stdout.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestStop(t *testing.T) {
|
||||
t.Run("stop nil cmd", func(t *testing.T) {
|
||||
p := &python{}
|
||||
err := p.Stop()
|
||||
if err != nil {
|
||||
t.Errorf("Expected nil error, got %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("stop with running process", func(t *testing.T) {
|
||||
p := &python{
|
||||
stderr: io.Discard,
|
||||
stdout: io.Discard,
|
||||
}
|
||||
|
||||
p.cmd = exec.Command("python3", "-c", "import time; time.sleep(10)")
|
||||
if err := p.cmd.Start(); err != nil {
|
||||
t.Fatalf("Failed to start command: %v", err)
|
||||
}
|
||||
|
||||
err := p.Stop()
|
||||
if err != nil {
|
||||
t.Errorf("Expected nil error, got %v", err)
|
||||
}
|
||||
|
||||
// Verify it actually stopped
|
||||
_ = p.cmd.Wait()
|
||||
})
|
||||
|
||||
t.Run("stop already exited", func(t *testing.T) {
|
||||
p := &python{}
|
||||
p.cmd = exec.Command("python3", "-c", "print(1)")
|
||||
if err := p.cmd.Run(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
err := p.Stop()
|
||||
if err != nil {
|
||||
t.Errorf("Expected nil error, got %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestRun_Errors(t *testing.T) {
|
||||
t.Run("invalid runtime error", func(t *testing.T) {
|
||||
algo := &python{
|
||||
algoFile: "algo.py",
|
||||
runtime: "non-existent-python",
|
||||
stderr: io.Discard,
|
||||
stdout: io.Discard,
|
||||
}
|
||||
err := algo.Run()
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "error creating virtual environment")
|
||||
})
|
||||
|
||||
t.Run("pip install failure", func(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "python-err-test")
|
||||
require.NoError(t, err)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
scriptPath := filepath.Join(tmpDir, "test.py")
|
||||
require.NoError(t, os.WriteFile(scriptPath, []byte("print(1)"), 0o644))
|
||||
|
||||
reqPath := filepath.Join(tmpDir, "requirements.txt")
|
||||
require.NoError(t, os.WriteFile(reqPath, []byte("non-existent-package==9.9.9"), 0o644))
|
||||
|
||||
algo := &python{
|
||||
algoFile: scriptPath,
|
||||
requirementsFile: reqPath,
|
||||
runtime: "python3",
|
||||
stderr: io.Discard,
|
||||
stdout: io.Discard,
|
||||
}
|
||||
err = algo.Run()
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "error installing requirements")
|
||||
})
|
||||
}
|
||||
|
||||
func TestNewAlgorithmEmptyRuntime(t *testing.T) {
|
||||
eventsSvc := new(mocks.Service)
|
||||
algo := NewAlgorithm(slog.Default(), eventsSvc, "", "req.txt", "algo.py", nil, "")
|
||||
p := algo.(*python)
|
||||
if p.runtime != PyRuntime {
|
||||
t.Errorf("Expected default runtime %s, got %s", PyRuntime, p.runtime)
|
||||
}
|
||||
}
|
||||
@@ -1,82 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package algorithm_test
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm"
|
||||
"github.com/ultravioletrs/cocos/internal"
|
||||
)
|
||||
|
||||
func TestZipDirectory(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
directories []string
|
||||
files []string
|
||||
expected []string
|
||||
}{
|
||||
{
|
||||
name: "empty directory",
|
||||
directories: []string{"testdata"},
|
||||
},
|
||||
{
|
||||
name: "single file",
|
||||
files: []string{"file1.txt"},
|
||||
},
|
||||
{
|
||||
name: "directory with single file",
|
||||
directories: []string{"testdata"},
|
||||
expected: []string{"testdata/file1.txt"},
|
||||
},
|
||||
{
|
||||
name: "directory with multiple files",
|
||||
directories: []string{"testdata"},
|
||||
expected: []string{
|
||||
"testdata/file1.txt",
|
||||
"testdata/file2.txt",
|
||||
"testdata/file3.txt",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "nested directories",
|
||||
directories: []string{"testdata", "testdata/nested"},
|
||||
expected: []string{
|
||||
"testdata/nested/file1.txt",
|
||||
"testdata/nested/file2.txt",
|
||||
"testdata/nested/file3.txt",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if err := os.Mkdir(algorithm.ResultsDir, 0o755); err != nil {
|
||||
t.Fatalf("error creating results directory: %s", err.Error())
|
||||
}
|
||||
defer func() {
|
||||
if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
|
||||
t.Fatalf("error removing results directory and its contents: %s", err.Error())
|
||||
}
|
||||
}()
|
||||
|
||||
for _, dir := range tc.directories {
|
||||
if dir != "" {
|
||||
if err := os.Mkdir(algorithm.ResultsDir+"/"+dir, 0o755); err != nil {
|
||||
t.Fatalf("error creating test directory: %s", err.Error())
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, file := range tc.files {
|
||||
if _, err := os.Create(algorithm.ResultsDir + "/" + file); err != nil {
|
||||
t.Fatalf("error creating test file: %s", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := internal.ZipDirectoryToMemory(algorithm.ResultsDir); err != nil {
|
||||
t.Errorf("ZipDirectory() error = %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,83 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package wasm
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"sync"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
)
|
||||
|
||||
var execCommand = exec.Command
|
||||
|
||||
const wasmRuntime = "wasmedge"
|
||||
|
||||
var mapDirOption = []string{"--dir", ".:" + algorithm.ResultsDir}
|
||||
|
||||
var _ algorithm.Algorithm = (*wasm)(nil)
|
||||
|
||||
type wasm struct {
|
||||
algoFile string
|
||||
stderr io.Writer
|
||||
stdout io.Writer
|
||||
args []string
|
||||
cmd *exec.Cmd
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, args []string, algoFile, cmpID string) algorithm.Algorithm {
|
||||
return &wasm{
|
||||
algoFile: algoFile,
|
||||
stderr: &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
|
||||
stdout: &logging.Stdout{Logger: logger},
|
||||
args: args,
|
||||
}
|
||||
}
|
||||
|
||||
func (w *wasm) Run() error {
|
||||
args := append(mapDirOption, w.algoFile)
|
||||
args = append(args, w.args...)
|
||||
w.mu.Lock()
|
||||
w.cmd = execCommand(wasmRuntime, args...)
|
||||
w.cmd.Stderr = w.stderr
|
||||
w.cmd.Stdout = w.stdout
|
||||
|
||||
if err := w.cmd.Start(); err != nil {
|
||||
w.mu.Unlock()
|
||||
return fmt.Errorf("error starting algorithm: %v", err)
|
||||
}
|
||||
w.mu.Unlock()
|
||||
|
||||
if err := w.cmd.Wait(); err != nil {
|
||||
return fmt.Errorf("algorithm execution error: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *wasm) Stop() error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
if w.cmd == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if w.cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := w.cmd.Process.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) {
|
||||
return fmt.Errorf("error stopping algorithm: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,179 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package wasm
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/logging"
|
||||
"github.com/ultravioletrs/cocos/agent/events/mocks"
|
||||
)
|
||||
|
||||
const testWasm = "test.wasm"
|
||||
|
||||
func TestNewAlgorithm(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventsSvc := new(mocks.Service)
|
||||
algoFile := testWasm
|
||||
args := []string{"arg1", "arg2"}
|
||||
|
||||
algo := NewAlgorithm(logger, eventsSvc, args, algoFile, "")
|
||||
|
||||
w, ok := algo.(*wasm)
|
||||
if !ok {
|
||||
t.Fatalf("NewAlgorithm did not return a *wasm")
|
||||
}
|
||||
|
||||
if w.algoFile != algoFile {
|
||||
t.Errorf("Expected algoFile to be %s, got %s", algoFile, w.algoFile)
|
||||
}
|
||||
|
||||
if len(w.args) != len(args) {
|
||||
t.Errorf("Expected %d args, got %d", len(args), len(w.args))
|
||||
}
|
||||
|
||||
_, ok = w.stderr.(*logging.Stderr)
|
||||
if !ok {
|
||||
t.Errorf("Expected stderr to be *algorithm.Stderr")
|
||||
}
|
||||
|
||||
_, ok = w.stdout.(*logging.Stdout)
|
||||
if !ok {
|
||||
t.Errorf("Expected stdout to be *algorithm.Stdout")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunError(t *testing.T) {
|
||||
// Mock exec.Command to return an error
|
||||
execCommand = mockExecCommandError
|
||||
defer func() { execCommand = exec.Command }()
|
||||
|
||||
algoFile := testWasm
|
||||
args := []string{"arg1", "arg2"}
|
||||
|
||||
w := &wasm{
|
||||
algoFile: algoFile,
|
||||
args: args,
|
||||
stderr: os.Stderr, // Use real stderr or io.Discard
|
||||
stdout: os.Stdout,
|
||||
}
|
||||
|
||||
err := w.Run()
|
||||
|
||||
if err == nil {
|
||||
t.Errorf("Run() should have returned an error")
|
||||
}
|
||||
}
|
||||
|
||||
func mockExecCommand(command string, args ...string) *exec.Cmd {
|
||||
cs := []string{"-test.run=TestHelperProcess", "--", command}
|
||||
cs = append(cs, args...)
|
||||
cmd := exec.Command(os.Args[0], cs...)
|
||||
cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1"}
|
||||
return cmd
|
||||
}
|
||||
|
||||
func mockExecCommandError(command string, args ...string) *exec.Cmd {
|
||||
cmd := mockExecCommand(command, args...)
|
||||
cmd.Env = append(cmd.Env, "GO_WANT_HELPER_PROCESS_ERROR=1")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func TestStop(t *testing.T) {
|
||||
t.Run("stop nil cmd", func(t *testing.T) {
|
||||
w := &wasm{}
|
||||
err := w.Stop()
|
||||
if err != nil {
|
||||
t.Errorf("Expected nil error, got %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("stop with running process", func(t *testing.T) {
|
||||
oldExecCommand := execCommand
|
||||
execCommand = mockExecCommand
|
||||
defer func() { execCommand = oldExecCommand }()
|
||||
|
||||
w := &wasm{
|
||||
algoFile: testWasm,
|
||||
stdout: os.Stdout,
|
||||
stderr: os.Stderr,
|
||||
}
|
||||
|
||||
// We need to simulate a running process.
|
||||
// mockExecCommand returns a command that runs TestHelperProcess.
|
||||
// If we don't call Wait(), it keeps running? No, TestHelperProcess exits immediately.
|
||||
// Let's modify TestHelperProcess to sleep if an env var is set.
|
||||
|
||||
w.cmd = mockExecCommand("sleep", "10")
|
||||
w.cmd.Env = append(w.cmd.Env, "GO_WANT_HELPER_PROCESS_SLEEP=1")
|
||||
if err := w.cmd.Start(); err != nil {
|
||||
t.Fatalf("Failed to start command: %v", err)
|
||||
}
|
||||
|
||||
err := w.Stop()
|
||||
if err != nil {
|
||||
t.Errorf("Expected nil error, got %v", err)
|
||||
}
|
||||
_ = w.cmd.Wait()
|
||||
})
|
||||
}
|
||||
|
||||
func TestStopAlreadyExited(t *testing.T) {
|
||||
oldExecCommand := execCommand
|
||||
execCommand = mockExecCommand
|
||||
defer func() { execCommand = oldExecCommand }()
|
||||
|
||||
w := &wasm{
|
||||
algoFile: testWasm,
|
||||
stdout: os.Stdout,
|
||||
stderr: os.Stderr,
|
||||
}
|
||||
|
||||
w.cmd = mockExecCommand("true")
|
||||
if err := w.cmd.Run(); err != nil {
|
||||
t.Fatalf("Failed to run command: %v", err)
|
||||
}
|
||||
|
||||
err := w.Stop()
|
||||
if err != nil {
|
||||
t.Errorf("Expected nil error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunSuccess(t *testing.T) {
|
||||
oldExecCommand := execCommand
|
||||
execCommand = mockExecCommand
|
||||
defer func() { execCommand = oldExecCommand }()
|
||||
|
||||
algoFile := testWasm
|
||||
args := []string{"arg1", "arg2"}
|
||||
|
||||
w := &wasm{
|
||||
algoFile: algoFile,
|
||||
args: args,
|
||||
stderr: os.Stderr,
|
||||
stdout: os.Stdout,
|
||||
}
|
||||
|
||||
err := w.Run()
|
||||
if err != nil {
|
||||
t.Errorf("Run() returned unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHelperProcess(t *testing.T) {
|
||||
if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
|
||||
return
|
||||
}
|
||||
if os.Getenv("GO_WANT_HELPER_PROCESS_SLEEP") == "1" {
|
||||
time.Sleep(10 * time.Second)
|
||||
}
|
||||
if os.Getenv("GO_WANT_HELPER_PROCESS_ERROR") == "1" {
|
||||
os.Exit(1)
|
||||
}
|
||||
os.Exit(0)
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/go-kit/kit/endpoint"
|
||||
kitgrpc "github.com/go-kit/kit/transport/grpc"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"google.golang.org/grpc"
|
||||
)
|
||||
|
||||
const svcName = "agent.AgentService"
|
||||
|
||||
type grpcClient struct {
|
||||
algo endpoint.Endpoint
|
||||
data endpoint.Endpoint
|
||||
result endpoint.Endpoint
|
||||
attestation endpoint.Endpoint
|
||||
timeout time.Duration
|
||||
}
|
||||
|
||||
// NewClient returns new gRPC client instance.
|
||||
func NewClient(conn *grpc.ClientConn, timeout time.Duration) agent.AgentServiceClient {
|
||||
return &grpcClient{
|
||||
algo: kitgrpc.NewClient(
|
||||
conn,
|
||||
svcName,
|
||||
"Algo",
|
||||
encodeAlgoRequest,
|
||||
decodeAlgoResponse,
|
||||
agent.AlgoResponse{},
|
||||
).Endpoint(),
|
||||
data: kitgrpc.NewClient(
|
||||
conn,
|
||||
svcName,
|
||||
"Data",
|
||||
encodeDataRequest,
|
||||
decodeDataResponse,
|
||||
agent.DataResponse{},
|
||||
).Endpoint(),
|
||||
result: kitgrpc.NewClient(
|
||||
conn,
|
||||
svcName,
|
||||
"Result",
|
||||
encodeResultRequest,
|
||||
decodeResultResponse,
|
||||
agent.ResultResponse{},
|
||||
).Endpoint(),
|
||||
attestation: kitgrpc.NewClient(
|
||||
conn,
|
||||
svcName,
|
||||
"Attestation",
|
||||
encodeAttestationRequest,
|
||||
decodeAttestationResponse,
|
||||
agent.AttestationResponse{},
|
||||
).Endpoint(),
|
||||
timeout: timeout,
|
||||
}
|
||||
}
|
||||
|
||||
// encodeAlgoRequest is a transport/grpc.EncodeRequestFunc that
|
||||
// converts a user-domain algoReq to a gRPC request.
|
||||
func encodeAlgoRequest(_ context.Context, request interface{}) (interface{}, error) {
|
||||
req, ok := request.(*algoReq)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid request type: %T", request)
|
||||
}
|
||||
|
||||
return &agent.AlgoRequest{
|
||||
Algorithm: req.Algorithm,
|
||||
Provider: req.Provider,
|
||||
Id: req.Id,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// decodeAlgoResponse is a transport/grpc.DecodeResponseFunc that
|
||||
// converts a gRPC AlgoResponse to a user-domain response.
|
||||
func decodeAlgoResponse(_ context.Context, grpcResponse interface{}) (interface{}, error) {
|
||||
_, ok := grpcResponse.(*agent.AlgoResponse)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid response type: %T", grpcResponse)
|
||||
}
|
||||
|
||||
return algoRes{}, nil
|
||||
}
|
||||
|
||||
// encodeDataRequest is a transport/grpc.EncodeRequestFunc that
|
||||
// converts a user-domain dataReq to a gRPC request.
|
||||
func encodeDataRequest(_ context.Context, request interface{}) (interface{}, error) {
|
||||
req, ok := request.(*dataReq)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid request type: %T", request)
|
||||
}
|
||||
|
||||
return &agent.DataRequest{
|
||||
Dataset: req.Dataset,
|
||||
Provider: req.Provider,
|
||||
Id: req.Id,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// decodeDataResponse is a transport/grpc.DecodeResponseFunc that
|
||||
// converts a gRPC DataResponse to a user-domain response.
|
||||
func decodeDataResponse(_ context.Context, grpcResponse interface{}) (interface{}, error) {
|
||||
_, ok := grpcResponse.(*agent.DataResponse)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid response type: %T", grpcResponse)
|
||||
}
|
||||
|
||||
return dataRes{}, nil
|
||||
}
|
||||
|
||||
// encodeResultRequest is a transport/grpc.EncodeRequestFunc that
|
||||
// converts a user-domain resultReq to a gRPC request.
|
||||
func encodeResultRequest(_ context.Context, request interface{}) (interface{}, error) {
|
||||
req, ok := request.(*resultReq)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid request type: %T", request)
|
||||
}
|
||||
|
||||
return &agent.ResultRequest{
|
||||
Consumer: req.Consumer,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// decodeResultResponse is a transport/grpc.DecodeResponseFunc that
|
||||
// converts a gRPC ResultResponse to a user-domain response.
|
||||
func decodeResultResponse(_ context.Context, grpcResponse interface{}) (interface{}, error) {
|
||||
response, ok := grpcResponse.(*agent.ResultResponse)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid response type: %T", grpcResponse)
|
||||
}
|
||||
|
||||
return resultRes{
|
||||
File: response.File,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// encodeAttestationRequest is a transport/grpc.EncodeRequestFunc that
|
||||
// converts a user-domain attestationReq to a gRPC request.
|
||||
func encodeAttestationRequest(_ context.Context, request interface{}) (interface{}, error) {
|
||||
req, ok := request.(*attestationReq)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid request type: %T", request)
|
||||
}
|
||||
return &agent.AttestationRequest{ReportData: req.ReportData[:]}, nil
|
||||
}
|
||||
|
||||
// decodeAttestationResponse is a transport/grpc.DecodeResponseFunc that
|
||||
// converts a gRPC AttestationResponse to a user-domain response.
|
||||
func decodeAttestationResponse(_ context.Context, grpcResponse interface{}) (interface{}, error) {
|
||||
response, ok := grpcResponse.(*agent.AttestationResponse)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid response type: %T", grpcResponse)
|
||||
}
|
||||
|
||||
return attestationRes{
|
||||
File: response.File,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Algo implements the Algo method of the agent.AgentServiceClient interface.
|
||||
func (c grpcClient) Algo(ctx context.Context, request *agent.AlgoRequest, _ ...grpc.CallOption) (*agent.AlgoResponse, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, c.timeout)
|
||||
defer cancel()
|
||||
|
||||
_, err := c.algo(ctx, &algoReq{Algorithm: request.Algorithm, Provider: request.Provider, Id: request.Id})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &agent.AlgoResponse{}, nil
|
||||
}
|
||||
|
||||
// Data implements the Data method of the agent.AgentServiceClient interface.
|
||||
func (c grpcClient) Data(ctx context.Context, request *agent.DataRequest, _ ...grpc.CallOption) (*agent.DataResponse, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, c.timeout)
|
||||
defer cancel()
|
||||
|
||||
_, err := c.data(ctx, &dataReq{Dataset: request.Dataset, Provider: request.Provider, Id: request.Id})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &agent.DataResponse{}, nil
|
||||
}
|
||||
|
||||
// Result implements the Result method of the agent.AgentServiceClient interface.
|
||||
func (c grpcClient) Result(ctx context.Context, request *agent.ResultRequest, _ ...grpc.CallOption) (*agent.ResultResponse, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, c.timeout)
|
||||
defer cancel()
|
||||
|
||||
res, err := c.result(ctx, &resultReq{Consumer: request.Consumer})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resultRes := res.(resultRes)
|
||||
return &agent.ResultResponse{File: resultRes.File}, nil
|
||||
}
|
||||
|
||||
// Result implements the Result method of the agent.AgentServiceClient interface.
|
||||
func (c grpcClient) Attestation(ctx context.Context, request *agent.AttestationRequest, _ ...grpc.CallOption) (*agent.AttestationResponse, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, c.timeout)
|
||||
defer cancel()
|
||||
|
||||
res, err := c.attestation(ctx, &attestationReq{ReportData: [agent.ReportDataSize]byte(request.ReportData)})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
attestationRes := res.(attestationRes)
|
||||
return &agent.AttestationResponse{File: attestationRes.File}, nil
|
||||
}
|
||||
@@ -7,18 +7,17 @@ import (
|
||||
|
||||
"github.com/go-kit/kit/endpoint"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
)
|
||||
|
||||
func algoEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
return func(ctx context.Context, request any) (any, error) {
|
||||
return func(ctx context.Context, request interface{}) (interface{}, error) {
|
||||
req := request.(algoReq)
|
||||
|
||||
if err := req.validate(); err != nil {
|
||||
return algoRes{}, err
|
||||
}
|
||||
|
||||
algo := agent.Algorithm{Algorithm: req.Algorithm, Requirements: req.Requirements}
|
||||
algo := agent.Algorithm{Algorithm: req.Algorithm, Provider: req.Provider, ID: req.Id}
|
||||
|
||||
err := svc.Algo(ctx, algo)
|
||||
if err != nil {
|
||||
@@ -30,14 +29,14 @@ func algoEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
}
|
||||
|
||||
func dataEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
return func(ctx context.Context, request any) (any, error) {
|
||||
return func(ctx context.Context, request interface{}) (interface{}, error) {
|
||||
req := request.(dataReq)
|
||||
|
||||
if err := req.validate(); err != nil {
|
||||
return dataRes{}, err
|
||||
}
|
||||
|
||||
dataset := agent.Dataset{Dataset: req.Dataset, Filename: req.Filename}
|
||||
dataset := agent.Dataset{Dataset: req.Dataset, Provider: req.Provider, ID: req.Id}
|
||||
|
||||
err := svc.Data(ctx, dataset)
|
||||
if err != nil {
|
||||
@@ -49,13 +48,13 @@ func dataEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
}
|
||||
|
||||
func resultEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
return func(ctx context.Context, request any) (any, error) {
|
||||
return func(ctx context.Context, request interface{}) (interface{}, error) {
|
||||
req := request.(resultReq)
|
||||
|
||||
if err := req.validate(); err != nil {
|
||||
return resultRes{}, err
|
||||
}
|
||||
file, err := svc.Result(ctx)
|
||||
file, err := svc.Result(ctx, req.Consumer)
|
||||
if err != nil {
|
||||
return resultRes{}, err
|
||||
}
|
||||
@@ -65,13 +64,13 @@ func resultEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
}
|
||||
|
||||
func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
return func(ctx context.Context, request any) (any, error) {
|
||||
return func(ctx context.Context, request interface{}) (interface{}, error) {
|
||||
req := request.(attestationReq)
|
||||
|
||||
if err := req.validate(); err != nil {
|
||||
return attestationRes{}, err
|
||||
}
|
||||
file, err := svc.Attestation(ctx, req.TeeNonce, req.VtpmNonce, attestation.PlatformType(req.AttType))
|
||||
file, err := svc.Attestation(ctx, req.ReportData)
|
||||
if err != nil {
|
||||
return attestationRes{}, err
|
||||
}
|
||||
@@ -79,33 +78,3 @@ func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
return attestationRes{File: file}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func imaMeasurementsEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
return func(ctx context.Context, request any) (any, error) {
|
||||
req := request.(imaMeasurementsReq)
|
||||
|
||||
if err := req.validate(); err != nil {
|
||||
return imaMeasurementsRes{}, err
|
||||
}
|
||||
file, pcr10, err := svc.IMAMeasurements(ctx)
|
||||
if err != nil {
|
||||
return imaMeasurementsRes{}, err
|
||||
}
|
||||
|
||||
return imaMeasurementsRes{File: file, PCR10: pcr10}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func azureAttestationTokenEndpoint(svc agent.Service) endpoint.Endpoint {
|
||||
return func(ctx context.Context, request any) (any, error) {
|
||||
req := request.(azureAttestationTokenReq)
|
||||
if err := req.validate(); err != nil {
|
||||
return fetchAttestationTokenRes{}, err
|
||||
}
|
||||
file, err := svc.AzureAttestationToken(ctx, req.tokenNonce)
|
||||
if err != nil {
|
||||
return fetchAttestationTokenRes{}, err
|
||||
}
|
||||
return fetchAttestationTokenRes{File: file}, nil
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,227 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/agent/mocks"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"golang.org/x/crypto/sha3"
|
||||
)
|
||||
|
||||
const svcErr = "Service Error"
|
||||
|
||||
func TestAlgoEndpoint(t *testing.T) {
|
||||
svc := new(mocks.Service)
|
||||
tests := []struct {
|
||||
name string
|
||||
req algoReq
|
||||
expectedErr bool
|
||||
}{
|
||||
{
|
||||
name: "Success",
|
||||
req: algoReq{Algorithm: []byte("algorithm")},
|
||||
},
|
||||
{
|
||||
name: "Validation Error",
|
||||
req: algoReq{},
|
||||
expectedErr: true,
|
||||
},
|
||||
{
|
||||
name: "Service Error",
|
||||
req: algoReq{Algorithm: []byte("algorithm")},
|
||||
expectedErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.name == svcErr {
|
||||
svc.On("Algo", context.Background(), agent.Algorithm{Algorithm: tt.req.Algorithm}).Return(errors.New("")).Once()
|
||||
} else {
|
||||
svc.On("Algo", context.Background(), agent.Algorithm{Algorithm: tt.req.Algorithm}).Return(nil).Once()
|
||||
}
|
||||
endpoint := algoEndpoint(svc)
|
||||
_, err := endpoint(context.Background(), tt.req)
|
||||
if (err != nil) != tt.expectedErr {
|
||||
t.Errorf("algoEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDataEndpoint(t *testing.T) {
|
||||
svc := new(mocks.Service)
|
||||
tests := []struct {
|
||||
name string
|
||||
req dataReq
|
||||
expectedErr bool
|
||||
}{
|
||||
{
|
||||
name: "Success",
|
||||
req: dataReq{Dataset: []byte("dataset")},
|
||||
},
|
||||
{
|
||||
name: "Validation Error",
|
||||
req: dataReq{},
|
||||
expectedErr: true,
|
||||
},
|
||||
{
|
||||
name: "Service Error",
|
||||
req: dataReq{Dataset: []byte("dataset")},
|
||||
expectedErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.name == svcErr {
|
||||
svc.On("Data", context.Background(), agent.Dataset{Dataset: tt.req.Dataset}).Return(errors.New("")).Once()
|
||||
} else {
|
||||
svc.On("Data", context.Background(), agent.Dataset{Dataset: tt.req.Dataset}).Return(nil).Once()
|
||||
}
|
||||
endpoint := dataEndpoint(svc)
|
||||
_, err := endpoint(context.Background(), tt.req)
|
||||
if (err != nil) != tt.expectedErr {
|
||||
t.Errorf("dataEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultEndpoint(t *testing.T) {
|
||||
svc := new(mocks.Service)
|
||||
tests := []struct {
|
||||
name string
|
||||
req resultReq
|
||||
expectedErr bool
|
||||
}{
|
||||
{
|
||||
name: "Success",
|
||||
req: resultReq{},
|
||||
},
|
||||
{
|
||||
name: "Service Error",
|
||||
req: resultReq{},
|
||||
expectedErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.name == svcErr {
|
||||
svc.On("Result", context.Background()).Return([]byte{}, errors.New("")).Once()
|
||||
} else {
|
||||
svc.On("Result", context.Background()).Return([]byte{}, nil).Once()
|
||||
}
|
||||
endpoint := resultEndpoint(svc)
|
||||
res, err := endpoint(context.Background(), tt.req)
|
||||
if (err != nil) != tt.expectedErr {
|
||||
t.Errorf("resultEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
|
||||
}
|
||||
if err == nil {
|
||||
_, ok := res.(resultRes)
|
||||
if !ok {
|
||||
t.Errorf("resultEndpoint() returned unexpected type %T", res)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAttestationEndpoint(t *testing.T) {
|
||||
svc := new(mocks.Service)
|
||||
tests := []struct {
|
||||
name string
|
||||
req attestationReq
|
||||
expectedErr bool
|
||||
}{
|
||||
{
|
||||
name: "Success",
|
||||
req: attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.SNP},
|
||||
},
|
||||
{
|
||||
name: "Service Error",
|
||||
req: attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.SNP},
|
||||
expectedErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.name == svcErr {
|
||||
svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, errors.New("")).Once()
|
||||
} else {
|
||||
svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, nil).Once()
|
||||
}
|
||||
endpoint := attestationEndpoint(svc)
|
||||
res, err := endpoint(context.Background(), tt.req)
|
||||
if (err != nil) != tt.expectedErr {
|
||||
t.Errorf("attestationEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
|
||||
}
|
||||
if err == nil {
|
||||
_, ok := res.(attestationRes)
|
||||
if !ok {
|
||||
t.Errorf("attestationEndpoint() returned unexpected type %T", res)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAttestationTokenEndpoint(t *testing.T) {
|
||||
svc := new(mocks.Service)
|
||||
tests := []struct {
|
||||
name string
|
||||
req azureAttestationTokenReq
|
||||
mockErr error
|
||||
expectedErr bool
|
||||
}{
|
||||
{
|
||||
name: "Success",
|
||||
req: azureAttestationTokenReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce"))},
|
||||
mockErr: nil,
|
||||
expectedErr: false,
|
||||
},
|
||||
{
|
||||
name: "Service Error",
|
||||
req: azureAttestationTokenReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce"))},
|
||||
mockErr: errors.New("mock failure"),
|
||||
expectedErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Only call service mock if validation is expected to pass
|
||||
if err := tt.req.validate(); err == nil {
|
||||
svc.On("AzureAttestationToken", mock.Anything, tt.req.tokenNonce).
|
||||
Return([]byte("mock file"), tt.mockErr).Once()
|
||||
}
|
||||
|
||||
endpoint := azureAttestationTokenEndpoint(svc)
|
||||
res, err := endpoint(context.Background(), tt.req)
|
||||
|
||||
if (err != nil) != tt.expectedErr {
|
||||
t.Errorf("attestationTokenEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
|
||||
}
|
||||
|
||||
if !tt.expectedErr {
|
||||
r, ok := res.(fetchAttestationTokenRes)
|
||||
if !ok {
|
||||
t.Errorf("attestationTokenEndpoint() returned unexpected type %T", res)
|
||||
}
|
||||
if string(r.File) != "mock file" {
|
||||
t.Errorf("expected file content 'mock file', got %s", r.File)
|
||||
}
|
||||
}
|
||||
|
||||
svc.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,74 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/agent/auth"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
type authInterceptor struct {
|
||||
auth auth.Authenticator
|
||||
}
|
||||
|
||||
type wrappedServerStream struct {
|
||||
grpc.ServerStream
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func (s *wrappedServerStream) Context() context.Context {
|
||||
return s.ctx
|
||||
}
|
||||
|
||||
func NewAuthInterceptor(authSvc auth.Authenticator) (grpc.UnaryServerInterceptor, grpc.StreamServerInterceptor) {
|
||||
ai := &authInterceptor{auth: authSvc}
|
||||
return ai.AuthUnaryInterceptor(), ai.AuthStreamInterceptor()
|
||||
}
|
||||
|
||||
func (s *authInterceptor) AuthStreamInterceptor() grpc.StreamServerInterceptor {
|
||||
return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
|
||||
switch info.FullMethod {
|
||||
case agent.AgentService_Algo_FullMethodName:
|
||||
if _, err := s.auth.AuthenticateUser(stream.Context(), auth.AlgorithmProviderRole); err != nil {
|
||||
return status.Errorf(codes.Unauthenticated, "%v", err.Error())
|
||||
}
|
||||
return handler(srv, stream)
|
||||
case agent.AgentService_Data_FullMethodName:
|
||||
ctx, err := s.auth.AuthenticateUser(stream.Context(), auth.DataProviderRole)
|
||||
if err != nil {
|
||||
return status.Errorf(codes.Unauthenticated, "%s", err.Error())
|
||||
}
|
||||
wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
|
||||
return handler(srv, wrapped)
|
||||
case agent.AgentService_Result_FullMethodName:
|
||||
ctx, err := s.auth.AuthenticateUser(stream.Context(), auth.ConsumerRole)
|
||||
if err != nil {
|
||||
return status.Errorf(codes.Unauthenticated, "%v", err.Error())
|
||||
}
|
||||
wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
|
||||
return handler(srv, wrapped)
|
||||
default:
|
||||
return handler(srv, stream)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *authInterceptor) AuthUnaryInterceptor() grpc.UnaryServerInterceptor {
|
||||
return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
|
||||
switch info.FullMethod {
|
||||
case agent.AgentService_Result_FullMethodName:
|
||||
ctx, err := s.auth.AuthenticateUser(ctx, auth.ConsumerRole)
|
||||
if err != nil {
|
||||
return nil, status.Errorf(codes.Unauthenticated, "%v", err.Error())
|
||||
}
|
||||
return handler(ctx, req)
|
||||
default:
|
||||
return handler(ctx, req)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,152 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/agent/auth"
|
||||
"github.com/ultravioletrs/cocos/agent/auth/mocks"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
func TestAuthUnaryInterceptor(t *testing.T) {
|
||||
authmock := new(mocks.Authenticator)
|
||||
tests := []struct {
|
||||
name string
|
||||
authorized bool
|
||||
method string
|
||||
role auth.UserRole
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "authorized result method",
|
||||
authorized: true,
|
||||
method: agent.AgentService_Result_FullMethodName,
|
||||
role: auth.ConsumerRole,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "unauthorized result method",
|
||||
authorized: false,
|
||||
method: agent.AgentService_Result_FullMethodName,
|
||||
role: auth.ConsumerRole,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "other method",
|
||||
authorized: false,
|
||||
method: "OtherMethod",
|
||||
role: auth.ConsumerRole,
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
switch tt.authorized {
|
||||
case true:
|
||||
mockCall := authmock.On("AuthenticateUser", context.Background(), tt.role).Return(context.Background(), nil)
|
||||
mockCall.Once()
|
||||
case false:
|
||||
mockCall := authmock.On("AuthenticateUser", context.Background(), tt.role).Return(context.Background(), auth.ErrMissingMetadata)
|
||||
mockCall.Once()
|
||||
}
|
||||
unaryInt, _ := NewAuthInterceptor(authmock)
|
||||
|
||||
_, err := unaryInt(context.Background(), nil, &grpc.UnaryServerInfo{FullMethod: tt.method}, func(ctx context.Context, req any) (any, error) {
|
||||
return nil, nil
|
||||
})
|
||||
|
||||
if tt.wantErr && err == nil {
|
||||
t.Errorf("expected error, got none")
|
||||
} else if !tt.wantErr && err != nil {
|
||||
t.Errorf("unexpected error: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthStreamInterceptor(t *testing.T) {
|
||||
authmock := new(mocks.Authenticator)
|
||||
tests := []struct {
|
||||
name string
|
||||
authorized bool
|
||||
method string
|
||||
role auth.UserRole
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "authorized algo method",
|
||||
authorized: true,
|
||||
method: agent.AgentService_Algo_FullMethodName,
|
||||
role: auth.AlgorithmProviderRole,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "unauthorized algo method",
|
||||
authorized: false,
|
||||
method: agent.AgentService_Algo_FullMethodName,
|
||||
role: auth.AlgorithmProviderRole,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "authorized data method",
|
||||
authorized: true,
|
||||
method: agent.AgentService_Data_FullMethodName,
|
||||
role: auth.DataProviderRole,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "unauthorized data method",
|
||||
authorized: false,
|
||||
method: agent.AgentService_Data_FullMethodName,
|
||||
role: auth.DataProviderRole,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "other method",
|
||||
authorized: false,
|
||||
method: "OtherMethod",
|
||||
role: auth.DataProviderRole,
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
switch tt.authorized {
|
||||
case true:
|
||||
mockCall := authmock.On("AuthenticateUser", mock.Anything, tt.role).Return(context.Background(), nil)
|
||||
mockCall.Once()
|
||||
case false:
|
||||
mockCall := authmock.On("AuthenticateUser", mock.Anything, tt.role).Return(context.Background(), auth.ErrMissingMetadata)
|
||||
mockCall.Once()
|
||||
}
|
||||
_, streamInt := NewAuthInterceptor(authmock)
|
||||
|
||||
err := streamInt(nil, &mockServerStream{ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs())}, &grpc.StreamServerInfo{FullMethod: tt.method}, func(srv any, stream grpc.ServerStream) error {
|
||||
return nil
|
||||
})
|
||||
|
||||
if tt.wantErr && err == nil {
|
||||
t.Errorf("expected error, got none")
|
||||
} else if !tt.wantErr && err != nil {
|
||||
t.Errorf("unexpected error: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type mockServerStream struct {
|
||||
grpc.ServerStream
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func (m *mockServerStream) Context() context.Context {
|
||||
return m.ctx
|
||||
}
|
||||
+22
-35
@@ -4,36 +4,49 @@ package grpc
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
)
|
||||
|
||||
type algoReq struct {
|
||||
Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
|
||||
Requirements []byte
|
||||
Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
|
||||
Provider string `protobuf:"bytes,2,opt,name=provider,proto3" json:"provider,omitempty"`
|
||||
Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"`
|
||||
}
|
||||
|
||||
func (req algoReq) validate() error {
|
||||
if len(req.Algorithm) == 0 {
|
||||
return errors.New("algorithm binary is required")
|
||||
}
|
||||
if req.Id == "" {
|
||||
return errors.New("malformed entity")
|
||||
}
|
||||
if req.Provider == "" {
|
||||
return errors.New("malformed entity")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type dataReq struct {
|
||||
Dataset []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
|
||||
Filename string
|
||||
Provider string `protobuf:"bytes,2,opt,name=provider,proto3" json:"provider,omitempty"`
|
||||
Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"`
|
||||
}
|
||||
|
||||
func (req dataReq) validate() error {
|
||||
if len(req.Dataset) == 0 {
|
||||
return errors.New("dataset is required")
|
||||
return errors.New("dataset CSV file is required")
|
||||
}
|
||||
if req.Id == "" {
|
||||
return errors.New("malformed entity")
|
||||
}
|
||||
if req.Provider == "" {
|
||||
return errors.New("malformed entity")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type resultReq struct{}
|
||||
type resultReq struct {
|
||||
Consumer string `protobuf:"bytes,1,opt,name=consumer,proto3" json:"consumer,omitempty"`
|
||||
}
|
||||
|
||||
func (req resultReq) validate() error {
|
||||
// No request parameters to validate, so no validation logic needed
|
||||
@@ -41,35 +54,9 @@ func (req resultReq) validate() error {
|
||||
}
|
||||
|
||||
type attestationReq struct {
|
||||
TeeNonce [vtpm.SEVNonce]byte
|
||||
VtpmNonce [vtpm.Nonce]byte
|
||||
AttType attestation.PlatformType
|
||||
}
|
||||
|
||||
type azureAttestationTokenReq struct {
|
||||
tokenNonce [vtpm.Nonce]byte
|
||||
ReportData [64]byte
|
||||
}
|
||||
|
||||
func (req attestationReq) validate() error {
|
||||
return validateAttestationType(req.AttType)
|
||||
}
|
||||
|
||||
func (req azureAttestationTokenReq) validate() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateAttestationType(attType attestation.PlatformType) error {
|
||||
switch attType {
|
||||
case attestation.SNP, attestation.VTPM, attestation.SNPvTPM, attestation.Azure, attestation.TDX:
|
||||
return nil
|
||||
default:
|
||||
return errors.New("invalid attestation type")
|
||||
}
|
||||
}
|
||||
|
||||
type imaMeasurementsReq struct{}
|
||||
|
||||
func (req imaMeasurementsReq) validate() error {
|
||||
// No request parameters to validate, so no validation logic needed
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -7,18 +7,9 @@ type algoRes struct{}
|
||||
type dataRes struct{}
|
||||
|
||||
type resultRes struct {
|
||||
File []byte
|
||||
File []byte `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"`
|
||||
}
|
||||
|
||||
type attestationRes struct {
|
||||
File []byte
|
||||
}
|
||||
|
||||
type imaMeasurementsRes struct {
|
||||
File []byte
|
||||
PCR10 []byte
|
||||
}
|
||||
|
||||
type fetchAttestationTokenRes struct {
|
||||
File []byte
|
||||
}
|
||||
|
||||
+68
-371
@@ -3,437 +3,134 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
|
||||
"github.com/go-kit/kit/endpoint"
|
||||
"github.com/go-kit/kit/transport/grpc"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/metadata"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
const (
|
||||
bufferSize = 1024 * 1024
|
||||
FileSizeKey = "file-size"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrTEENonceLength = errors.New("malformed report data, expect less or equal to 64 bytes")
|
||||
ErrVTPMNonceLength = errors.New("malformed vTPM nonce, expect less or equal to 32 bytes")
|
||||
ErrTokenNonceLength = errors.New("malformed token nonce, expect less or equal to 32 bytes")
|
||||
)
|
||||
|
||||
var _ agent.AgentServiceServer = (*grpcServer)(nil)
|
||||
|
||||
type grpcServer struct {
|
||||
handlers map[string]grpc.Handler
|
||||
algo grpc.Handler
|
||||
data grpc.Handler
|
||||
result grpc.Handler
|
||||
attestation grpc.Handler
|
||||
agent.UnimplementedAgentServiceServer
|
||||
}
|
||||
|
||||
type endpointConfig struct {
|
||||
endpoint func(agent.Service) endpoint.Endpoint
|
||||
decodeRequest grpc.DecodeRequestFunc
|
||||
encodeResponse grpc.EncodeResponseFunc
|
||||
}
|
||||
|
||||
// NewServer returns new AgentServiceServer instance.
|
||||
func NewServer(svc agent.Service) agent.AgentServiceServer {
|
||||
// Define endpoint configurations
|
||||
endpoints := map[string]endpointConfig{
|
||||
"algo": {
|
||||
endpoint: algoEndpoint,
|
||||
decodeRequest: decodeAlgoRequest,
|
||||
encodeResponse: encodeAlgoResponse,
|
||||
},
|
||||
"data": {
|
||||
endpoint: dataEndpoint,
|
||||
decodeRequest: decodeDataRequest,
|
||||
encodeResponse: encodeDataResponse,
|
||||
},
|
||||
"result": {
|
||||
endpoint: resultEndpoint,
|
||||
decodeRequest: decodeResultRequest,
|
||||
encodeResponse: encodeResultResponse,
|
||||
},
|
||||
"attestation": {
|
||||
endpoint: attestationEndpoint,
|
||||
decodeRequest: decodeAttestationRequest,
|
||||
encodeResponse: encodeAttestationResponse,
|
||||
},
|
||||
"imaMeasurements": {
|
||||
endpoint: imaMeasurementsEndpoint,
|
||||
decodeRequest: decodeIMAMeasurementsRequest,
|
||||
encodeResponse: encodeIMAMeasurementsResponse,
|
||||
},
|
||||
"azureAttestationToken": {
|
||||
endpoint: azureAttestationTokenEndpoint,
|
||||
decodeRequest: decodeAttestationTokenRequest,
|
||||
encodeResponse: encodeAttestationTokenResponse,
|
||||
},
|
||||
}
|
||||
|
||||
// Create handlers using the configurations
|
||||
handlers := make(map[string]grpc.Handler)
|
||||
for name, config := range endpoints {
|
||||
handlers[name] = grpc.NewServer(
|
||||
config.endpoint(svc),
|
||||
config.decodeRequest,
|
||||
config.encodeResponse,
|
||||
)
|
||||
}
|
||||
|
||||
return &grpcServer{
|
||||
handlers: handlers,
|
||||
algo: grpc.NewServer(
|
||||
algoEndpoint(svc),
|
||||
decodeAlgoRequest,
|
||||
encodeAlgoResponse,
|
||||
),
|
||||
data: grpc.NewServer(
|
||||
dataEndpoint(svc),
|
||||
decodeDataRequest,
|
||||
encodeDataResponse,
|
||||
),
|
||||
result: grpc.NewServer(
|
||||
resultEndpoint(svc),
|
||||
decodeResultRequest,
|
||||
encodeResultResponse,
|
||||
),
|
||||
attestation: grpc.NewServer(
|
||||
attestationEndpoint(svc),
|
||||
decodeAttestationRequest,
|
||||
encodeAttestationResponse,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
func decodeAlgoRequest(_ context.Context, grpcReq any) (any, error) {
|
||||
func decodeAlgoRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
|
||||
req := grpcReq.(*agent.AlgoRequest)
|
||||
|
||||
return algoReq{
|
||||
Algorithm: req.Algorithm,
|
||||
Requirements: req.Requirements,
|
||||
Algorithm: req.Algorithm,
|
||||
Provider: req.Provider,
|
||||
Id: req.Id,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func encodeAlgoResponse(_ context.Context, response any) (any, error) {
|
||||
func encodeAlgoResponse(_ context.Context, response interface{}) (interface{}, error) {
|
||||
return &agent.AlgoResponse{}, nil
|
||||
}
|
||||
|
||||
func decodeDataRequest(_ context.Context, grpcReq any) (any, error) {
|
||||
func decodeDataRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
|
||||
req := grpcReq.(*agent.DataRequest)
|
||||
|
||||
return dataReq{
|
||||
Dataset: req.Dataset,
|
||||
Filename: req.Filename,
|
||||
Provider: req.Provider,
|
||||
Id: req.Id,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func encodeDataResponse(_ context.Context, response any) (any, error) {
|
||||
func encodeDataResponse(_ context.Context, response interface{}) (interface{}, error) {
|
||||
return &agent.DataResponse{}, nil
|
||||
}
|
||||
|
||||
func decodeResultRequest(_ context.Context, grpcReq any) (any, error) {
|
||||
return resultReq{}, nil
|
||||
func decodeResultRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
|
||||
req := grpcReq.(*agent.ResultRequest)
|
||||
return resultReq{Consumer: req.Consumer}, nil
|
||||
}
|
||||
|
||||
func encodeResultResponse(_ context.Context, response any) (any, error) {
|
||||
func encodeResultResponse(_ context.Context, response interface{}) (interface{}, error) {
|
||||
res := response.(resultRes)
|
||||
return &agent.ResultResponse{
|
||||
File: res.File,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func validateNonce(nonce []byte, maxLen int, target any) error {
|
||||
if len(nonce) > maxLen {
|
||||
switch maxLen {
|
||||
case vtpm.SEVNonce:
|
||||
return ErrTEENonceLength
|
||||
case vtpm.Nonce:
|
||||
return ErrVTPMNonceLength
|
||||
default:
|
||||
return ErrTokenNonceLength
|
||||
}
|
||||
}
|
||||
|
||||
switch t := target.(type) {
|
||||
case *[vtpm.SEVNonce]byte:
|
||||
copy(t[:], nonce)
|
||||
case *[vtpm.Nonce]byte:
|
||||
copy(t[:], nonce)
|
||||
default:
|
||||
return fmt.Errorf("unsupported target type for nonce validation: %T", target)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeAttestationRequest(_ context.Context, grpcReq any) (any, error) {
|
||||
func decodeAttestationRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
|
||||
req := grpcReq.(*agent.AttestationRequest)
|
||||
var reportData [vtpm.SEVNonce]byte
|
||||
var nonce [vtpm.Nonce]byte
|
||||
|
||||
if err := validateNonce(req.TeeNonce, vtpm.SEVNonce, &reportData); err != nil {
|
||||
return nil, err
|
||||
if len(req.ReportData) != agent.ReportDataSize {
|
||||
return nil, errors.New("malformed report data, expect 64 bytes")
|
||||
}
|
||||
|
||||
if err := validateNonce(req.VtpmNonce, vtpm.Nonce, &nonce); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return attestationReq{
|
||||
TeeNonce: reportData,
|
||||
VtpmNonce: nonce,
|
||||
AttType: attestation.PlatformType(req.Type),
|
||||
}, nil
|
||||
return attestationReq{ReportData: [agent.ReportDataSize]byte(req.ReportData)}, nil
|
||||
}
|
||||
|
||||
func encodeAttestationResponse(_ context.Context, response any) (any, error) {
|
||||
func encodeAttestationResponse(_ context.Context, response interface{}) (interface{}, error) {
|
||||
res := response.(attestationRes)
|
||||
return &agent.AttestationResponse{
|
||||
File: res.File,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func decodeAttestationTokenRequest(_ context.Context, grpcReq any) (any, error) {
|
||||
req := grpcReq.(*agent.AttestationTokenRequest)
|
||||
var nonce [vtpm.Nonce]byte
|
||||
|
||||
if err := validateNonce(req.TokenNonce, vtpm.Nonce, &nonce); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return azureAttestationTokenReq{
|
||||
tokenNonce: nonce,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func encodeAttestationTokenResponse(_ context.Context, response any) (any, error) {
|
||||
res := response.(fetchAttestationTokenRes)
|
||||
return &agent.AttestationTokenResponse{
|
||||
File: res.File,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func decodeIMAMeasurementsRequest(_ context.Context, grpcReq any) (any, error) {
|
||||
return imaMeasurementsReq{}, nil
|
||||
}
|
||||
|
||||
func encodeIMAMeasurementsResponse(_ context.Context, response any) (any, error) {
|
||||
res := response.(imaMeasurementsRes)
|
||||
return &agent.IMAMeasurementsResponse{
|
||||
File: res.File,
|
||||
Pcr10: res.PCR10,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *grpcServer) streamingHandler(
|
||||
ctx context.Context,
|
||||
handlerName string,
|
||||
req any,
|
||||
stream any,
|
||||
sendFn func([]byte) error,
|
||||
getFileData func(any) []byte,
|
||||
) error {
|
||||
handler, ok := s.handlers[handlerName]
|
||||
if !ok {
|
||||
return status.Errorf(codes.NotFound, "handler %q not found", handlerName)
|
||||
}
|
||||
|
||||
_, res, err := handler.ServeGRPC(ctx, req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
fileData := getFileData(res)
|
||||
|
||||
// Set file size header
|
||||
if setter, ok := stream.(interface{ SetHeader(metadata.MD) error }); ok {
|
||||
if err := setter.SetHeader(metadata.New(map[string]string{
|
||||
FileSizeKey: fmt.Sprint(len(fileData)),
|
||||
})); err != nil {
|
||||
return status.Error(codes.Internal, err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// Stream the file data
|
||||
return s.streamFileData(bytes.NewBuffer(fileData), sendFn)
|
||||
}
|
||||
|
||||
func (s *grpcServer) streamFileData(buffer *bytes.Buffer, sendFn func([]byte) error) error {
|
||||
buf := make([]byte, bufferSize)
|
||||
for {
|
||||
n, err := buffer.Read(buf)
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return status.Error(codes.Internal, err.Error())
|
||||
}
|
||||
|
||||
if err := sendFn(buf[:n]); err != nil {
|
||||
return status.Error(codes.Internal, err.Error())
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func receiveStreamingData(getData func() ([]byte, string, error)) ([]byte, string, error) {
|
||||
var data []byte
|
||||
var filename string
|
||||
|
||||
for {
|
||||
chunk, fname, err := getData()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, "", status.Error(codes.Internal, err.Error())
|
||||
}
|
||||
data = append(data, chunk...)
|
||||
if fname != "" {
|
||||
filename = fname
|
||||
}
|
||||
}
|
||||
return data, filename, nil
|
||||
}
|
||||
|
||||
// Algo implements agent.AgentServiceServer.
|
||||
func (s *grpcServer) Algo(stream agent.AgentService_AlgoServer) error {
|
||||
algoFile, reqFile, err := s.receiveAlgoData(stream)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, res, err := s.handlers["algo"].ServeGRPC(stream.Context(), &agent.AlgoRequest{
|
||||
Algorithm: algoFile,
|
||||
Requirements: reqFile,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return stream.SendAndClose(res.(*agent.AlgoResponse))
|
||||
}
|
||||
|
||||
func (s *grpcServer) receiveAlgoData(stream agent.AgentService_AlgoServer) ([]byte, []byte, error) {
|
||||
var algoFile, reqFile []byte
|
||||
for {
|
||||
chunk, err := stream.Recv()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, nil, status.Error(codes.Internal, err.Error())
|
||||
}
|
||||
algoFile = append(algoFile, chunk.Algorithm...)
|
||||
reqFile = append(reqFile, chunk.Requirements...)
|
||||
}
|
||||
return algoFile, reqFile, nil
|
||||
}
|
||||
|
||||
// Data implements agent.AgentServiceServer.
|
||||
func (s *grpcServer) Data(stream agent.AgentService_DataServer) error {
|
||||
dataFile, filename, err := receiveStreamingData(func() ([]byte, string, error) {
|
||||
chunk, err := stream.Recv()
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return chunk.Dataset, chunk.Filename, nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, res, err := s.handlers["data"].ServeGRPC(stream.Context(), &agent.DataRequest{
|
||||
Dataset: dataFile,
|
||||
Filename: filename,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return stream.SendAndClose(res.(*agent.DataResponse))
|
||||
}
|
||||
|
||||
func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_ResultServer) error {
|
||||
return s.streamingHandler(
|
||||
stream.Context(),
|
||||
"result",
|
||||
req,
|
||||
stream,
|
||||
func(data []byte) error {
|
||||
return stream.Send(&agent.ResultResponse{File: data})
|
||||
},
|
||||
func(res any) []byte {
|
||||
return res.(*agent.ResultResponse).File
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func (s *grpcServer) Attestation(req *agent.AttestationRequest, stream agent.AgentService_AttestationServer) error {
|
||||
return s.streamingHandler(
|
||||
stream.Context(),
|
||||
"attestation",
|
||||
req,
|
||||
stream,
|
||||
func(data []byte) error {
|
||||
return stream.Send(&agent.AttestationResponse{File: data})
|
||||
},
|
||||
func(res any) []byte {
|
||||
return res.(*agent.AttestationResponse).File
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func (s *grpcServer) IMAMeasurements(req *agent.IMAMeasurementsRequest, stream agent.AgentService_IMAMeasurementsServer) error {
|
||||
_, res, err := s.handlers["imaMeasurements"].ServeGRPC(stream.Context(), req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rr := res.(*agent.IMAMeasurementsResponse)
|
||||
|
||||
if err := stream.SetHeader(metadata.New(map[string]string{
|
||||
FileSizeKey: strconv.Itoa(len(rr.File)),
|
||||
})); err != nil {
|
||||
return status.Error(codes.Internal, err.Error())
|
||||
}
|
||||
|
||||
return s.streamDualBuffers(
|
||||
bytes.NewBuffer(rr.File),
|
||||
bytes.NewBuffer(rr.Pcr10),
|
||||
func(fileData, pcr10Data []byte) error {
|
||||
return stream.Send(&agent.IMAMeasurementsResponse{
|
||||
File: fileData,
|
||||
Pcr10: pcr10Data,
|
||||
})
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func (s *grpcServer) AzureAttestationToken(ctx context.Context, req *agent.AttestationTokenRequest) (*agent.AttestationTokenResponse, error) {
|
||||
_, res, err := s.handlers["azureAttestationToken"].ServeGRPC(ctx, req)
|
||||
func (s *grpcServer) Algo(ctx context.Context, req *agent.AlgoRequest) (*agent.AlgoResponse, error) {
|
||||
_, res, err := s.algo.ServeGRPC(ctx, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ar := res.(*agent.AlgoResponse)
|
||||
return ar, nil
|
||||
}
|
||||
|
||||
rr, ok := res.(*agent.AttestationTokenResponse)
|
||||
if !ok {
|
||||
return nil, status.Error(codes.Internal, "failed to cast response to AttestationTokenResponse")
|
||||
func (s *grpcServer) Data(ctx context.Context, req *agent.DataRequest) (*agent.DataResponse, error) {
|
||||
_, res, err := s.data.ServeGRPC(ctx, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dr := res.(*agent.DataResponse)
|
||||
return dr, nil
|
||||
}
|
||||
|
||||
func (s *grpcServer) Result(ctx context.Context, req *agent.ResultRequest) (*agent.ResultResponse, error) {
|
||||
_, res, err := s.result.ServeGRPC(ctx, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rr := res.(*agent.ResultResponse)
|
||||
return rr, nil
|
||||
}
|
||||
|
||||
func (s *grpcServer) streamDualBuffers(
|
||||
buf1, buf2 *bytes.Buffer,
|
||||
sendFn func([]byte, []byte) error,
|
||||
) error {
|
||||
buff1 := make([]byte, bufferSize)
|
||||
buff2 := make([]byte, bufferSize)
|
||||
|
||||
for {
|
||||
n1, err1 := buf1.Read(buff1)
|
||||
if err1 != nil && err1 != io.EOF {
|
||||
return status.Error(codes.Internal, err1.Error())
|
||||
}
|
||||
|
||||
n2, err2 := buf2.Read(buff2)
|
||||
if err2 != nil && err2 != io.EOF {
|
||||
return status.Error(codes.Internal, err2.Error())
|
||||
}
|
||||
|
||||
if n1 == 0 && err1 == io.EOF && n2 == 0 && err2 == io.EOF {
|
||||
break
|
||||
}
|
||||
|
||||
if err := sendFn(buff1[:n1], buff2[:n2]); err != nil {
|
||||
return status.Error(codes.Internal, err.Error())
|
||||
}
|
||||
func (s *grpcServer) Attestation(ctx context.Context, req *agent.AttestationRequest) (*agent.AttestationResponse, error) {
|
||||
_, res, err := s.attestation.ServeGRPC(ctx, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil
|
||||
rr := res.(*agent.AttestationResponse)
|
||||
return rr, nil
|
||||
}
|
||||
|
||||
@@ -1,499 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/agent/mocks"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
type MockAgentService_AlgoServer struct {
|
||||
grpc.ServerStream
|
||||
mock.Mock
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func (m *MockAgentService_AlgoServer) Context() context.Context {
|
||||
return m.ctx
|
||||
}
|
||||
|
||||
func (m *MockAgentService_AlgoServer) Recv() (*agent.AlgoRequest, error) {
|
||||
args := m.Called()
|
||||
return args.Get(0).(*agent.AlgoRequest), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockAgentService_AlgoServer) SendAndClose(resp *agent.AlgoResponse) error {
|
||||
args := m.Called(resp)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
type MockAgentService_DataServer struct {
|
||||
grpc.ServerStream
|
||||
mock.Mock
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func (m *MockAgentService_DataServer) Context() context.Context {
|
||||
return m.ctx
|
||||
}
|
||||
|
||||
func (m *MockAgentService_DataServer) Recv() (*agent.DataRequest, error) {
|
||||
args := m.Called()
|
||||
return args.Get(0).(*agent.DataRequest), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockAgentService_DataServer) SendAndClose(resp *agent.DataResponse) error {
|
||||
args := m.Called(resp)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
type MockAgentService_ResultServer struct {
|
||||
grpc.ServerStream
|
||||
mock.Mock
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func (m *MockAgentService_ResultServer) Context() context.Context {
|
||||
return m.ctx
|
||||
}
|
||||
|
||||
func (m *MockAgentService_ResultServer) SetHeader(md metadata.MD) error {
|
||||
args := m.Called(md)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockAgentService_ResultServer) Send(resp *agent.ResultResponse) error {
|
||||
args := m.Called(resp)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
type MockAgentService_AttestationServer struct {
|
||||
grpc.ServerStream
|
||||
mock.Mock
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func (m *MockAgentService_AttestationServer) Context() context.Context {
|
||||
return m.ctx
|
||||
}
|
||||
|
||||
func (m *MockAgentService_AttestationServer) Send(resp *agent.AttestationResponse) error {
|
||||
args := m.Called(resp)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockAgentService_AttestationServer) SetHeader(md metadata.MD) error {
|
||||
args := m.Called(md)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
type MockAgentService_IMAMeasurementsServer struct {
|
||||
grpc.ServerStream
|
||||
mock.Mock
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func (m *MockAgentService_IMAMeasurementsServer) Context() context.Context {
|
||||
return m.ctx
|
||||
}
|
||||
|
||||
func (m *MockAgentService_IMAMeasurementsServer) Send(resp *agent.IMAMeasurementsResponse) error {
|
||||
args := m.Called(resp)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockAgentService_IMAMeasurementsServer) SetHeader(md metadata.MD) error {
|
||||
args := m.Called(md)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func TestNewServer(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
grpcServer, ok := server.(*grpcServer)
|
||||
assert.True(t, ok)
|
||||
assert.NotNil(t, grpcServer.handlers)
|
||||
assert.Len(t, grpcServer.handlers, 6) // Should have 6 handlers
|
||||
|
||||
// Check that all expected handlers are present
|
||||
expectedHandlers := []string{"algo", "data", "result", "attestation", "imaMeasurements", "azureAttestationToken"}
|
||||
for _, handler := range expectedHandlers {
|
||||
assert.Contains(t, grpcServer.handlers, handler)
|
||||
assert.NotNil(t, grpcServer.handlers[handler])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAlgo(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
mockStream := &MockAgentService_AlgoServer{ctx: context.Background()}
|
||||
mockStream.On("Recv").Return(&agent.AlgoRequest{Algorithm: []byte("algo"), Requirements: []byte("req")}, nil).Once()
|
||||
mockStream.On("Recv").Return(&agent.AlgoRequest{}, io.EOF).Once()
|
||||
mockStream.On("SendAndClose", &agent.AlgoResponse{}).Return(nil).Once()
|
||||
|
||||
mockService.On("Algo", context.Background(), agent.Algorithm{Algorithm: []byte("algo"), Requirements: []byte("req")}).Return(nil)
|
||||
|
||||
err := server.Algo(mockStream)
|
||||
assert.NoError(t, err)
|
||||
|
||||
mockStream.AssertExpectations(t)
|
||||
mockService.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestAlgoWithMultipleChunks(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
mockStream := &MockAgentService_AlgoServer{ctx: context.Background()}
|
||||
mockStream.On("Recv").Return(&agent.AlgoRequest{Algorithm: []byte("algo"), Requirements: []byte("req")}, nil).Once()
|
||||
mockStream.On("Recv").Return(&agent.AlgoRequest{Algorithm: []byte("2"), Requirements: []byte("2")}, nil).Once()
|
||||
mockStream.On("Recv").Return(&agent.AlgoRequest{}, io.EOF).Once()
|
||||
mockStream.On("SendAndClose", &agent.AlgoResponse{}).Return(nil).Once()
|
||||
|
||||
mockService.On("Algo", context.Background(), agent.Algorithm{Algorithm: []byte("algo2"), Requirements: []byte("req2")}).Return(nil)
|
||||
|
||||
err := server.Algo(mockStream)
|
||||
assert.NoError(t, err)
|
||||
|
||||
mockStream.AssertExpectations(t)
|
||||
mockService.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestData(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
mockStream := &MockAgentService_DataServer{ctx: context.Background()}
|
||||
mockStream.On("Recv").Return(&agent.DataRequest{Dataset: []byte("data"), Filename: "test.txt"}, nil).Once()
|
||||
mockStream.On("Recv").Return(&agent.DataRequest{}, io.EOF).Once()
|
||||
mockStream.On("SendAndClose", &agent.DataResponse{}).Return(nil).Once()
|
||||
|
||||
mockService.On("Data", context.Background(), agent.Dataset{Dataset: []byte("data"), Filename: "test.txt"}).Return(nil)
|
||||
|
||||
err := server.Data(mockStream)
|
||||
assert.NoError(t, err)
|
||||
|
||||
mockStream.AssertExpectations(t)
|
||||
mockService.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestResult(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
resultData := []byte("result data")
|
||||
mockStream := &MockAgentService_ResultServer{ctx: context.Background()}
|
||||
|
||||
// Mock the SetHeader call
|
||||
mockStream.On("SetHeader", mock.AnythingOfType("metadata.MD")).Return(nil).Once()
|
||||
|
||||
// Mock the Send call - it should be called with the result data
|
||||
mockStream.On("Send", mock.MatchedBy(func(resp *agent.ResultResponse) bool {
|
||||
return len(resp.File) > 0
|
||||
})).Return(nil).Once()
|
||||
|
||||
mockService.On("Result", mock.Anything).Return(resultData, nil)
|
||||
|
||||
err := server.Result(&agent.ResultRequest{}, mockStream)
|
||||
assert.NoError(t, err)
|
||||
|
||||
mockStream.AssertExpectations(t)
|
||||
mockService.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestAttestation(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
attestationData := []byte("attestation data")
|
||||
mockStream := &MockAgentService_AttestationServer{ctx: context.Background()}
|
||||
|
||||
// Mock the SetHeader call
|
||||
mockStream.On("SetHeader", mock.AnythingOfType("metadata.MD")).Return(nil).Once()
|
||||
|
||||
// Mock the Send call
|
||||
mockStream.On("Send", mock.MatchedBy(func(resp *agent.AttestationResponse) bool {
|
||||
return len(resp.File) > 0
|
||||
})).Return(nil).Once()
|
||||
|
||||
reportData := [vtpm.SEVNonce]byte{}
|
||||
vtpmNonce := [vtpm.Nonce]byte{}
|
||||
attestationType := attestation.SNP
|
||||
mockService.On("Attestation", mock.Anything, reportData, vtpmNonce, attestationType).Return(attestationData, nil)
|
||||
|
||||
err := server.Attestation(&agent.AttestationRequest{TeeNonce: reportData[:], Type: int32(attestationType)}, mockStream)
|
||||
assert.NoError(t, err)
|
||||
|
||||
mockService.AssertExpectations(t)
|
||||
mockStream.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestIMAMeasurements(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
imaData := []byte("ima data")
|
||||
pcr10Data := []byte("pcr10 data")
|
||||
|
||||
mockStream := &MockAgentService_IMAMeasurementsServer{ctx: context.Background()}
|
||||
|
||||
// Mock the SetHeader call
|
||||
mockStream.On("SetHeader", mock.AnythingOfType("metadata.MD")).Return(nil).Once()
|
||||
|
||||
// Mock the Send call
|
||||
mockStream.On("Send", mock.MatchedBy(func(resp *agent.IMAMeasurementsResponse) bool {
|
||||
return len(resp.File) > 0 || len(resp.Pcr10) > 0
|
||||
})).Return(nil).Once()
|
||||
|
||||
mockService.On("IMAMeasurements", mock.Anything).Return(imaData, pcr10Data, nil)
|
||||
|
||||
err := server.IMAMeasurements(&agent.IMAMeasurementsRequest{}, mockStream)
|
||||
assert.NoError(t, err)
|
||||
|
||||
mockService.AssertExpectations(t)
|
||||
mockStream.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestAttestationToken(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
attestationData := []byte("attestation token data")
|
||||
vtpmNonce := [vtpm.Nonce]byte{}
|
||||
attestationType := attestation.SNP
|
||||
|
||||
mockService.On("AzureAttestationToken", mock.Anything, vtpmNonce).Return(attestationData, nil)
|
||||
|
||||
resp, err := server.AzureAttestationToken(context.Background(), &agent.AttestationTokenRequest{
|
||||
TokenNonce: vtpmNonce[:],
|
||||
Type: int32(attestationType),
|
||||
})
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, attestationData, resp.File)
|
||||
|
||||
mockService.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestValidateNonce(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
nonce []byte
|
||||
maxLen int
|
||||
shouldError bool
|
||||
expectedErr error
|
||||
}{
|
||||
{
|
||||
name: "valid TEE nonce",
|
||||
nonce: make([]byte, vtpm.SEVNonce),
|
||||
maxLen: vtpm.SEVNonce,
|
||||
shouldError: false,
|
||||
},
|
||||
{
|
||||
name: "valid vTPM nonce",
|
||||
nonce: make([]byte, vtpm.Nonce),
|
||||
maxLen: vtpm.Nonce,
|
||||
shouldError: false,
|
||||
},
|
||||
{
|
||||
name: "TEE nonce too long",
|
||||
nonce: make([]byte, vtpm.SEVNonce+1),
|
||||
maxLen: vtpm.SEVNonce,
|
||||
shouldError: true,
|
||||
expectedErr: ErrTEENonceLength,
|
||||
},
|
||||
{
|
||||
name: "vTPM nonce too long",
|
||||
nonce: make([]byte, vtpm.Nonce+1),
|
||||
maxLen: vtpm.Nonce,
|
||||
shouldError: true,
|
||||
expectedErr: ErrVTPMNonceLength,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.maxLen == vtpm.SEVNonce {
|
||||
var target [vtpm.SEVNonce]byte
|
||||
err := validateNonce(tt.nonce, tt.maxLen, &target)
|
||||
if tt.shouldError {
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, tt.expectedErr, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
} else {
|
||||
var target [vtpm.Nonce]byte
|
||||
err := validateNonce(tt.nonce, tt.maxLen, &target)
|
||||
if tt.shouldError {
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, tt.expectedErr, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeAlgoRequest(t *testing.T) {
|
||||
req := &agent.AlgoRequest{Algorithm: []byte("algo"), Requirements: []byte("req")}
|
||||
decoded, err := decodeAlgoRequest(context.Background(), req)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, algoReq{Algorithm: []byte("algo"), Requirements: []byte("req")}, decoded)
|
||||
}
|
||||
|
||||
func TestEncodeAlgoResponse(t *testing.T) {
|
||||
encoded, err := encodeAlgoResponse(context.Background(), algoRes{})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, &agent.AlgoResponse{}, encoded)
|
||||
}
|
||||
|
||||
func TestDecodeDataRequest(t *testing.T) {
|
||||
req := &agent.DataRequest{Dataset: []byte("data"), Filename: "test.txt"}
|
||||
decoded, err := decodeDataRequest(context.Background(), req)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, dataReq{Dataset: []byte("data"), Filename: "test.txt"}, decoded)
|
||||
}
|
||||
|
||||
func TestEncodeDataResponse(t *testing.T) {
|
||||
encoded, err := encodeDataResponse(context.Background(), dataRes{})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, &agent.DataResponse{}, encoded)
|
||||
}
|
||||
|
||||
func TestDecodeResultRequest(t *testing.T) {
|
||||
decoded, err := decodeResultRequest(context.Background(), &agent.ResultRequest{})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, resultReq{}, decoded)
|
||||
}
|
||||
|
||||
func TestEncodeResultResponse(t *testing.T) {
|
||||
encoded, err := encodeResultResponse(context.Background(), resultRes{File: []byte("result")})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, &agent.ResultResponse{File: []byte("result")}, encoded)
|
||||
}
|
||||
|
||||
func TestDecodeAttestationRequest(t *testing.T) {
|
||||
teeNonce := make([]byte, vtpm.SEVNonce)
|
||||
vtpmNonce := make([]byte, vtpm.Nonce)
|
||||
|
||||
req := &agent.AttestationRequest{
|
||||
TeeNonce: teeNonce,
|
||||
VtpmNonce: vtpmNonce,
|
||||
Type: int32(attestation.SNP),
|
||||
}
|
||||
|
||||
decoded, err := decodeAttestationRequest(context.Background(), req)
|
||||
assert.NoError(t, err)
|
||||
|
||||
decodedReq := decoded.(attestationReq)
|
||||
assert.Equal(t, attestation.SNP, decodedReq.AttType)
|
||||
}
|
||||
|
||||
func TestDecodeAttestationRequestWithInvalidNonce(t *testing.T) {
|
||||
// Test with TEE nonce too long
|
||||
teeNonce := make([]byte, vtpm.SEVNonce+1)
|
||||
req := &agent.AttestationRequest{TeeNonce: teeNonce}
|
||||
|
||||
_, err := decodeAttestationRequest(context.Background(), req)
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, ErrTEENonceLength, err)
|
||||
|
||||
// Test with vTPM nonce too long
|
||||
vtpmNonce := make([]byte, vtpm.Nonce+1)
|
||||
req = &agent.AttestationRequest{VtpmNonce: vtpmNonce}
|
||||
|
||||
_, err = decodeAttestationRequest(context.Background(), req)
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, ErrVTPMNonceLength, err)
|
||||
}
|
||||
|
||||
func TestEncodeAttestationResponse(t *testing.T) {
|
||||
encoded, err := encodeAttestationResponse(context.Background(), attestationRes{File: []byte("attestation")})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, &agent.AttestationResponse{File: []byte("attestation")}, encoded)
|
||||
}
|
||||
|
||||
func TestDecodeAttestationTokenRequest(t *testing.T) {
|
||||
tokenNonce := make([]byte, vtpm.Nonce)
|
||||
req := &agent.AttestationTokenRequest{
|
||||
TokenNonce: tokenNonce,
|
||||
Type: int32(attestation.SNP),
|
||||
}
|
||||
|
||||
_, err := decodeAttestationTokenRequest(context.Background(), req)
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestDecodeAttestationTokenRequestWithInvalidNonce(t *testing.T) {
|
||||
// Test with token nonce too long
|
||||
tokenNonce := make([]byte, vtpm.Nonce+1)
|
||||
req := &agent.AttestationTokenRequest{TokenNonce: tokenNonce}
|
||||
|
||||
_, err := decodeAttestationTokenRequest(context.Background(), req)
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, ErrVTPMNonceLength, err)
|
||||
}
|
||||
|
||||
func TestEncodeAttestationTokenResponse(t *testing.T) {
|
||||
encoded, err := encodeAttestationTokenResponse(context.Background(), fetchAttestationTokenRes{File: []byte("attestation")})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, &agent.AttestationTokenResponse{File: []byte("attestation")}, encoded)
|
||||
}
|
||||
|
||||
func TestDecodeIMAMeasurementsRequest(t *testing.T) {
|
||||
decoded, err := decodeIMAMeasurementsRequest(context.Background(), &agent.IMAMeasurementsRequest{})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, imaMeasurementsReq{}, decoded)
|
||||
}
|
||||
|
||||
func TestEncodeIMAMeasurementsResponse(t *testing.T) {
|
||||
encoded, err := encodeIMAMeasurementsResponse(context.Background(), imaMeasurementsRes{
|
||||
File: []byte("ima"),
|
||||
PCR10: []byte("pcr10"),
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, &agent.IMAMeasurementsResponse{
|
||||
File: []byte("ima"),
|
||||
Pcr10: []byte("pcr10"),
|
||||
}, encoded)
|
||||
}
|
||||
|
||||
func TestAlgoWithStreamError(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
mockStream := &MockAgentService_AlgoServer{ctx: context.Background()}
|
||||
mockStream.On("Recv").Return(&agent.AlgoRequest{}, assert.AnError).Once()
|
||||
|
||||
err := server.Algo(mockStream)
|
||||
assert.Error(t, err)
|
||||
|
||||
mockStream.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestDataWithStreamError(t *testing.T) {
|
||||
mockService := new(mocks.Service)
|
||||
server := NewServer(mockService)
|
||||
|
||||
mockStream := &MockAgentService_DataServer{ctx: context.Background()}
|
||||
mockStream.On("Recv").Return(&agent.DataRequest{}, assert.AnError).Once()
|
||||
|
||||
err := server.Data(mockStream)
|
||||
assert.Error(t, err)
|
||||
|
||||
mockStream.AssertExpectations(t)
|
||||
}
|
||||
+5
-69
@@ -2,6 +2,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//go:build !test
|
||||
// +build !test
|
||||
|
||||
package api
|
||||
|
||||
@@ -12,8 +13,6 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
)
|
||||
|
||||
var _ agent.Service = (*loggingMiddleware)(nil)
|
||||
@@ -28,43 +27,6 @@ func LoggingMiddleware(svc agent.Service, logger *slog.Logger) agent.Service {
|
||||
return &loggingMiddleware{logger, svc}
|
||||
}
|
||||
|
||||
// State implements agent.Service.
|
||||
func (lm *loggingMiddleware) State() (state string) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method State took %s to complete with state %s", time.Since(begin), state)
|
||||
lm.logger.Debug(message)
|
||||
}(time.Now())
|
||||
return lm.svc.State()
|
||||
}
|
||||
|
||||
// InitComputation implements agent.Service.
|
||||
func (lm *loggingMiddleware) InitComputation(ctx context.Context, cmp agent.Computation) (err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method InitComputation for computation id %s took %s to complete", cmp.ID, time.Since(begin))
|
||||
if err != nil {
|
||||
lm.logger.WithGroup(cmp.ID).Warn(fmt.Sprintf("%s with error: %s", message, err))
|
||||
return
|
||||
}
|
||||
lm.logger.WithGroup(cmp.ID).Info(fmt.Sprintf("%s without errors", message))
|
||||
}(time.Now())
|
||||
|
||||
return lm.svc.InitComputation(ctx, cmp)
|
||||
}
|
||||
|
||||
// StopComputation implements agent.Service.
|
||||
func (lm *loggingMiddleware) StopComputation(ctx context.Context) (err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method StopComputation took %s to complete", time.Since(begin))
|
||||
if err != nil {
|
||||
lm.logger.Warn(fmt.Sprintf("%s with error: %s", message, err))
|
||||
return
|
||||
}
|
||||
lm.logger.Info(fmt.Sprintf("%s without errors", message))
|
||||
}(time.Now())
|
||||
|
||||
return lm.svc.StopComputation(ctx)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) Algo(ctx context.Context, algorithm agent.Algorithm) (err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method Algo took %s to complete", time.Since(begin))
|
||||
@@ -91,7 +53,7 @@ func (lm *loggingMiddleware) Data(ctx context.Context, dataset agent.Dataset) (e
|
||||
return lm.svc.Data(ctx, dataset)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) Result(ctx context.Context) (response []byte, err error) {
|
||||
func (lm *loggingMiddleware) Result(ctx context.Context, consumer string) (response []byte, err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method Result took %s to complete", time.Since(begin))
|
||||
if err != nil {
|
||||
@@ -101,10 +63,10 @@ func (lm *loggingMiddleware) Result(ctx context.Context) (response []byte, err e
|
||||
lm.logger.Info(fmt.Sprintf("%s without errors", message))
|
||||
}(time.Now())
|
||||
|
||||
return lm.svc.Result(ctx)
|
||||
return lm.svc.Result(ctx, consumer)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [vtpm.SEVNonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) (response []byte, err error) {
|
||||
func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [agent.ReportDataSize]byte) (response []byte, err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method Attestation took %s to complete", time.Since(begin))
|
||||
if err != nil {
|
||||
@@ -114,31 +76,5 @@ func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [vtpm.S
|
||||
lm.logger.Info(fmt.Sprintf("%s without errors", message))
|
||||
}(time.Now())
|
||||
|
||||
return lm.svc.Attestation(ctx, reportData, nonce, attType)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) IMAMeasurements(ctx context.Context) (file []byte, pcr10 []byte, err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method IMAMeasurements took %s to complete", time.Since(begin))
|
||||
if err != nil {
|
||||
lm.logger.Warn(fmt.Sprintf("%s with error: %s", message, err))
|
||||
return
|
||||
}
|
||||
lm.logger.Info(fmt.Sprintf("%s without errors", message))
|
||||
}(time.Now())
|
||||
|
||||
return lm.svc.IMAMeasurements(ctx)
|
||||
}
|
||||
|
||||
func (lm *loggingMiddleware) AzureAttestationToken(ctx context.Context, nonce [vtpm.Nonce]byte) (response []byte, err error) {
|
||||
defer func(begin time.Time) {
|
||||
message := fmt.Sprintf("Method AzureAttestationToken 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.AzureAttestationToken(ctx, nonce)
|
||||
return lm.svc.Attestation(ctx, reportData)
|
||||
}
|
||||
|
||||
+5
-54
@@ -2,6 +2,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//go:build !test
|
||||
// +build !test
|
||||
|
||||
package api
|
||||
|
||||
@@ -11,8 +12,6 @@ import (
|
||||
|
||||
"github.com/go-kit/kit/metrics"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
)
|
||||
|
||||
var _ agent.Service = (*metricsMiddleware)(nil)
|
||||
@@ -33,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)
|
||||
@@ -81,38 +50,20 @@ func (ms *metricsMiddleware) Data(ctx context.Context, dataset agent.Dataset) er
|
||||
return ms.svc.Data(ctx, dataset)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) Result(ctx context.Context) ([]byte, error) {
|
||||
func (ms *metricsMiddleware) Result(ctx context.Context, consumer string) ([]byte, error) {
|
||||
defer func(begin time.Time) {
|
||||
ms.counter.With("method", "result").Add(1)
|
||||
ms.latency.With("method", "result").Observe(time.Since(begin).Seconds())
|
||||
}(time.Now())
|
||||
|
||||
return ms.svc.Result(ctx)
|
||||
return ms.svc.Result(ctx, consumer)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) Attestation(ctx context.Context, reportData [vtpm.SEVNonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error) {
|
||||
func (ms *metricsMiddleware) Attestation(ctx context.Context, reportData [agent.ReportDataSize]byte) ([]byte, error) {
|
||||
defer func(begin time.Time) {
|
||||
ms.counter.With("method", "attestation").Add(1)
|
||||
ms.latency.With("method", "attestation").Observe(time.Since(begin).Seconds())
|
||||
}(time.Now())
|
||||
|
||||
return ms.svc.Attestation(ctx, reportData, nonce, attType)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) AzureAttestationToken(ctx context.Context, nonce [vtpm.Nonce]byte) ([]byte, error) {
|
||||
defer func(begin time.Time) {
|
||||
ms.counter.With("method", "attestation_token").Add(1)
|
||||
ms.latency.With("method", "attestation_token").Observe(time.Since(begin).Seconds())
|
||||
}(time.Now())
|
||||
|
||||
return ms.svc.AzureAttestationToken(ctx, nonce)
|
||||
}
|
||||
|
||||
func (ms *metricsMiddleware) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
|
||||
defer func(begin time.Time) {
|
||||
ms.counter.With("method", "imameasurements").Add(1)
|
||||
ms.latency.With("method", "imameasurements").Observe(time.Since(begin).Seconds())
|
||||
}(time.Now())
|
||||
|
||||
return ms.svc.IMAMeasurements(ctx)
|
||||
return ms.svc.Attestation(ctx, reportData)
|
||||
}
|
||||
|
||||
@@ -1,173 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto"
|
||||
"crypto/ecdsa"
|
||||
"crypto/ed25519"
|
||||
"crypto/rsa"
|
||||
"crypto/sha256"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
|
||||
"github.com/absmach/magistrala/pkg/errors"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/metadata"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
type UserRole string
|
||||
|
||||
const (
|
||||
UserMetadataKey = "user-id"
|
||||
SignatureMetadataKey = "signature"
|
||||
ConsumerRole UserRole = "consumer"
|
||||
DataProviderRole UserRole = "data-provider"
|
||||
AlgorithmProviderRole UserRole = "algorithm-provider"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrMissingMetadata = errors.New("missing metadata")
|
||||
ErrInvalidMetadata = errors.New("invalid metadata")
|
||||
ErrSignatureVerificationFailed = errors.New("signature verification failed")
|
||||
)
|
||||
|
||||
type Authenticator interface {
|
||||
AuthenticateUser(ctx context.Context, role UserRole) (context.Context, error)
|
||||
}
|
||||
|
||||
type service struct {
|
||||
resultConsumers []any
|
||||
datasetProviders []any
|
||||
algorithmProvider any
|
||||
}
|
||||
|
||||
func New(manifest agent.Computation) (Authenticator, error) {
|
||||
s := &service{}
|
||||
for _, rc := range manifest.ResultConsumers {
|
||||
pubKey, err := x509.ParsePKIXPublicKey(rc.UserKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pKey, err := decodePublicKey(pubKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
s.resultConsumers = append(s.resultConsumers, pKey)
|
||||
}
|
||||
|
||||
for _, dp := range manifest.Datasets {
|
||||
pubKey, err := x509.ParsePKIXPublicKey(dp.UserKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pKey, err := decodePublicKey(pubKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
s.datasetProviders = append(s.datasetProviders, pKey)
|
||||
}
|
||||
|
||||
pubKey, err := x509.ParsePKIXPublicKey(manifest.Algorithm.UserKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pKey, err := decodePublicKey(pubKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
s.algorithmProvider = pKey
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func extractSignature(md metadata.MD) (string, error) {
|
||||
signature := md.Get(SignatureMetadataKey)
|
||||
if len(signature) != 1 {
|
||||
return "", status.Errorf(codes.Unauthenticated, "invalid metadata")
|
||||
}
|
||||
|
||||
return signature[0], nil
|
||||
}
|
||||
|
||||
func verifySignature(role UserRole, signature string, publicKey any) error {
|
||||
hash := sha256.Sum256([]byte(role))
|
||||
sigByte, err := base64.StdEncoding.DecodeString(signature)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var ok bool
|
||||
|
||||
switch publicKey := publicKey.(type) {
|
||||
case *rsa.PublicKey:
|
||||
if err = rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, hash[:], sigByte); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
case *ecdsa.PublicKey:
|
||||
ok = ecdsa.VerifyASN1(publicKey, hash[:], sigByte)
|
||||
case ed25519.PublicKey:
|
||||
ok = ed25519.Verify(publicKey, []byte(role), sigByte)
|
||||
}
|
||||
|
||||
if !ok {
|
||||
return ErrSignatureVerificationFailed
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *service) AuthenticateUser(ctx context.Context, role UserRole) (context.Context, error) {
|
||||
md, ok := metadata.FromIncomingContext(ctx)
|
||||
if !ok {
|
||||
return nil, ErrMissingMetadata
|
||||
}
|
||||
signature, err := extractSignature(md)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, ErrInvalidMetadata)
|
||||
}
|
||||
|
||||
switch role {
|
||||
case ConsumerRole:
|
||||
for i, rc := range s.resultConsumers {
|
||||
if err := verifySignature(role, signature, rc); err == nil {
|
||||
return agent.IndexToContext(ctx, i), nil
|
||||
}
|
||||
}
|
||||
case DataProviderRole:
|
||||
for _, dp := range s.datasetProviders {
|
||||
if err := verifySignature(role, signature, dp); err == nil {
|
||||
return ctx, nil
|
||||
}
|
||||
}
|
||||
case AlgorithmProviderRole:
|
||||
if err := verifySignature(role, signature, s.algorithmProvider); err == nil {
|
||||
return ctx, nil
|
||||
}
|
||||
}
|
||||
|
||||
return ctx, ErrSignatureVerificationFailed
|
||||
}
|
||||
|
||||
func decodePublicKey(key any) (pubKey any, err error) {
|
||||
switch key := key.(type) {
|
||||
case *rsa.PublicKey:
|
||||
return key, nil
|
||||
case *ecdsa.PublicKey:
|
||||
return key, nil
|
||||
case ed25519.PublicKey:
|
||||
return key, nil
|
||||
default:
|
||||
return nil, errors.New("unsupported public key type")
|
||||
}
|
||||
}
|
||||
@@ -1,158 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto"
|
||||
"crypto/ecdsa"
|
||||
"crypto/ed25519"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/rsa"
|
||||
"crypto/sha256"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
"testing"
|
||||
|
||||
"github.com/absmach/magistrala/pkg/errors"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
func TestAuthenticateUser(t *testing.T) {
|
||||
resultConsumerKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
require.NoError(t, err)
|
||||
|
||||
dataProviderKey, err := rsa.GenerateKey(rand.Reader, 2048)
|
||||
require.NoError(t, err)
|
||||
|
||||
pubEd25519Key, algorithmProviderKey, err := ed25519.GenerateKey(rand.Reader)
|
||||
require.NoError(t, err)
|
||||
|
||||
resultConsumerPubKey, err := x509.MarshalPKIXPublicKey(&resultConsumerKey.PublicKey)
|
||||
require.NoError(t, err)
|
||||
|
||||
dataProviderPubKey, err := x509.MarshalPKIXPublicKey(&dataProviderKey.PublicKey)
|
||||
require.NoError(t, err)
|
||||
|
||||
algorithmProviderPubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
|
||||
require.NoError(t, err)
|
||||
|
||||
manifest := agent.Computation{
|
||||
ResultConsumers: []agent.ResultConsumer{{UserKey: resultConsumerPubKey}},
|
||||
Datasets: []agent.Dataset{{UserKey: dataProviderPubKey}},
|
||||
Algorithm: &agent.Algorithm{UserKey: algorithmProviderPubKey},
|
||||
}
|
||||
|
||||
auth, err := New(manifest)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create authenticator: %v", err)
|
||||
}
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
role UserRole
|
||||
key any
|
||||
expectedErr error
|
||||
}{
|
||||
{
|
||||
name: "valid result consumer",
|
||||
role: ConsumerRole,
|
||||
key: resultConsumerKey,
|
||||
expectedErr: nil,
|
||||
},
|
||||
{
|
||||
name: "valid data provider",
|
||||
role: DataProviderRole,
|
||||
key: dataProviderKey,
|
||||
expectedErr: nil,
|
||||
},
|
||||
{
|
||||
name: "valid algorithm provider",
|
||||
role: AlgorithmProviderRole,
|
||||
key: algorithmProviderKey,
|
||||
expectedErr: nil,
|
||||
},
|
||||
{
|
||||
name: "invalid role",
|
||||
role: "invalid-role",
|
||||
key: resultConsumerKey,
|
||||
expectedErr: ErrSignatureVerificationFailed,
|
||||
},
|
||||
{
|
||||
name: "invalid key",
|
||||
role: ConsumerRole,
|
||||
key: dataProviderKey,
|
||||
expectedErr: ErrSignatureVerificationFailed,
|
||||
},
|
||||
{
|
||||
name: "missing signature",
|
||||
role: ConsumerRole,
|
||||
key: resultConsumerKey,
|
||||
expectedErr: ErrInvalidMetadata,
|
||||
},
|
||||
{
|
||||
name: "missing metadata",
|
||||
role: ConsumerRole,
|
||||
key: resultConsumerKey,
|
||||
expectedErr: ErrMissingMetadata,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
signature, err := signRole(tc.role, tc.key)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to sign role: %v", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
switch tc.name {
|
||||
case "missing signature":
|
||||
ctx = metadata.NewIncomingContext(ctx, metadata.Pairs())
|
||||
case "missing metadata":
|
||||
default:
|
||||
ctx = metadata.NewIncomingContext(ctx, metadata.Pairs(SignatureMetadataKey, signature))
|
||||
}
|
||||
|
||||
ctx, err = auth.AuthenticateUser(ctx, tc.role)
|
||||
assert.True(t, errors.Contains(err, tc.expectedErr), "expected error %v, got %v", tc.expectedErr, err)
|
||||
|
||||
if err == nil {
|
||||
switch id, ok := agent.IndexFromContext(ctx); {
|
||||
case tc.role == ConsumerRole:
|
||||
assert.True(t, ok, "expected index in context")
|
||||
assert.Equal(t, 0, id, "expected index 0 in context")
|
||||
default:
|
||||
assert.False(t, ok, "expected no index in context")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func signRole(role UserRole, key crypto.PrivateKey) (string, error) {
|
||||
var signature []byte
|
||||
var err error
|
||||
|
||||
switch k := key.(type) {
|
||||
case ed25519.PrivateKey:
|
||||
signature, err = k.Sign(rand.Reader, []byte(role), crypto.Hash(0))
|
||||
case *rsa.PrivateKey, *ecdsa.PrivateKey:
|
||||
hash := sha256.Sum256([]byte(role))
|
||||
signer := key.(crypto.Signer)
|
||||
signature, err = signer.Sign(rand.Reader, hash[:], crypto.SHA256)
|
||||
default:
|
||||
return "", errors.New("unsupported key type")
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return base64.StdEncoding.EncodeToString(signature), nil
|
||||
}
|
||||
@@ -1,110 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent/auth"
|
||||
)
|
||||
|
||||
// NewAuthenticator creates a new instance of Authenticator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
|
||||
// The first argument is typically a *testing.T value.
|
||||
func NewAuthenticator(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *Authenticator {
|
||||
mock := &Authenticator{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
|
||||
// Authenticator is an autogenerated mock type for the Authenticator type
|
||||
type Authenticator struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type Authenticator_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *Authenticator) EXPECT() *Authenticator_Expecter {
|
||||
return &Authenticator_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// AuthenticateUser provides a mock function for the type Authenticator
|
||||
func (_mock *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRole) (context.Context, error) {
|
||||
ret := _mock.Called(ctx, role)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for AuthenticateUser")
|
||||
}
|
||||
|
||||
var r0 context.Context
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, auth.UserRole) (context.Context, error)); ok {
|
||||
return returnFunc(ctx, role)
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, auth.UserRole) context.Context); ok {
|
||||
r0 = returnFunc(ctx, role)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(context.Context)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func(context.Context, auth.UserRole) error); ok {
|
||||
r1 = returnFunc(ctx, role)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// Authenticator_AuthenticateUser_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AuthenticateUser'
|
||||
type Authenticator_AuthenticateUser_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// AuthenticateUser is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - role auth.UserRole
|
||||
func (_e *Authenticator_Expecter) AuthenticateUser(ctx interface{}, role interface{}) *Authenticator_AuthenticateUser_Call {
|
||||
return &Authenticator_AuthenticateUser_Call{Call: _e.mock.On("AuthenticateUser", ctx, role)}
|
||||
}
|
||||
|
||||
func (_c *Authenticator_AuthenticateUser_Call) Run(run func(ctx context.Context, role auth.UserRole)) *Authenticator_AuthenticateUser_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
var arg1 auth.UserRole
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].(auth.UserRole)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Authenticator_AuthenticateUser_Call) Return(context1 context.Context, err error) *Authenticator_AuthenticateUser_Call {
|
||||
_c.Call.Return(context1, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Authenticator_AuthenticateUser_Call) RunAndReturn(run func(ctx context.Context, role auth.UserRole) (context.Context, error)) *Authenticator_AuthenticateUser_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
+47
-83
@@ -3,61 +3,34 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"google.golang.org/grpc/metadata"
|
||||
"reflect"
|
||||
)
|
||||
|
||||
var _ fmt.Stringer = (*Datasets)(nil)
|
||||
var (
|
||||
_ fmt.Stringer = (*Datasets)(nil)
|
||||
_ fmt.Stringer = (*Algorithms)(nil)
|
||||
)
|
||||
|
||||
type AgentConfig struct {
|
||||
CertFile string `json:"cert_file,omitempty"`
|
||||
KeyFile string `json:"server_key,omitempty"`
|
||||
ServerCAFile string `json:"server_ca_file,omitempty"`
|
||||
ClientCAFile string `json:"client_ca_file,omitempty"`
|
||||
AttestedTls bool `json:"attested_tls,omitempty"`
|
||||
}
|
||||
|
||||
// ResourceSource specifies the location of a remote encrypted resource.
|
||||
type ResourceSource struct {
|
||||
// Type is the type of resource source.
|
||||
// Supported values: "oci-image", "s3", "gcs", "https", "http"
|
||||
Type string `json:"type,omitempty"`
|
||||
// URL is the location of the resource.
|
||||
// Examples:
|
||||
// - OCI: "docker://registry/repo:tag"
|
||||
// - S3: "s3://bucket/key"
|
||||
// - GCS: "gs://bucket/key"
|
||||
// - HTTPS: "https://host/path/to/file"
|
||||
// - HTTP: "http://host/path/to/file"
|
||||
URL string `json:"url,omitempty"`
|
||||
// KBSResourcePath is the path to the decryption key in KBS (e.g., "default/key/my-key")
|
||||
KBSResourcePath string `json:"kbs_resource_path,omitempty"`
|
||||
// Encrypted indicates whether the resource is encrypted and requires KBS
|
||||
Encrypted bool `json:"encrypted,omitempty"`
|
||||
}
|
||||
|
||||
// KBSConfig holds configuration for Key Broker Service.
|
||||
type KBSConfig struct {
|
||||
// URL is the KBS endpoint (e.g., "https://kbs.example.com")
|
||||
URL string `json:"url,omitempty"`
|
||||
// Enabled indicates whether to use KBS for key retrieval
|
||||
Enabled bool `json:"enabled,omitempty"`
|
||||
LogLevel string `json:"log_level"`
|
||||
Host string `json:"host"`
|
||||
Port string `json:"port"`
|
||||
CertFile string `json:"cert_file"`
|
||||
KeyFile string `json:"server_key"`
|
||||
ServerCAFile string `json:"server_ca_file"`
|
||||
ClientCAFile string `json:"client_ca_file"`
|
||||
}
|
||||
|
||||
type Computation struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Datasets Datasets `json:"datasets,omitempty"`
|
||||
Algorithm *Algorithm `json:"algorithm,omitempty"`
|
||||
ResultConsumers []ResultConsumer `json:"result_consumers,omitempty"`
|
||||
}
|
||||
|
||||
type ResultConsumer struct {
|
||||
UserKey []byte `json:"user_key,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Datasets Datasets `json:"datasets,omitempty"`
|
||||
Algorithms Algorithms `json:"algorithms,omitempty"`
|
||||
ResultConsumers []string `json:"result_consumers,omitempty"`
|
||||
AgentConfig AgentConfig `json:"agent_config,omitempty"`
|
||||
}
|
||||
|
||||
func (d *Datasets) String() string {
|
||||
@@ -68,51 +41,42 @@ func (d *Datasets) String() string {
|
||||
return string(dat)
|
||||
}
|
||||
|
||||
func (a *Algorithms) String() string {
|
||||
dat, err := json.Marshal(a)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return string(dat)
|
||||
}
|
||||
|
||||
type Dataset struct {
|
||||
Dataset []byte `json:"-"`
|
||||
Hash [32]byte `json:"hash,omitempty"`
|
||||
UserKey []byte `json:"user_key,omitempty"`
|
||||
Filename string `json:"filename,omitempty"`
|
||||
Source *ResourceSource `json:"source,omitempty"` // Optional remote source
|
||||
Decompress bool `json:"decompress,omitempty"`
|
||||
KBS *KBSConfig `json:"kbs,omitempty"`
|
||||
Dataset []byte `json:"-"`
|
||||
Hash [32]byte `json:"hash,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
}
|
||||
|
||||
type Datasets []Dataset
|
||||
|
||||
type Algorithm struct {
|
||||
Algorithm []byte `json:"-"`
|
||||
Hash [32]byte `json:"hash,omitempty"`
|
||||
UserKey []byte `json:"user_key,omitempty"`
|
||||
Requirements []byte `json:"-"`
|
||||
Source *ResourceSource `json:"source,omitempty"` // Optional remote source
|
||||
AlgoType string `json:"algo_type,omitempty"`
|
||||
AlgoArgs []string `json:"algo_args,omitempty"`
|
||||
KBS *KBSConfig `json:"kbs,omitempty"`
|
||||
Algorithm []byte `json:"-"`
|
||||
Hash [32]byte `json:"hash,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
}
|
||||
|
||||
type ManifestIndexKey struct{}
|
||||
type Algorithms []Algorithm
|
||||
|
||||
func IndexToContext(ctx context.Context, index int) context.Context {
|
||||
return context.WithValue(ctx, ManifestIndexKey{}, index)
|
||||
}
|
||||
|
||||
func IndexFromContext(ctx context.Context) (int, bool) {
|
||||
index, ok := ctx.Value(ManifestIndexKey{}).(int)
|
||||
return index, ok
|
||||
}
|
||||
|
||||
const DecompressKey = "decompress"
|
||||
|
||||
func DecompressFromContext(ctx context.Context) bool {
|
||||
vals := metadata.ValueFromIncomingContext(ctx, DecompressKey)
|
||||
if len(vals) == 0 {
|
||||
return false
|
||||
func containsID(slice interface{}, id string) int {
|
||||
rangeOnMe := reflect.ValueOf(slice)
|
||||
for i := 0; i < rangeOnMe.Len(); i++ {
|
||||
s := rangeOnMe.Index(i)
|
||||
f := s.FieldByName("ID")
|
||||
if f.IsValid() {
|
||||
if f.Interface() == id {
|
||||
return i
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return vals[0] == "true"
|
||||
}
|
||||
|
||||
func DecompressToContext(ctx context.Context, decompress bool) context.Context {
|
||||
return metadata.AppendToOutgoingContext(ctx, DecompressKey, fmt.Sprintf("%t", decompress))
|
||||
return -1
|
||||
}
|
||||
|
||||
@@ -1,130 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
func TestDatasetsString(t *testing.T) {
|
||||
datasets := Datasets{
|
||||
{
|
||||
Hash: [32]byte{1, 2, 3},
|
||||
UserKey: []byte("user_key"),
|
||||
Filename: "test.dat",
|
||||
},
|
||||
}
|
||||
|
||||
expected := `[{"hash":[1,2,3,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],"user_key":"dXNlcl9rZXk=","filename":"test.dat"}]`
|
||||
result := datasets.String()
|
||||
|
||||
if result != expected {
|
||||
t.Errorf("Datasets.String() = %v, want %v", result, expected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIndexToContext(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
index := 5
|
||||
|
||||
newCtx := IndexToContext(ctx, index)
|
||||
result, ok := IndexFromContext(newCtx)
|
||||
|
||||
if !ok {
|
||||
t.Errorf("IndexFromContext() ok = false, want true")
|
||||
}
|
||||
|
||||
if result != index {
|
||||
t.Errorf("IndexFromContext() = %v, want %v", result, index)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecompressFromContext(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
ctx context.Context
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "No decompress metadata",
|
||||
ctx: context.Background(),
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "Decompress true",
|
||||
ctx: metadata.NewIncomingContext(
|
||||
context.Background(),
|
||||
metadata.Pairs(DecompressKey, "true"),
|
||||
),
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Decompress false",
|
||||
ctx: metadata.NewIncomingContext(
|
||||
context.Background(),
|
||||
metadata.Pairs(DecompressKey, "false"),
|
||||
),
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := DecompressFromContext(tt.ctx)
|
||||
if result != tt.expected {
|
||||
t.Errorf("DecompressFromContext() = %v, want %v", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecompressToContext(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
decompress := true
|
||||
|
||||
newCtx := DecompressToContext(ctx, decompress)
|
||||
md, ok := metadata.FromOutgoingContext(newCtx)
|
||||
|
||||
if !ok {
|
||||
t.Errorf("metadata.FromOutgoingContext() ok = false, want true")
|
||||
}
|
||||
|
||||
vals := md.Get(DecompressKey)
|
||||
if len(vals) != 1 {
|
||||
t.Errorf("len(md.Get(DecompressKey)) = %v, want 1", len(vals))
|
||||
}
|
||||
|
||||
if vals[0] != "true" {
|
||||
t.Errorf("md.Get(DecompressKey)[0] = %v, want 'true'", vals[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentConfigJSON(t *testing.T) {
|
||||
cfg := AgentConfig{
|
||||
CertFile: "cert.pem",
|
||||
KeyFile: "key.pem",
|
||||
ServerCAFile: "server-ca.pem",
|
||||
ClientCAFile: "client-ca.pem",
|
||||
AttestedTls: true,
|
||||
}
|
||||
|
||||
data, err := json.Marshal(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to marshal AgentConfig: %v", err)
|
||||
}
|
||||
|
||||
var unmarshaledConfig AgentConfig
|
||||
err = json.Unmarshal(data, &unmarshaledConfig)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to unmarshal AgentConfig: %v", err)
|
||||
}
|
||||
|
||||
if !reflect.DeepEqual(cfg, unmarshaledConfig) {
|
||||
t.Errorf("Unmarshaled config does not match original. Got %+v, want %+v", unmarshaledConfig, cfg)
|
||||
}
|
||||
}
|
||||
@@ -1,451 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"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"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"github.com/ultravioletrs/cocos/pkg/clients/grpc"
|
||||
"github.com/ultravioletrs/cocos/pkg/ingress"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
const (
|
||||
reconnectInterval = 5 * time.Second
|
||||
sendTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
var (
|
||||
errCorruptedManifest = errors.New("received manifest may be corrupted")
|
||||
errUnknownMessageType = 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
|
||||
ingressProxy ingress.ProxyServer
|
||||
storage storage.Storage
|
||||
reconnectFn func(context.Context) (grpc.Client, cvms.Service_ProcessClient, error)
|
||||
grpcClient grpc.Client
|
||||
}
|
||||
|
||||
// NewClient returns new gRPC client instance.
|
||||
func NewClient(stream cvms.Service_ProcessClient, svc agent.Service, messageQueue chan *cvms.ClientStreamMessage, logger *slog.Logger, sp server.AgentServer, ingressProxy ingress.ProxyServer, storageDir string, reconnectFn func(context.Context) (grpc.Client, cvms.Service_ProcessClient, error), grpcClient grpc.Client) (*CVMSClient, error) {
|
||||
store, err := storage.NewFileStorage(storageDir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &CVMSClient{
|
||||
stream: stream,
|
||||
svc: svc,
|
||||
messageQueue: messageQueue,
|
||||
logger: logger,
|
||||
runReqManager: newRunRequestManager(),
|
||||
sp: sp,
|
||||
ingressProxy: ingressProxy,
|
||||
storage: store,
|
||||
reconnectFn: reconnectFn,
|
||||
grpcClient: grpcClient,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (client *CVMSClient) Process(ctx context.Context, cancel context.CancelFunc) error {
|
||||
for {
|
||||
err := client.processWithRetry(ctx)
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
client.logger.Info("Connection lost, attempting to reconnect...", "error", err)
|
||||
time.Sleep(reconnectInterval)
|
||||
|
||||
grpcClient, stream, err := client.reconnectFn(ctx)
|
||||
if err != nil {
|
||||
client.logger.Error("Failed to reconnect", "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
client.mu.Lock()
|
||||
client.stream = stream
|
||||
client.grpcClient = grpcClient
|
||||
client.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func (client *CVMSClient) processWithRetry(ctx context.Context) error {
|
||||
eg, ctx := errgroup.WithContext(ctx)
|
||||
|
||||
eg.Go(func() error {
|
||||
return client.handleIncomingMessages(ctx)
|
||||
})
|
||||
|
||||
eg.Go(func() error {
|
||||
return client.handleOutgoingMessages(ctx)
|
||||
})
|
||||
|
||||
return eg.Wait()
|
||||
}
|
||||
|
||||
func (client *CVMSClient) handleIncomingMessages(ctx context.Context) error {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
req, err := client.stream.Recv()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := client.processIncomingMessage(ctx, req); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (client *CVMSClient) handleOutgoingMessages(ctx context.Context) error {
|
||||
pendingMsgs, err := client.storage.Load()
|
||||
if err != nil {
|
||||
client.logger.Error("Failed to load pending messages", "error", err)
|
||||
} else {
|
||||
client.sendPendingMessages(pendingMsgs)
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case msg := <-client.messageQueue:
|
||||
if err := client.sendStreamMessage(msg); err != nil {
|
||||
if err := client.storage.Add(msg); err != nil {
|
||||
client.logger.Error("Failed to store pending message", "error", err)
|
||||
}
|
||||
client.logger.Error("Failed to send message, stored for retry", "error", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (client *CVMSClient) sendStreamMessage(msg *cvms.ClientStreamMessage) error {
|
||||
client.mu.Lock()
|
||||
defer client.mu.Unlock()
|
||||
|
||||
return client.stream.Send(msg)
|
||||
}
|
||||
|
||||
func (client *CVMSClient) sendPendingMessages(pending []storage.Message) {
|
||||
for _, pm := range pending {
|
||||
if err := client.sendStreamMessage(pm.Message); err != nil {
|
||||
if err := client.storage.Add(pm.Message); err != nil {
|
||||
client.logger.Error("Failed to store pending message", "error", err)
|
||||
}
|
||||
client.logger.Error("Failed to resend pending message", "error", err)
|
||||
} else {
|
||||
client.logger.Info("Successfully resent pending message")
|
||||
}
|
||||
}
|
||||
|
||||
if err := client.storage.Clear(); err != nil {
|
||||
client.logger.Error("Failed to clear pending messages", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (client *CVMSClient) processIncomingMessage(ctx context.Context, req *cvms.ServerStreamMessage) error {
|
||||
switch mes := req.Message.(type) {
|
||||
case *cvms.ServerStreamMessage_RunReqChunks:
|
||||
return client.handleRunReqChunks(ctx, mes)
|
||||
case *cvms.ServerStreamMessage_RunReq:
|
||||
client.logger.Info("Starting computation execution from non-chunked run request", "computationId", mes.RunReq.Id, "name", mes.RunReq.Name)
|
||||
go client.executeRun(ctx, mes.RunReq)
|
||||
case *cvms.ServerStreamMessage_StopComputation:
|
||||
go client.handleStopComputation(ctx, mes)
|
||||
case *cvms.ServerStreamMessage_AgentStateReq:
|
||||
client.handleAgentStateReq(mes)
|
||||
case *cvms.ServerStreamMessage_DisconnectReq:
|
||||
client.logger.Info("Received disconnect request")
|
||||
client.mu.Lock()
|
||||
if err := client.grpcClient.Close(); err != nil {
|
||||
client.logger.Error("Failed to close gRPC client", "error", err)
|
||||
}
|
||||
client.mu.Unlock()
|
||||
default:
|
||||
return errUnknownMessageType
|
||||
}
|
||||
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 {
|
||||
client.logger.Debug("Received RunReq chunk", "id", msg.RunReqChunks.Id, "size", len(msg.RunReqChunks.Data), "isLast", msg.RunReqChunks.IsLast)
|
||||
buffer, complete := client.runReqManager.addChunk(msg.RunReqChunks.Id, msg.RunReqChunks.Data, msg.RunReqChunks.IsLast)
|
||||
|
||||
if complete {
|
||||
client.logger.Info("Received complete computation run request", "id", msg.RunReqChunks.Id, "totalSize", len(buffer))
|
||||
var runReq cvms.ComputationRunReq
|
||||
if err := proto.Unmarshal(buffer, &runReq); err != nil {
|
||||
return errors.Wrap(err, errCorruptedManifest)
|
||||
}
|
||||
|
||||
client.logger.Info("Starting computation execution", "computationId", runReq.Id, "name", runReq.Name)
|
||||
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,
|
||||
AlgoType: runReq.Algorithm.AlgoType,
|
||||
}
|
||||
// Copy remote source if configured
|
||||
if runReq.Algorithm.Source != nil {
|
||||
ac.Algorithm.Source = &agent.ResourceSource{
|
||||
URL: runReq.Algorithm.Source.Url,
|
||||
KBSResourcePath: runReq.Algorithm.Source.KbsResourcePath,
|
||||
Encrypted: runReq.Algorithm.Source.Encrypted,
|
||||
}
|
||||
}
|
||||
ac.Algorithm.AlgoArgs = runReq.Algorithm.AlgoArgs
|
||||
if runReq.Algorithm.Kbs != nil {
|
||||
ac.Algorithm.KBS = &agent.KBSConfig{
|
||||
URL: runReq.Algorithm.Kbs.Url,
|
||||
Enabled: runReq.Algorithm.Kbs.Enabled,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, ds := range runReq.Datasets {
|
||||
dataset := agent.Dataset{
|
||||
Hash: [32]byte(ds.Hash),
|
||||
UserKey: ds.UserKey,
|
||||
Filename: ds.Filename,
|
||||
}
|
||||
// Copy remote source if configured
|
||||
if ds.Source != nil {
|
||||
dataset.Source = &agent.ResourceSource{
|
||||
URL: ds.Source.Url,
|
||||
KBSResourcePath: ds.Source.KbsResourcePath,
|
||||
Encrypted: ds.Source.Encrypted,
|
||||
}
|
||||
}
|
||||
dataset.Decompress = ds.Decompress
|
||||
if ds.Kbs != nil {
|
||||
dataset.KBS = &agent.KBSConfig{
|
||||
URL: ds.Kbs.Url,
|
||||
Enabled: ds.Kbs.Enabled,
|
||||
}
|
||||
}
|
||||
ac.Datasets = append(ac.Datasets, dataset)
|
||||
}
|
||||
|
||||
for _, rc := range runReq.ResultConsumers {
|
||||
ac.ResultConsumers = append(ac.ResultConsumers, agent.ResultConsumer{
|
||||
UserKey: rc.UserKey,
|
||||
})
|
||||
}
|
||||
|
||||
// Check if the agent is in the correct state to initialize a new computation.
|
||||
// If the agent is already processing this computation (e.g., after a reconnection),
|
||||
// skip initialization to avoid state errors.
|
||||
currentState := client.svc.State()
|
||||
if currentState != "ReceivingManifest" {
|
||||
client.logger.Info("Agent already processing computation, skipping initialization", "state", currentState, "computationId", runReq.Id)
|
||||
return
|
||||
}
|
||||
|
||||
if err := client.svc.InitComputation(ctx, ac); err != nil {
|
||||
client.logger.Warn(err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
ccPlatform := attestation.CCPlatform()
|
||||
|
||||
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{
|
||||
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()
|
||||
}
|
||||
|
||||
// Start ingress proxy if available
|
||||
if client.ingressProxy != nil {
|
||||
if err := client.ingressProxy.Start(
|
||||
ingress.AgentConfigToProxyConfig(agent.AgentConfig{
|
||||
CertFile: runReq.AgentConfig.CertFile,
|
||||
KeyFile: runReq.AgentConfig.KeyFile,
|
||||
ServerCAFile: runReq.AgentConfig.ServerCaFile,
|
||||
ClientCAFile: runReq.AgentConfig.ClientCaFile,
|
||||
AttestedTls: runReq.AgentConfig.AttestedTls,
|
||||
}),
|
||||
ingress.ComputationToProxyContext(ac),
|
||||
); err != nil {
|
||||
client.logger.Warn(fmt.Sprintf("failed to start ingress proxy: %s", err.Error()))
|
||||
}
|
||||
}
|
||||
|
||||
defer func() {
|
||||
if ccPlatform == attestation.Azure || ccPlatform == attestation.SNPvTPM {
|
||||
cmpJson, err := json.Marshal(ac)
|
||||
if err != nil {
|
||||
client.logger.Error(err.Error())
|
||||
return
|
||||
}
|
||||
if err = vtpm.ExtendPCR(vtpm.PCR16, cmpJson); err != nil {
|
||||
client.logger.Error(err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
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()
|
||||
}
|
||||
// Stop ingress proxy if available
|
||||
if client.ingressProxy != nil {
|
||||
if err := client.ingressProxy.Stop(); err != nil {
|
||||
client.logger.Warn(fmt.Sprintf("failed to stop ingress proxy: %s", 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,654 +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"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
|
||||
servermocks "github.com/ultravioletrs/cocos/agent/cvms/server/mocks"
|
||||
"github.com/ultravioletrs/cocos/agent/mocks"
|
||||
pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
|
||||
clientmocks "github.com/ultravioletrs/cocos/pkg/clients/grpc/mocks"
|
||||
"github.com/ultravioletrs/cocos/pkg/ingress"
|
||||
"golang.org/x/crypto/sha3"
|
||||
"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)
|
||||
}
|
||||
|
||||
// mockIngressProxy is a mock implementation of the ingress proxy.
|
||||
type mockIngressProxy struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *mockIngressProxy) Start(config ingress.ProxyConfig, ctx ingress.ProxyContext) error {
|
||||
args := m.Called(config, ctx)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *mockIngressProxy) Stop() error {
|
||||
args := m.Called()
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func TestManagerClient_Process(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setupMocks func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client)
|
||||
expectError bool
|
||||
errorMsg string
|
||||
}{
|
||||
{
|
||||
name: "Stop computation",
|
||||
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
|
||||
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, grpcClient *clientmocks.Client) {
|
||||
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: "Agent state request",
|
||||
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
|
||||
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
|
||||
Message: &cvms.ServerStreamMessage_AgentStateReq{
|
||||
AgentStateReq: &cvms.AgentStateReq{
|
||||
Id: "test-agent",
|
||||
},
|
||||
},
|
||||
}, nil)
|
||||
mockStream.On("Send", mock.Anything).Return(nil)
|
||||
mockSvc.On("State").Return("test-state")
|
||||
},
|
||||
expectError: true,
|
||||
errorMsg: "context deadline exceeded",
|
||||
},
|
||||
{
|
||||
name: "Disconnect request",
|
||||
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
|
||||
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
|
||||
Message: &cvms.ServerStreamMessage_DisconnectReq{},
|
||||
}, nil)
|
||||
mockStream.On("Send", mock.Anything).Return(nil)
|
||||
grpcClient.On("Close").Return(nil)
|
||||
},
|
||||
expectError: true,
|
||||
errorMsg: "context deadline exceeded",
|
||||
},
|
||||
{
|
||||
name: "Receive error",
|
||||
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
|
||||
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{}, assert.AnError)
|
||||
},
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage)
|
||||
logger := mglog.NewMock()
|
||||
|
||||
go func() {
|
||||
<-messageQueue
|
||||
}()
|
||||
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
tc.setupMocks(mockStream, mockSvc, mockServerSvc, grpcClient)
|
||||
|
||||
err = client.Process(ctx, cancel)
|
||||
|
||||
if tc.expectError {
|
||||
assert.Error(t, err)
|
||||
if tc.errorMsg != "" {
|
||||
assert.Contains(t, err.Error(), tc.errorMsg)
|
||||
}
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerClient_handleRunReqChunks(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
runReq := &cvms.ComputationRunReq{
|
||||
Id: "test-id",
|
||||
Datasets: []*cvms.Dataset{
|
||||
{
|
||||
Hash: sha3.New256().Sum([]byte("test-dataset")),
|
||||
},
|
||||
},
|
||||
Algorithm: &cvms.Algorithm{
|
||||
Hash: sha3.New256().Sum([]byte("test-algorithm")),
|
||||
},
|
||||
ResultConsumers: []*cvms.ResultConsumer{
|
||||
{
|
||||
UserKey: []byte("test-consumer"),
|
||||
},
|
||||
},
|
||||
}
|
||||
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("State").Return("ReceivingManifest")
|
||||
mockSvc.On("InitComputation", mock.Anything, mock.Anything).Return(nil)
|
||||
mockServerSvc.On("Start", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
|
||||
err = client.handleRunReqChunks(context.Background(), chunk1)
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = client.handleRunReqChunks(context.Background(), chunk2)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Wait for the goroutine to finish
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
mockSvc.AssertExpectations(t)
|
||||
assert.Len(t, messageQueue, 1)
|
||||
|
||||
msg := <-messageQueue
|
||||
runRes, ok := msg.Message.(*cvms.ClientStreamMessage_RunRes)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "test-id", runRes.RunRes.ComputationId)
|
||||
}
|
||||
|
||||
func TestManagerClient_handleStopComputation(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
stopReq := &cvms.ServerStreamMessage_StopComputation{
|
||||
StopComputation: &cvms.StopComputation{
|
||||
ComputationId: "test-comp-id",
|
||||
},
|
||||
}
|
||||
|
||||
mockSvc.On("StopComputation", mock.Anything).Return(nil)
|
||||
mockServerSvc.On("Stop").Return(nil)
|
||||
|
||||
client.handleStopComputation(context.Background(), stopReq)
|
||||
|
||||
// Wait for the goroutine to finish
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
mockSvc.AssertExpectations(t)
|
||||
assert.Len(t, messageQueue, 1)
|
||||
|
||||
msg := <-messageQueue
|
||||
stopRes, ok := msg.Message.(*cvms.ClientStreamMessage_StopComputationRes)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "test-comp-id", stopRes.StopComputationRes.ComputationId)
|
||||
assert.Empty(t, stopRes.StopComputationRes.Message)
|
||||
}
|
||||
|
||||
func TestManagerClient_timeoutRequest(t *testing.T) {
|
||||
rm := newRunRequestManager()
|
||||
rm.requests["test-id"] = &runRequest{
|
||||
timer: time.NewTimer(100 * time.Millisecond),
|
||||
buffer: []byte("test-data"),
|
||||
lastChunk: time.Now(),
|
||||
}
|
||||
|
||||
rm.timeoutRequest("test-id")
|
||||
|
||||
assert.Len(t, rm.requests, 0)
|
||||
}
|
||||
|
||||
// TestManagerClient_sendPendingMessages tests sending pending messages on reconnection.
|
||||
func TestManagerClient_sendPendingMessages(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Add a pending message to storage
|
||||
testMsg := &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_RunRes{
|
||||
RunRes: &cvms.RunResponse{
|
||||
ComputationId: "test-id",
|
||||
},
|
||||
},
|
||||
}
|
||||
err = client.storage.Add(testMsg)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Mock successful send
|
||||
mockStream.On("Send", mock.Anything).Return(nil).Once()
|
||||
|
||||
// Load and send pending messages
|
||||
pending, err := client.storage.Load()
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, pending, 1)
|
||||
|
||||
client.sendPendingMessages(pending)
|
||||
|
||||
mockStream.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// TestManagerClient_sendPendingMessagesWithError tests pending message send failure.
|
||||
func TestManagerClient_sendPendingMessagesWithError(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
testMsg := &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_RunRes{
|
||||
RunRes: &cvms.RunResponse{
|
||||
ComputationId: "test-id",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Mock failed send
|
||||
mockStream.On("Send", mock.Anything).Return(assert.AnError)
|
||||
|
||||
pending := []storage.Message{
|
||||
{
|
||||
Message: testMsg,
|
||||
Time: time.Now(),
|
||||
},
|
||||
}
|
||||
|
||||
client.sendPendingMessages(pending)
|
||||
|
||||
mockStream.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// TestManagerClient_addChunkTimeout tests chunk timeout in runRequestManager.
|
||||
func TestManagerClient_addChunkTimeout(t *testing.T) {
|
||||
rm := newRunRequestManager()
|
||||
|
||||
// Add first chunk
|
||||
chunk1 := []byte("chunk1")
|
||||
buffer, complete := rm.addChunk("test-id", chunk1, false)
|
||||
assert.Nil(t, buffer)
|
||||
assert.False(t, complete)
|
||||
|
||||
// Verify request exists
|
||||
rm.mu.Lock()
|
||||
assert.Contains(t, rm.requests, "test-id")
|
||||
rm.mu.Unlock()
|
||||
|
||||
// Wait for timeout
|
||||
time.Sleep(35 * time.Second) // runReqTimeout is 30 seconds
|
||||
|
||||
// Verify request was removed
|
||||
rm.mu.Lock()
|
||||
assert.NotContains(t, rm.requests, "test-id")
|
||||
rm.mu.Unlock()
|
||||
}
|
||||
|
||||
// TestManagerClient_addChunkMultiple tests adding multiple chunks.
|
||||
func TestManagerClient_addChunkMultiple(t *testing.T) {
|
||||
rm := newRunRequestManager()
|
||||
|
||||
chunk1 := []byte("chunk1")
|
||||
chunk2 := []byte("chunk2")
|
||||
chunk3 := []byte("chunk3")
|
||||
|
||||
// Add chunks
|
||||
buffer, complete := rm.addChunk("test-id", chunk1, false)
|
||||
assert.Nil(t, buffer)
|
||||
assert.False(t, complete)
|
||||
|
||||
buffer, complete = rm.addChunk("test-id", chunk2, false)
|
||||
assert.Nil(t, buffer)
|
||||
assert.False(t, complete)
|
||||
|
||||
buffer, complete = rm.addChunk("test-id", chunk3, true)
|
||||
assert.NotNil(t, buffer)
|
||||
assert.True(t, complete)
|
||||
|
||||
expected := append(append(chunk1, chunk2...), chunk3...)
|
||||
assert.Equal(t, expected, buffer)
|
||||
}
|
||||
|
||||
// TestManagerClient_handleStopComputationWithIngressProxy tests stop with ingress proxy.
|
||||
func TestManagerClient_handleStopComputationWithIngressProxy(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
mockIngressProxy := new(mockIngressProxy)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, mockIngressProxy, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
stopReq := &cvms.ServerStreamMessage_StopComputation{
|
||||
StopComputation: &cvms.StopComputation{
|
||||
ComputationId: "test-comp-id",
|
||||
},
|
||||
}
|
||||
|
||||
mockSvc.On("StopComputation", mock.Anything).Return(nil)
|
||||
mockServerSvc.On("Stop").Return(nil)
|
||||
mockIngressProxy.On("Stop").Return(nil)
|
||||
|
||||
client.handleStopComputation(context.Background(), stopReq)
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
mockSvc.AssertExpectations(t)
|
||||
mockServerSvc.AssertExpectations(t)
|
||||
mockIngressProxy.AssertExpectations(t)
|
||||
assert.Len(t, messageQueue, 1)
|
||||
}
|
||||
|
||||
// TestManagerClient_handleStopComputationWithIngressProxyError tests stop with ingress proxy error.
|
||||
func TestManagerClient_handleStopComputationWithIngressProxyError(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
mockIngressProxy := new(mockIngressProxy)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, mockIngressProxy, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
stopReq := &cvms.ServerStreamMessage_StopComputation{
|
||||
StopComputation: &cvms.StopComputation{
|
||||
ComputationId: "test-comp-id",
|
||||
},
|
||||
}
|
||||
|
||||
mockSvc.On("StopComputation", mock.Anything).Return(nil)
|
||||
mockServerSvc.On("Stop").Return(nil)
|
||||
mockIngressProxy.On("Stop").Return(assert.AnError)
|
||||
|
||||
client.handleStopComputation(context.Background(), stopReq)
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
mockIngressProxy.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// TestManagerClient_sendMessage tests sendMessage with timeout.
|
||||
func TestManagerClient_sendMessage(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 1)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
msg := &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_RunRes{
|
||||
RunRes: &cvms.RunResponse{
|
||||
ComputationId: "test-id",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
client.sendMessage(msg)
|
||||
|
||||
select {
|
||||
case received := <-messageQueue:
|
||||
assert.Equal(t, msg, received)
|
||||
case <-time.After(1 * time.Second):
|
||||
t.Fatal("Message not received")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManagerClient_sendMessageTimeout tests sendMessage timeout when queue is full.
|
||||
func TestManagerClient_sendMessageTimeout(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage) // No buffer
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
msg := &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_RunRes{
|
||||
RunRes: &cvms.RunResponse{
|
||||
ComputationId: "test-id",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Don't read from queue, so sendMessage will timeout
|
||||
client.sendMessage(msg)
|
||||
|
||||
// Should complete without blocking
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
|
||||
// TestManagerClient_handleRunReqChunksWithRemoteSource tests handling run request with remote source.
|
||||
func TestManagerClient_handleRunReqChunksWithRemoteSource(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
runReq := &cvms.ComputationRunReq{
|
||||
Id: "test-id-remote",
|
||||
Name: "test-computation",
|
||||
Description: "test description",
|
||||
Datasets: []*cvms.Dataset{
|
||||
{
|
||||
Hash: sha3.New256().Sum([]byte("test-dataset")),
|
||||
Filename: "data.csv",
|
||||
Source: &cvms.Source{
|
||||
Type: "oci-image",
|
||||
Url: "docker://registry.example.com/data:v1",
|
||||
KbsResourcePath: "default/key/data-key",
|
||||
Encrypted: true,
|
||||
},
|
||||
Decompress: true,
|
||||
},
|
||||
},
|
||||
Algorithm: &cvms.Algorithm{
|
||||
Hash: sha3.New256().Sum([]byte("test-algorithm")),
|
||||
AlgoType: "python",
|
||||
AlgoArgs: []string{"--verbose"},
|
||||
Source: &cvms.Source{
|
||||
Type: "oci-image",
|
||||
Url: "docker://registry.example.com/algo:v1",
|
||||
KbsResourcePath: "default/key/algo-key",
|
||||
Encrypted: true,
|
||||
},
|
||||
Kbs: &cvms.KBSConfig{
|
||||
Url: "https://kbs.example.com:8080",
|
||||
Enabled: true,
|
||||
},
|
||||
},
|
||||
|
||||
ResultConsumers: []*cvms.ResultConsumer{
|
||||
{
|
||||
UserKey: []byte("test-consumer"),
|
||||
},
|
||||
},
|
||||
}
|
||||
runReqBytes, _ := proto.Marshal(runReq)
|
||||
|
||||
chunk := &cvms.ServerStreamMessage_RunReqChunks{
|
||||
RunReqChunks: &cvms.RunReqChunks{
|
||||
Id: "chunk-remote-1",
|
||||
Data: runReqBytes,
|
||||
IsLast: true,
|
||||
},
|
||||
}
|
||||
|
||||
mockSvc.On("State").Return("ReceivingManifest")
|
||||
mockSvc.On("InitComputation", mock.Anything, mock.MatchedBy(func(c agent.Computation) bool {
|
||||
// Verify Algorithm KBS config is passed
|
||||
if c.Algorithm.KBS == nil || !c.Algorithm.KBS.Enabled || c.Algorithm.KBS.URL != "https://kbs.example.com:8080" {
|
||||
return false
|
||||
}
|
||||
// Verify algorithm source is passed
|
||||
if c.Algorithm.Source == nil ||
|
||||
c.Algorithm.Source.URL != "docker://registry.example.com/algo:v1" ||
|
||||
c.Algorithm.Source.KBSResourcePath != "default/key/algo-key" ||
|
||||
!c.Algorithm.Source.Encrypted {
|
||||
return false
|
||||
}
|
||||
// Verify algorithm type and args
|
||||
if c.Algorithm.AlgoType != "python" || len(c.Algorithm.AlgoArgs) != 1 || c.Algorithm.AlgoArgs[0] != "--verbose" {
|
||||
return false
|
||||
}
|
||||
// Verify dataset source is passed
|
||||
if len(c.Datasets) != 1 ||
|
||||
c.Datasets[0].Source == nil ||
|
||||
c.Datasets[0].Source.URL != "docker://registry.example.com/data:v1" ||
|
||||
c.Datasets[0].Filename != "data.csv" ||
|
||||
!c.Datasets[0].Decompress {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})).Return(nil)
|
||||
mockServerSvc.On("Start", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
|
||||
err = client.handleRunReqChunks(context.Background(), chunk)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Wait for the goroutine to finish
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
mockSvc.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// TestManagerClient_handleRunReqChunksAlreadyProcessing tests skipping init when already processing.
|
||||
func TestManagerClient_handleRunReqChunksAlreadyProcessing(t *testing.T) {
|
||||
mockStream := new(mockStream)
|
||||
mockSvc := new(mocks.Service)
|
||||
mockServerSvc := new(servermocks.AgentServer)
|
||||
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
logger := mglog.NewMock()
|
||||
grpcClient := new(clientmocks.Client)
|
||||
|
||||
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
runReq := &cvms.ComputationRunReq{
|
||||
Id: "test-id-processing",
|
||||
Name: "test-computation",
|
||||
}
|
||||
runReqBytes, _ := proto.Marshal(runReq)
|
||||
|
||||
chunk := &cvms.ServerStreamMessage_RunReqChunks{
|
||||
RunReqChunks: &cvms.RunReqChunks{
|
||||
Id: "chunk-processing-1",
|
||||
Data: runReqBytes,
|
||||
IsLast: true,
|
||||
},
|
||||
}
|
||||
|
||||
// Simulate agent already processing a computation
|
||||
mockSvc.On("State").Return("Running")
|
||||
|
||||
err = client.handleRunReqChunks(context.Background(), chunk)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Wait for the goroutine to finish
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// InitComputation should NOT be called since state is not ReceivingManifest
|
||||
mockSvc.AssertNotCalled(t, "InitComputation")
|
||||
}
|
||||
@@ -1,141 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"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")
|
||||
}
|
||||
|
||||
slog.Info("client connected to cvms server", "address", client.Addr.String())
|
||||
|
||||
eg, ctx := errgroup.WithContext(stream.Context())
|
||||
|
||||
eg.Go(func() error {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
slog.Info("receive goroutine context done", "address", client.Addr.String())
|
||||
return ctx.Err()
|
||||
default:
|
||||
req, err := stream.Recv()
|
||||
if err != nil {
|
||||
slog.Error("failed to receive from stream", "address", client.Addr.String(), "error", err)
|
||||
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)
|
||||
slog.Info("send goroutine Run() returned", "address", client.Addr.String())
|
||||
return nil
|
||||
})
|
||||
|
||||
err := eg.Wait()
|
||||
slog.Info("stream closed", "address", client.Addr.String(), "error", err)
|
||||
return err
|
||||
}
|
||||
|
||||
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,273 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/absmach/magistrala/pkg/errors"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"google.golang.org/grpc/credentials"
|
||||
"google.golang.org/grpc/peer"
|
||||
)
|
||||
|
||||
type mockServerStream struct {
|
||||
mock.Mock
|
||||
cvms.Service_ProcessServer
|
||||
}
|
||||
|
||||
func (m *mockServerStream) Send(msg *cvms.ServerStreamMessage) error {
|
||||
args := m.Called(msg)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *mockServerStream) Recv() (*cvms.ClientStreamMessage, error) {
|
||||
args := m.Called()
|
||||
return args.Get(0).(*cvms.ClientStreamMessage), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *mockServerStream) Context() context.Context {
|
||||
args := m.Called()
|
||||
return args.Get(0).(context.Context)
|
||||
}
|
||||
|
||||
type mockService struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *mockService) Run(ctx context.Context, ipAddress string, sendMessage SendFunc, authInfo credentials.AuthInfo) {
|
||||
m.Called(ctx, ipAddress, sendMessage, authInfo)
|
||||
}
|
||||
|
||||
func TestNewServer(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage)
|
||||
mockSvc := new(mockService)
|
||||
|
||||
server := NewServer(incoming, mockSvc)
|
||||
|
||||
assert.NotNil(t, server)
|
||||
assert.IsType(t, &grpcServer{}, server)
|
||||
}
|
||||
|
||||
func TestGrpcServer_Process(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
recvReturn *cvms.ClientStreamMessage
|
||||
recvError error
|
||||
expectedError string
|
||||
}{
|
||||
{
|
||||
name: "Process with context deadline exceeded",
|
||||
recvReturn: &cvms.ClientStreamMessage{},
|
||||
recvError: nil,
|
||||
expectedError: "context deadline exceeded",
|
||||
},
|
||||
{
|
||||
name: "Process with Recv error",
|
||||
recvReturn: &cvms.ClientStreamMessage{},
|
||||
recvError: errors.New("recv error"),
|
||||
expectedError: "recv error",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage, 1)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
mockStream := new(mockServerStream)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
mockStream.On("Context").Return(peer.NewContext(ctx, &peer.Peer{
|
||||
Addr: mockAddr{},
|
||||
AuthInfo: mockAuthInfo{},
|
||||
}))
|
||||
|
||||
if tt.recvError == nil {
|
||||
go func() {
|
||||
for mes := range incoming {
|
||||
assert.NotNil(t, mes)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
mockStream.On("Recv").Return(tt.recvReturn, tt.recvError)
|
||||
mockSvc.On("Run", mock.Anything, "test", mock.Anything, mock.AnythingOfType("mockAuthInfo")).Return()
|
||||
|
||||
err := server.Process(mockStream)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), tt.expectedError)
|
||||
mockStream.AssertExpectations(t)
|
||||
mockSvc.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
mockStream := new(mockServerStream)
|
||||
|
||||
runReq := &cvms.ComputationRunReq{
|
||||
Id: "test-id",
|
||||
}
|
||||
|
||||
largePayload := make([]byte, bufferSize*2)
|
||||
for i := range largePayload {
|
||||
largePayload[i] = byte(i % 256)
|
||||
}
|
||||
runReq.Algorithm = &cvms.Algorithm{}
|
||||
runReq.Algorithm.UserKey = largePayload
|
||||
|
||||
mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(nil).Times(4)
|
||||
|
||||
err := server.sendRunReqInChunks(mockStream, runReq)
|
||||
|
||||
assert.NoError(t, err)
|
||||
mockStream.AssertExpectations(t)
|
||||
|
||||
calls := mockStream.Calls
|
||||
assert.Equal(t, 4, len(calls))
|
||||
|
||||
for i, call := range calls {
|
||||
msg := call.Arguments[0].(*cvms.ServerStreamMessage)
|
||||
chunk := msg.GetRunReqChunks()
|
||||
|
||||
assert.NotNil(t, chunk)
|
||||
assert.Equal(t, "test-id", chunk.Id)
|
||||
|
||||
if i < 3 {
|
||||
assert.False(t, chunk.IsLast)
|
||||
} else {
|
||||
assert.Equal(t, 0, len(chunk.Data))
|
||||
assert.True(t, chunk.IsLast)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type mockAddr struct{}
|
||||
|
||||
func (mockAddr) Network() string { return "test network" }
|
||||
func (mockAddr) String() string { return "test" }
|
||||
|
||||
type mockAuthInfo struct{}
|
||||
|
||||
func (mockAuthInfo) AuthType() string { return "test auth" }
|
||||
|
||||
func TestGrpcServer_ProcessWithMockService(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setupMockFn func(*mockService, *mockServerStream)
|
||||
}{
|
||||
{
|
||||
name: "Run Request Test",
|
||||
setupMockFn: func(mockSvc *mockService, mockStream *mockServerStream) {
|
||||
mockSvc.On("Run", mock.Anything, "test", mock.Anything, mock.AnythingOfType("mockAuthInfo")).
|
||||
Run(func(args mock.Arguments) {
|
||||
sendFunc := args.Get(2).(SendFunc)
|
||||
runReq := &cvms.ComputationRunReq{Id: "test-run-id"}
|
||||
err := sendFunc(&cvms.ServerStreamMessage{
|
||||
Message: &cvms.ServerStreamMessage_RunReq{
|
||||
RunReq: runReq,
|
||||
},
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
}).
|
||||
Return()
|
||||
|
||||
mockStream.On("Send", mock.MatchedBy(func(msg *cvms.ServerStreamMessage) bool {
|
||||
chunks := msg.GetRunReqChunks()
|
||||
return chunks != nil && chunks.Id == "test-run-id"
|
||||
})).Return(nil)
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage, 10)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
go func() {
|
||||
for mes := range incoming {
|
||||
assert.NotNil(t, mes)
|
||||
}
|
||||
}()
|
||||
|
||||
mockStream := new(mockServerStream)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
peerCtx := peer.NewContext(ctx, &peer.Peer{
|
||||
Addr: mockAddr{},
|
||||
AuthInfo: mockAuthInfo{},
|
||||
})
|
||||
|
||||
mockStream.On("Context").Return(peerCtx)
|
||||
mockStream.On("Recv").Return(&cvms.ClientStreamMessage{}, nil).Maybe()
|
||||
|
||||
tt.setupMockFn(mockSvc, mockStream)
|
||||
|
||||
go func() {
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
cancel()
|
||||
}()
|
||||
|
||||
err := server.Process(mockStream)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "context canceled")
|
||||
mockStream.AssertExpectations(t)
|
||||
mockSvc.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrpcServer_sendRunReqInChunksError(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
mockStream := new(mockServerStream)
|
||||
|
||||
runReq := &cvms.ComputationRunReq{
|
||||
Id: "test-id",
|
||||
}
|
||||
|
||||
// Simulate an error when sending
|
||||
mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(errors.New("send error")).Once()
|
||||
|
||||
err := server.sendRunReqInChunks(mockStream, runReq)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "send error")
|
||||
mockStream.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestGrpcServer_ProcessMissingPeerInfo(t *testing.T) {
|
||||
incoming := make(chan *cvms.ClientStreamMessage)
|
||||
mockSvc := new(mockService)
|
||||
server := NewServer(incoming, mockSvc).(*grpcServer)
|
||||
|
||||
mockStream := new(mockServerStream)
|
||||
ctx := context.Background()
|
||||
|
||||
// Return a context without peer info
|
||||
mockStream.On("Context").Return(ctx)
|
||||
|
||||
err := server.Process(mockStream)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "failed to get peer info")
|
||||
mockStream.AssertExpectations(t)
|
||||
}
|
||||
@@ -1,242 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
|
||||
)
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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 for the type Storage
|
||||
func (_mock *Storage) Add(msg *cvms.ClientStreamMessage) error {
|
||||
ret := _mock.Called(msg)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Add")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(*cvms.ClientStreamMessage) error); ok {
|
||||
r0 = returnFunc(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) {
|
||||
var arg0 *cvms.ClientStreamMessage
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(*cvms.ClientStreamMessage)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Storage_Add_Call) Return(err error) *Storage_Add_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Storage_Add_Call) RunAndReturn(run func(msg *cvms.ClientStreamMessage) error) *Storage_Add_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Clear provides a mock function for the type Storage
|
||||
func (_mock *Storage) Clear() error {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Clear")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func() error); ok {
|
||||
r0 = returnFunc()
|
||||
} 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(err error) *Storage_Clear_Call {
|
||||
_c.Call.Return(err)
|
||||
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 for the type Storage
|
||||
func (_mock *Storage) Load() ([]storage.Message, error) {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Load")
|
||||
}
|
||||
|
||||
var r0 []storage.Message
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func() ([]storage.Message, error)); ok {
|
||||
return returnFunc()
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func() []storage.Message); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).([]storage.Message)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func() error); ok {
|
||||
r1 = returnFunc()
|
||||
} 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(messages []storage.Message, err error) *Storage_Load_Call {
|
||||
_c.Call.Return(messages, err)
|
||||
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 for the type Storage
|
||||
func (_mock *Storage) Save(messages []storage.Message) error {
|
||||
ret := _mock.Called(messages)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Save")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func([]storage.Message) error); ok {
|
||||
r0 = returnFunc(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) {
|
||||
var arg0 []storage.Message
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].([]storage.Message)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Storage_Save_Call) Return(err error) *Storage_Save_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Storage_Save_Call) RunAndReturn(run func(messages []storage.Message) error) *Storage_Save_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -1,450 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package storage
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
)
|
||||
|
||||
func createTempDir(t *testing.T) string {
|
||||
tmpDir, err := os.MkdirTemp("", "storage_test_*")
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
os.RemoveAll(tmpDir)
|
||||
})
|
||||
return tmpDir
|
||||
}
|
||||
|
||||
func createTestMessage(content string) *cvms.ClientStreamMessage {
|
||||
return &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_RunRes{
|
||||
RunRes: &cvms.RunResponse{
|
||||
Error: "",
|
||||
ComputationId: content,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewFileStorage(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
storageDir string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "valid directory",
|
||||
storageDir: createTempDir(t),
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "non-existent directory gets created",
|
||||
storageDir: filepath.Join(createTempDir(t), "subdir"),
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "invalid directory path",
|
||||
storageDir: "/invalid/path/that/cannot/be/created",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
storage, err := NewFileStorage(tt.storageDir)
|
||||
|
||||
if tt.expectError {
|
||||
assert.Error(t, err)
|
||||
assert.Nil(t, storage)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, storage)
|
||||
assert.Equal(t, filepath.Join(tt.storageDir, "pending_messages.json"), storage.path)
|
||||
assert.Empty(t, storage.msgs)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileStorage_Load(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setupFile func(string) error
|
||||
expectedMsgs int
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "load from non-existent file",
|
||||
setupFile: func(path string) error {
|
||||
// Don't create file
|
||||
return nil
|
||||
},
|
||||
expectedMsgs: 0,
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "load from empty file",
|
||||
setupFile: func(path string) error {
|
||||
return os.WriteFile(path, []byte("[]"), 0o644)
|
||||
},
|
||||
expectedMsgs: 0,
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "load from corrupted file",
|
||||
setupFile: func(path string) error {
|
||||
return os.WriteFile(path, []byte("invalid json"), 0o644)
|
||||
},
|
||||
expectedMsgs: 0,
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
tmpDir := createTempDir(t)
|
||||
storage, err := NewFileStorage(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = tt.setupFile(storage.path)
|
||||
require.NoError(t, err)
|
||||
|
||||
msgs, err := storage.Load()
|
||||
|
||||
if tt.expectError {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, msgs, tt.expectedMsgs)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileStorage_Save(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
messages []Message
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "save empty messages",
|
||||
messages: []Message{},
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "save single message",
|
||||
messages: []Message{
|
||||
{
|
||||
Message: createTestMessage("test"),
|
||||
Time: time.Now(),
|
||||
},
|
||||
},
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "save multiple messages",
|
||||
messages: []Message{
|
||||
{
|
||||
Message: createTestMessage("test1"),
|
||||
Time: time.Now(),
|
||||
},
|
||||
{
|
||||
Message: createTestMessage("test2"),
|
||||
Time: time.Now().Add(time.Second),
|
||||
},
|
||||
},
|
||||
expectError: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
tmpDir := createTempDir(t)
|
||||
storage, err := NewFileStorage(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = storage.Save(tt.messages)
|
||||
|
||||
if tt.expectError {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify file was written correctly
|
||||
_, err := os.ReadFile(storage.path)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify internal state was updated
|
||||
assert.Equal(t, tt.messages, storage.msgs)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileStorage_Add(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialMsgs []Message
|
||||
newMessage *cvms.ClientStreamMessage
|
||||
expectError bool
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "add to empty storage",
|
||||
initialMsgs: []Message{},
|
||||
newMessage: createTestMessage("new"),
|
||||
expectError: false,
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "add to existing messages",
|
||||
initialMsgs: []Message{
|
||||
{
|
||||
Message: createTestMessage("existing"),
|
||||
Time: time.Now(),
|
||||
},
|
||||
},
|
||||
newMessage: createTestMessage("new"),
|
||||
expectError: false,
|
||||
expectedCount: 2,
|
||||
},
|
||||
{
|
||||
name: "add nil message",
|
||||
initialMsgs: []Message{},
|
||||
newMessage: nil,
|
||||
expectError: false,
|
||||
expectedCount: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
tmpDir := createTempDir(t)
|
||||
storage, err := NewFileStorage(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Setup initial messages
|
||||
if len(tt.initialMsgs) > 0 {
|
||||
err = storage.Save(tt.initialMsgs)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
beforeTime := time.Now()
|
||||
err = storage.Add(tt.newMessage)
|
||||
afterTime := time.Now()
|
||||
|
||||
if tt.expectError {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify message was added to internal state
|
||||
assert.Len(t, storage.msgs, tt.expectedCount)
|
||||
|
||||
// Verify timestamp is reasonable
|
||||
if tt.expectedCount > 0 {
|
||||
lastMsg := storage.msgs[len(storage.msgs)-1]
|
||||
assert.True(t, lastMsg.Time.After(beforeTime) || lastMsg.Time.Equal(beforeTime))
|
||||
assert.True(t, lastMsg.Time.Before(afterTime) || lastMsg.Time.Equal(afterTime))
|
||||
assert.Equal(t, tt.newMessage, lastMsg.Message)
|
||||
}
|
||||
|
||||
_, err := os.ReadFile(storage.path)
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileStorage_Clear(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialMsgs []Message
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "clear empty storage",
|
||||
initialMsgs: []Message{},
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "clear storage with messages",
|
||||
initialMsgs: []Message{
|
||||
{
|
||||
Message: createTestMessage("test1"),
|
||||
Time: time.Now(),
|
||||
},
|
||||
{
|
||||
Message: createTestMessage("test2"),
|
||||
Time: time.Now(),
|
||||
},
|
||||
},
|
||||
expectError: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
tmpDir := createTempDir(t)
|
||||
storage, err := NewFileStorage(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Setup initial messages
|
||||
if len(tt.initialMsgs) > 0 {
|
||||
err = storage.Save(tt.initialMsgs)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
err = storage.Clear()
|
||||
|
||||
if tt.expectError {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify internal state is cleared
|
||||
assert.Empty(t, storage.msgs)
|
||||
|
||||
// Verify file contains empty array
|
||||
data, err := os.ReadFile(storage.path)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "[]", string(data))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileStorage_ConcurrentAccess(t *testing.T) {
|
||||
tmpDir := createTempDir(t)
|
||||
storage, err := NewFileStorage(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Test concurrent Add operations
|
||||
numGoroutines := 10
|
||||
done := make(chan bool, numGoroutines)
|
||||
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func(id int) {
|
||||
defer func() { done <- true }()
|
||||
|
||||
msg := createTestMessage(string(rune('A' + id)))
|
||||
err := storage.Add(msg)
|
||||
assert.NoError(t, err)
|
||||
}(i)
|
||||
}
|
||||
|
||||
// Wait for all goroutines to complete
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
<-done
|
||||
}
|
||||
|
||||
// Verify all messages were added
|
||||
msgs, err := storage.Load()
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, msgs, numGoroutines)
|
||||
}
|
||||
|
||||
func TestFileStorage_IntegrationFlow(t *testing.T) {
|
||||
tmpDir := createTempDir(t)
|
||||
storage, err := NewFileStorage(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Test full workflow
|
||||
|
||||
// 1. Load from empty storage
|
||||
msgs, err := storage.Load()
|
||||
assert.NoError(t, err)
|
||||
assert.Empty(t, msgs)
|
||||
|
||||
// 2. Add some messages
|
||||
msg1 := createTestMessage("message1")
|
||||
err = storage.Add(msg1)
|
||||
assert.NoError(t, err)
|
||||
|
||||
msg2 := createTestMessage("message2")
|
||||
err = storage.Add(msg2)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 3. Load and verify
|
||||
msgs, err = storage.Load()
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, msgs, 2)
|
||||
|
||||
// 4. Save new set of messages
|
||||
newMsgs := []Message{
|
||||
{
|
||||
Message: createTestMessage("new1"),
|
||||
Time: time.Now(),
|
||||
},
|
||||
}
|
||||
err = storage.Save(newMsgs)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 5. Load and verify replacement
|
||||
msgs, err = storage.Load()
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, msgs, 1)
|
||||
|
||||
// 6. Clear storage
|
||||
err = storage.Clear()
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 7. Verify empty
|
||||
msgs, err = storage.Load()
|
||||
assert.NoError(t, err)
|
||||
assert.Empty(t, msgs)
|
||||
}
|
||||
|
||||
func TestFileStorage_FilePermissions(t *testing.T) {
|
||||
tmpDir := createTempDir(t)
|
||||
storage, err := NewFileStorage(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Add a message to create the file
|
||||
msg := createTestMessage("test")
|
||||
err = storage.Add(msg)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Check file permissions
|
||||
info, err := os.Stat(storage.path)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, os.FileMode(0o644), info.Mode().Perm())
|
||||
}
|
||||
|
||||
func TestFileStorage_ErrorHandling(t *testing.T) {
|
||||
tmpDir := createTempDir(t)
|
||||
storage, err := NewFileStorage(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Make directory read-only to trigger write errors
|
||||
err = os.Chmod(tmpDir, 0o555)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Restore permissions for cleanup
|
||||
t.Cleanup(func() {
|
||||
if err := os.Chmod(tmpDir, 0o755); err != nil {
|
||||
t.Errorf("Failed to restore permissions: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
// Try to add a message - should fail due to write permissions
|
||||
msg := createTestMessage("test")
|
||||
err = storage.Add(msg)
|
||||
assert.Error(t, err)
|
||||
|
||||
// Try to save - should fail due to write permissions
|
||||
err = storage.Save([]Message{})
|
||||
assert.Error(t, err)
|
||||
|
||||
// Try to clear - should fail due to write permissions
|
||||
err = storage.Clear()
|
||||
assert.Error(t, err)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,149 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
import "google/protobuf/timestamp.proto";
|
||||
|
||||
package cvms;
|
||||
|
||||
option go_package = "./cvms";
|
||||
|
||||
service Service {
|
||||
rpc Process(stream ClientStreamMessage) returns (stream ServerStreamMessage) {}
|
||||
}
|
||||
|
||||
message AgentStateReq {
|
||||
string id = 1;
|
||||
}
|
||||
|
||||
message AgentStateRes {
|
||||
string id = 1;
|
||||
string state = 2;
|
||||
}
|
||||
|
||||
message StopComputation {
|
||||
string computation_id = 1;
|
||||
}
|
||||
|
||||
message StopComputationResponse {
|
||||
string computation_id = 1;
|
||||
string message = 2;
|
||||
}
|
||||
|
||||
message RunResponse{
|
||||
string computation_id = 1;
|
||||
string error = 2;
|
||||
}
|
||||
|
||||
message AgentEvent {
|
||||
string event_type = 1;
|
||||
google.protobuf.Timestamp timestamp = 2;
|
||||
string computation_id = 3;
|
||||
bytes details = 4;
|
||||
string originator = 5;
|
||||
string status = 6;
|
||||
}
|
||||
|
||||
message AgentLog {
|
||||
string message = 1;
|
||||
string computation_id = 2;
|
||||
string level = 3;
|
||||
google.protobuf.Timestamp timestamp = 4;
|
||||
}
|
||||
|
||||
message ClientStreamMessage {
|
||||
oneof message {
|
||||
AgentLog agent_log = 1;
|
||||
AgentEvent agent_event = 2;
|
||||
RunResponse run_res = 3;
|
||||
StopComputationResponse stopComputationRes = 4;
|
||||
AgentStateRes agentStateRes = 5;
|
||||
AttestationResponse vTPMattestationReport = 6;
|
||||
azureAttestationToken azureAttestationToken = 7;
|
||||
}
|
||||
}
|
||||
|
||||
message ServerStreamMessage {
|
||||
oneof message {
|
||||
RunReqChunks runReqChunks = 1;
|
||||
ComputationRunReq runReq = 2;
|
||||
StopComputation stopComputation = 3;
|
||||
AgentStateReq agentStateReq = 4;
|
||||
DisconnectReq disconnectReq = 5;
|
||||
}
|
||||
}
|
||||
|
||||
message DisconnectReq {
|
||||
string id = 1;
|
||||
}
|
||||
|
||||
message RunReqChunks {
|
||||
bytes data = 1;
|
||||
string id = 2;
|
||||
bool is_last = 3;
|
||||
}
|
||||
|
||||
message ComputationRunReq {
|
||||
string id = 1;
|
||||
string name = 2;
|
||||
string description = 3;
|
||||
repeated Dataset datasets = 4;
|
||||
Algorithm algorithm = 5;
|
||||
repeated ResultConsumer result_consumers = 6;
|
||||
AgentConfig agent_config = 7;
|
||||
}
|
||||
|
||||
message ResultConsumer {
|
||||
bytes userKey = 1;
|
||||
}
|
||||
|
||||
message Dataset {
|
||||
bytes hash = 1; // should be sha3.Sum256, 32 byte length.
|
||||
bytes userKey = 2;
|
||||
string filename = 3;
|
||||
Source source = 4; // Optional remote source for encrypted dataset
|
||||
bool decompress = 5;
|
||||
KBSConfig kbs = 6; // Optional KBS configuration override
|
||||
}
|
||||
|
||||
message Algorithm {
|
||||
bytes hash = 1; // should be sha3.Sum256, 32 byte length.
|
||||
bytes userKey = 2;
|
||||
Source source = 3; // Optional remote source for encrypted algorithm
|
||||
string algo_type = 4;
|
||||
repeated string algo_args = 5;
|
||||
KBSConfig kbs = 6; // Optional KBS configuration override
|
||||
}
|
||||
|
||||
message Source {
|
||||
string type = 1; // Type of source: "oci-image", "s3", "gcs", "https", "http"
|
||||
string url = 2; // URL of the resource (e.g., docker://registry/repo:tag, s3://bucket/key, https://host/path)
|
||||
string kbs_resource_path = 3; // Path to decryption key in KBS (e.g., "default/key/my-key")
|
||||
bool encrypted = 4; // Whether the resource is encrypted (requires KBS)
|
||||
}
|
||||
|
||||
message KBSConfig {
|
||||
string url = 1; // KBS endpoint URL (e.g., "https://kbs.example.com")
|
||||
bool enabled = 2; // Whether to use KBS for key retrieval
|
||||
}
|
||||
|
||||
message AgentConfig {
|
||||
string port = 1;
|
||||
string cert_file = 2;
|
||||
string key_file = 3;
|
||||
string client_ca_file = 4;
|
||||
string server_ca_file = 5;
|
||||
string log_level = 6;
|
||||
bool attested_tls = 7;
|
||||
}
|
||||
|
||||
message AttestationResponse {
|
||||
bytes file = 1;
|
||||
string certSerialNumber = 2;
|
||||
}
|
||||
|
||||
message azureAttestationToken {
|
||||
bytes file = 1;
|
||||
string certSerialNumber = 2;
|
||||
}
|
||||
@@ -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.6.0
|
||||
// - protoc v7.35.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.Error(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 panics, 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,119 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
|
||||
"github.com/ultravioletrs/cocos/agent/auth"
|
||||
"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/credentials/insecure"
|
||||
"google.golang.org/grpc/health"
|
||||
"google.golang.org/grpc/health/grpc_health_v1"
|
||||
"google.golang.org/grpc/reflection"
|
||||
)
|
||||
|
||||
const (
|
||||
svcName = "agent"
|
||||
defSvcGRPCSocket = "/run/cocos/agent.sock"
|
||||
)
|
||||
|
||||
type AgentServer interface {
|
||||
Start(cfg agent.AgentConfig, cmp agent.Computation) error
|
||||
Stop() error
|
||||
}
|
||||
|
||||
type agentServer struct {
|
||||
mu sync.Mutex
|
||||
gs *grpc.Server
|
||||
logger *slog.Logger
|
||||
svc agent.Service
|
||||
host string
|
||||
}
|
||||
|
||||
func NewServer(logger *slog.Logger, svc agent.Service, host string) AgentServer {
|
||||
return &agentServer{
|
||||
logger: logger,
|
||||
svc: svc,
|
||||
host: host,
|
||||
}
|
||||
}
|
||||
|
||||
func (as *agentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
|
||||
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
|
||||
}
|
||||
|
||||
grpcServerOptions := []grpc.ServerOption{
|
||||
grpc.StatsHandler(otelgrpc.NewServerHandler()),
|
||||
}
|
||||
|
||||
// Add authentication interceptors
|
||||
unary, stream := agentgrpc.NewAuthInterceptor(authSvc)
|
||||
grpcServerOptions = append(grpcServerOptions, grpc.UnaryInterceptor(unary))
|
||||
grpcServerOptions = append(grpcServerOptions, grpc.StreamInterceptor(stream))
|
||||
|
||||
// Internal Unix socket is pure plaintext HTTP/2; Ingress Proxy handles external aTLS termination
|
||||
grpcServerOptions = append(grpcServerOptions, grpc.Creds(insecure.NewCredentials()))
|
||||
|
||||
as.mu.Lock()
|
||||
as.gs = grpc.NewServer(grpcServerOptions...)
|
||||
gs := as.gs
|
||||
as.mu.Unlock()
|
||||
|
||||
reflection.Register(gs)
|
||||
agent.RegisterAgentServiceServer(gs, agentgrpc.NewServer(as.svc))
|
||||
|
||||
healthServer := health.NewServer()
|
||||
healthServer.SetServingStatus("agent", grpc_health_v1.HealthCheckResponse_SERVING)
|
||||
grpc_health_v1.RegisterHealthServer(gs, healthServer)
|
||||
|
||||
socketPath := as.host
|
||||
if socketPath == "" || socketPath == "0.0.0.0" {
|
||||
socketPath = defSvcGRPCSocket
|
||||
}
|
||||
|
||||
var listener net.Listener
|
||||
if socketPath[0] == '/' || socketPath[0] == '.' {
|
||||
// Remove existing socket file if it exists
|
||||
_ = os.Remove(socketPath)
|
||||
listener, err = net.Listen("unix", socketPath)
|
||||
} else {
|
||||
listener, err = net.Listen("tcp", socketPath)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
as.logger.Error(fmt.Sprintf("failed to listen on %s: %s", socketPath, err))
|
||||
return err
|
||||
}
|
||||
|
||||
as.logger.Info(fmt.Sprintf("agent service gRPC server listening at %s without TLS", socketPath))
|
||||
|
||||
go func() {
|
||||
err := gs.Serve(listener)
|
||||
if err != nil && err != grpc.ErrServerStopped {
|
||||
as.logger.Error(fmt.Sprintf("failed to start grpc server %s", err.Error()))
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (as *agentServer) Stop() error {
|
||||
as.mu.Lock()
|
||||
defer as.mu.Unlock()
|
||||
if as.gs != nil {
|
||||
as.gs.GracefulStop()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -1,518 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/x509"
|
||||
"log/slog"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/agent/mocks"
|
||||
)
|
||||
|
||||
func setupTest(t *testing.T) (*slog.Logger, *mocks.Service, string, []byte) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError}))
|
||||
mockSvc := new(mocks.Service)
|
||||
host := "localhost:0"
|
||||
|
||||
privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
assert.NoError(t, err, "Failed to generate ECDSA key")
|
||||
|
||||
pubkey, err := x509.MarshalPKIXPublicKey(privateKey.Public())
|
||||
assert.NoError(t, err, "Failed to marshal public key")
|
||||
|
||||
return logger, mockSvc, host, pubkey
|
||||
}
|
||||
|
||||
func TestNewServer(t *testing.T) {
|
||||
logger, svc, host, _ := setupTest(t)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
logger *slog.Logger
|
||||
svc agent.Service
|
||||
host string
|
||||
expected AgentServer
|
||||
}{
|
||||
{
|
||||
name: "valid server creation",
|
||||
logger: logger,
|
||||
svc: svc,
|
||||
host: host,
|
||||
},
|
||||
{
|
||||
name: "server with empty host",
|
||||
logger: logger,
|
||||
svc: svc,
|
||||
host: "",
|
||||
},
|
||||
{
|
||||
name: "server with empty caUrl",
|
||||
logger: logger,
|
||||
svc: svc,
|
||||
host: host,
|
||||
},
|
||||
{
|
||||
name: "server with empty cvmId",
|
||||
logger: logger,
|
||||
svc: svc,
|
||||
host: host,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
server := NewServer(tt.logger, tt.svc, tt.host)
|
||||
|
||||
assert.NotNil(t, server)
|
||||
|
||||
agentSrv, ok := server.(*agentServer)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, tt.logger, agentSrv.logger)
|
||||
assert.Equal(t, tt.svc, agentSrv.svc)
|
||||
assert.Equal(t, tt.host, agentSrv.host)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentServer_Start(t *testing.T) {
|
||||
logger, svc, host, pubKey := setupTest(t)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
cfg agent.AgentConfig
|
||||
cmp agent.Computation
|
||||
setupMocks func(*mocks.Service)
|
||||
expectedError bool
|
||||
errorContains string
|
||||
}{
|
||||
{
|
||||
name: "successful start with default port",
|
||||
cfg: agent.AgentConfig{
|
||||
CertFile: "cert.pem",
|
||||
KeyFile: "key.pem",
|
||||
ServerCAFile: "server-ca.pem",
|
||||
ClientCAFile: "client-ca.pem",
|
||||
AttestedTls: true,
|
||||
},
|
||||
cmp: agent.Computation{
|
||||
ID: "test-computation-1",
|
||||
Name: "Test Computation",
|
||||
Description: "A test computation",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x01, 0x02, 0x03},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x04, 0x05, 0x06},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
},
|
||||
setupMocks: func(m *mocks.Service) {
|
||||
},
|
||||
expectedError: false,
|
||||
},
|
||||
{
|
||||
name: "successful start with custom port",
|
||||
cfg: agent.AgentConfig{
|
||||
CertFile: "cert.pem",
|
||||
KeyFile: "key.pem",
|
||||
ServerCAFile: "server-ca.pem",
|
||||
ClientCAFile: "client-ca.pem",
|
||||
AttestedTls: false,
|
||||
},
|
||||
cmp: agent.Computation{
|
||||
ID: "test-computation-2",
|
||||
Name: "Test Computation 2",
|
||||
Description: "Another test computation",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x07, 0x08, 0x09},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x0a, 0x0b, 0x0c},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
},
|
||||
setupMocks: func(m *mocks.Service) {
|
||||
},
|
||||
expectedError: false,
|
||||
},
|
||||
{
|
||||
name: "start with minimal config",
|
||||
cfg: agent.AgentConfig{
|
||||
AttestedTls: false,
|
||||
},
|
||||
cmp: agent.Computation{
|
||||
ID: "test-computation-3",
|
||||
Name: "Minimal Test",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x0d, 0x0e, 0x0f},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x10, 0x11, 0x12},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
},
|
||||
setupMocks: func(m *mocks.Service) {
|
||||
},
|
||||
expectedError: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
tt.setupMocks(svc)
|
||||
|
||||
server := NewServer(logger, svc, host)
|
||||
|
||||
err := server.Start(tt.cfg, tt.cmp)
|
||||
|
||||
if tt.expectedError {
|
||||
assert.Error(t, err)
|
||||
if tt.errorContains != "" {
|
||||
assert.Contains(t, err.Error(), tt.errorContains)
|
||||
}
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify the port was set correctly
|
||||
agentSrv := server.(*agentServer)
|
||||
assert.NotNil(t, agentSrv.gs)
|
||||
|
||||
if err := server.Stop(); err != nil {
|
||||
t.Fatalf("Failed to stop server after start: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
svc.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentServer_Stop(t *testing.T) {
|
||||
logger, svc, host, pubKey := setupTest(t)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
setupServer func(AgentServer) error
|
||||
expectedError bool
|
||||
errorContains string
|
||||
}{
|
||||
{
|
||||
name: "stop unstarted server",
|
||||
setupServer: func(server AgentServer) error {
|
||||
// Don't start the server
|
||||
return nil
|
||||
},
|
||||
expectedError: false,
|
||||
},
|
||||
{
|
||||
name: "stop started server",
|
||||
setupServer: func(server AgentServer) error {
|
||||
cfg := agent.AgentConfig{}
|
||||
cmp := agent.Computation{
|
||||
ID: "test-stop-computation",
|
||||
Name: "Stop Test",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x19, 0x1a, 0x1b},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x1c, 0x1d, 0x1e},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
}
|
||||
return server.Start(cfg, cmp)
|
||||
},
|
||||
expectedError: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
server := NewServer(logger, svc, host)
|
||||
|
||||
err := tt.setupServer(server)
|
||||
if err != nil {
|
||||
t.Fatalf("Setup failed: %v", err)
|
||||
}
|
||||
|
||||
// Give the server a moment to start if it was started
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
err = server.Stop()
|
||||
|
||||
if tt.expectedError {
|
||||
assert.Error(t, err)
|
||||
if tt.errorContains != "" {
|
||||
assert.Contains(t, err.Error(), tt.errorContains)
|
||||
}
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
svc.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentServer_StopMultipleTimes(t *testing.T) {
|
||||
logger, svc, host, pubKey := setupTest(t)
|
||||
server := NewServer(logger, svc, host)
|
||||
|
||||
// Start the server
|
||||
cfg := agent.AgentConfig{}
|
||||
cmp := agent.Computation{
|
||||
ID: "test-multiple-stop",
|
||||
Name: "Multiple Stop Test",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x1f, 0x20, 0x21},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x22, 0x23, 0x24},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
err := server.Start(cfg, cmp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Give the server a moment to start
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
// Stop the server multiple times
|
||||
err1 := server.Stop()
|
||||
err2 := server.Stop()
|
||||
err3 := server.Stop()
|
||||
|
||||
assert.NoError(t, err1)
|
||||
assert.NoError(t, err2)
|
||||
assert.NoError(t, err3)
|
||||
|
||||
svc.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestAgentServer_StartAfterStop(t *testing.T) {
|
||||
logger, svc, host, pubKey := setupTest(t)
|
||||
server := NewServer(logger, svc, host)
|
||||
|
||||
cfg := agent.AgentConfig{}
|
||||
cmp := agent.Computation{
|
||||
ID: "test-restart",
|
||||
Name: "Restart Test",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x25, 0x26, 0x27},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x28, 0x29, 0x2a},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Start, stop, then start again
|
||||
err := server.Start(cfg, cmp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
err = server.Stop()
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Start again with different config
|
||||
cfg2 := agent.AgentConfig{}
|
||||
cmp2 := agent.Computation{
|
||||
ID: "test-restart-2",
|
||||
Name: "Restart Test 2",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x2b, 0x2c, 0x2d},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x2e, 0x2f, 0x30},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
err = server.Start(cfg2, cmp2)
|
||||
assert.NoError(t, err)
|
||||
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
err = server.Stop()
|
||||
assert.NoError(t, err)
|
||||
|
||||
svc.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestAgentServer_ConfigValidation(t *testing.T) {
|
||||
logger, svc, host, pubKey := setupTest(t)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
config agent.AgentConfig
|
||||
cmp agent.Computation
|
||||
valid bool
|
||||
}{
|
||||
{
|
||||
name: "valid config with all fields",
|
||||
config: agent.AgentConfig{
|
||||
CertFile: "cert.pem",
|
||||
KeyFile: "key.pem",
|
||||
ServerCAFile: "server-ca.pem",
|
||||
ClientCAFile: "client-ca.pem",
|
||||
AttestedTls: true,
|
||||
},
|
||||
cmp: agent.Computation{
|
||||
ID: "valid-config-test",
|
||||
Name: "Valid Config Test",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x31, 0x32, 0x33},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x34, 0x35, 0x36},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
},
|
||||
valid: true,
|
||||
},
|
||||
{
|
||||
name: "valid config with minimal fields",
|
||||
config: agent.AgentConfig{},
|
||||
cmp: agent.Computation{
|
||||
ID: "minimal-config-test",
|
||||
Name: "Minimal Config Test",
|
||||
Algorithm: &agent.Algorithm{
|
||||
Hash: [32]byte{0x37, 0x38, 0x39},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
Datasets: []agent.Dataset{
|
||||
{
|
||||
Hash: [32]byte{0x3a, 0x3b, 0x3c},
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{
|
||||
UserKey: pubKey,
|
||||
},
|
||||
},
|
||||
},
|
||||
valid: true,
|
||||
},
|
||||
{
|
||||
name: "config with empty port uses default",
|
||||
config: agent.AgentConfig{},
|
||||
cmp: agent.Computation{
|
||||
ID: "default-port-test",
|
||||
Name: "Default Port Test",
|
||||
Algorithm: &agent.Algorithm{Hash: [32]byte{0x3d, 0x3e, 0x3f}, UserKey: pubKey},
|
||||
Datasets: []agent.Dataset{
|
||||
{Hash: [32]byte{0x40, 0x41, 0x42}, UserKey: pubKey},
|
||||
},
|
||||
ResultConsumers: []agent.ResultConsumer{
|
||||
{UserKey: pubKey},
|
||||
},
|
||||
},
|
||||
valid: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
server := NewServer(logger, svc, host)
|
||||
|
||||
err := server.Start(tt.config, tt.cmp)
|
||||
|
||||
if tt.valid {
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify server started successfully
|
||||
agentSrv := server.(*agentServer)
|
||||
assert.NotNil(t, agentSrv.gs)
|
||||
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
if err := server.Stop(); err != nil {
|
||||
t.Fatalf("Failed to stop server after start: %v", err)
|
||||
}
|
||||
} else {
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
svc.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConstants(t *testing.T) {
|
||||
assert.Equal(t, "agent", svcName)
|
||||
assert.Equal(t, "/run/cocos/agent.sock", defSvcGRPCSocket)
|
||||
}
|
||||
@@ -1,141 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
)
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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 for the type AgentServer
|
||||
func (_mock *AgentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
|
||||
ret := _mock.Called(cfg, cmp)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Start")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(agent.AgentConfig, agent.Computation) error); ok {
|
||||
r0 = returnFunc(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) {
|
||||
var arg0 agent.AgentConfig
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(agent.AgentConfig)
|
||||
}
|
||||
var arg1 agent.Computation
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].(agent.Computation)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentServer_Start_Call) Return(err error) *AgentServer_Start_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentServer_Start_Call) RunAndReturn(run func(cfg agent.AgentConfig, cmp agent.Computation) error) *AgentServer_Start_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Stop provides a mock function for the type AgentServer
|
||||
func (_mock *AgentServer) Stop() error {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Stop")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func() error); ok {
|
||||
r0 = returnFunc()
|
||||
} 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(err error) *AgentServer_Stop_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentServer_Stop_Call) RunAndReturn(run func() error) *AgentServer_Stop_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
+46
-19
@@ -4,38 +4,65 @@ package events
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"github.com/mdlayher/vsock"
|
||||
"github.com/ultravioletrs/cocos/pkg/manager"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
)
|
||||
|
||||
type service struct {
|
||||
service string
|
||||
queue chan *cvms.ClientStreamMessage
|
||||
service string
|
||||
computationID string
|
||||
conn *vsock.Conn
|
||||
}
|
||||
|
||||
type AgentEvent struct {
|
||||
EventType string `json:"event_type"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
ComputationID string `json:"computation_id,omitempty"`
|
||||
Details json.RawMessage `json:"details,omitempty"`
|
||||
Originator string `json:"originator"`
|
||||
Status string `json:"status,omitempty"`
|
||||
}
|
||||
|
||||
type Service interface {
|
||||
SendEvent(cmpID, event, status string, details json.RawMessage)
|
||||
SendEvent(event, status string, details json.RawMessage) error
|
||||
Close() error
|
||||
}
|
||||
|
||||
func New(svc string, queue chan *cvms.ClientStreamMessage) (Service, error) {
|
||||
func New(svc, computationID string, sockPort uint32) (Service, error) {
|
||||
conn, err := vsock.Dial(vsock.Host, sockPort, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &service{
|
||||
service: svc,
|
||||
queue: queue,
|
||||
service: svc,
|
||||
computationID: computationID,
|
||||
conn: conn,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *service) SendEvent(cmpID, event, status string, details json.RawMessage) {
|
||||
s.queue <- &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_AgentEvent{
|
||||
AgentEvent: &cvms.AgentEvent{
|
||||
EventType: event,
|
||||
Timestamp: timestamppb.Now(),
|
||||
ComputationId: cmpID,
|
||||
Originator: s.service,
|
||||
Status: status,
|
||||
Details: details,
|
||||
},
|
||||
},
|
||||
func (s *service) SendEvent(event, status string, details json.RawMessage) error {
|
||||
body := manager.ClientStreamMessage{Message: &manager.ClientStreamMessage_AgentEvent{AgentEvent: &manager.AgentEvent{
|
||||
EventType: event,
|
||||
Timestamp: timestamppb.Now(),
|
||||
ComputationId: s.computationID,
|
||||
Originator: s.service,
|
||||
Status: status,
|
||||
Details: details,
|
||||
}}}
|
||||
protoBody, err := proto.Marshal(&body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := s.conn.Write(protoBody); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *service) Close() error {
|
||||
return s.conn.Close()
|
||||
}
|
||||
|
||||
@@ -1,348 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.11
|
||||
// protoc v7.35.0
|
||||
// source: agent/events/events.proto
|
||||
|
||||
package events
|
||||
|
||||
import (
|
||||
protoreflect "google.golang.org/protobuf/reflect/protoreflect"
|
||||
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
|
||||
timestamppb "google.golang.org/protobuf/types/known/timestamppb"
|
||||
reflect "reflect"
|
||||
sync "sync"
|
||||
unsafe "unsafe"
|
||||
)
|
||||
|
||||
const (
|
||||
// Verify that this generated code is sufficiently up-to-date.
|
||||
_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
|
||||
// Verify that runtime/protoimpl is sufficiently up-to-date.
|
||||
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
|
||||
)
|
||||
|
||||
type AgentEvent struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
EventType string `protobuf:"bytes,1,opt,name=event_type,json=eventType,proto3" json:"event_type,omitempty"`
|
||||
Timestamp *timestamppb.Timestamp `protobuf:"bytes,2,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
|
||||
ComputationId string `protobuf:"bytes,3,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
|
||||
Details []byte `protobuf:"bytes,4,opt,name=details,proto3" json:"details,omitempty"`
|
||||
Originator string `protobuf:"bytes,5,opt,name=originator,proto3" json:"originator,omitempty"`
|
||||
Status string `protobuf:"bytes,6,opt,name=status,proto3" json:"status,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *AgentEvent) Reset() {
|
||||
*x = AgentEvent{}
|
||||
mi := &file_agent_events_events_proto_msgTypes[0]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *AgentEvent) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*AgentEvent) ProtoMessage() {}
|
||||
|
||||
func (x *AgentEvent) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_events_events_proto_msgTypes[0]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use AgentEvent.ProtoReflect.Descriptor instead.
|
||||
func (*AgentEvent) Descriptor() ([]byte, []int) {
|
||||
return file_agent_events_events_proto_rawDescGZIP(), []int{0}
|
||||
}
|
||||
|
||||
func (x *AgentEvent) GetEventType() string {
|
||||
if x != nil {
|
||||
return x.EventType
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *AgentEvent) GetTimestamp() *timestamppb.Timestamp {
|
||||
if x != nil {
|
||||
return x.Timestamp
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *AgentEvent) GetComputationId() string {
|
||||
if x != nil {
|
||||
return x.ComputationId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *AgentEvent) GetDetails() []byte {
|
||||
if x != nil {
|
||||
return x.Details
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *AgentEvent) GetOriginator() string {
|
||||
if x != nil {
|
||||
return x.Originator
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *AgentEvent) GetStatus() string {
|
||||
if x != nil {
|
||||
return x.Status
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type AgentLog struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Message string `protobuf:"bytes,1,opt,name=message,proto3" json:"message,omitempty"`
|
||||
ComputationId string `protobuf:"bytes,2,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
|
||||
Level string `protobuf:"bytes,3,opt,name=level,proto3" json:"level,omitempty"`
|
||||
Timestamp *timestamppb.Timestamp `protobuf:"bytes,4,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *AgentLog) Reset() {
|
||||
*x = AgentLog{}
|
||||
mi := &file_agent_events_events_proto_msgTypes[1]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *AgentLog) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*AgentLog) ProtoMessage() {}
|
||||
|
||||
func (x *AgentLog) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_events_events_proto_msgTypes[1]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use AgentLog.ProtoReflect.Descriptor instead.
|
||||
func (*AgentLog) Descriptor() ([]byte, []int) {
|
||||
return file_agent_events_events_proto_rawDescGZIP(), []int{1}
|
||||
}
|
||||
|
||||
func (x *AgentLog) GetMessage() string {
|
||||
if x != nil {
|
||||
return x.Message
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *AgentLog) GetComputationId() string {
|
||||
if x != nil {
|
||||
return x.ComputationId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *AgentLog) GetLevel() string {
|
||||
if x != nil {
|
||||
return x.Level
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *AgentLog) GetTimestamp() *timestamppb.Timestamp {
|
||||
if x != nil {
|
||||
return x.Timestamp
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type EventsLogs struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
// Types that are valid to be assigned to Message:
|
||||
//
|
||||
// *EventsLogs_AgentLog
|
||||
// *EventsLogs_AgentEvent
|
||||
Message isEventsLogs_Message `protobuf_oneof:"message"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *EventsLogs) Reset() {
|
||||
*x = EventsLogs{}
|
||||
mi := &file_agent_events_events_proto_msgTypes[2]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *EventsLogs) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*EventsLogs) ProtoMessage() {}
|
||||
|
||||
func (x *EventsLogs) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_events_events_proto_msgTypes[2]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use EventsLogs.ProtoReflect.Descriptor instead.
|
||||
func (*EventsLogs) Descriptor() ([]byte, []int) {
|
||||
return file_agent_events_events_proto_rawDescGZIP(), []int{2}
|
||||
}
|
||||
|
||||
func (x *EventsLogs) GetMessage() isEventsLogs_Message {
|
||||
if x != nil {
|
||||
return x.Message
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *EventsLogs) GetAgentLog() *AgentLog {
|
||||
if x != nil {
|
||||
if x, ok := x.Message.(*EventsLogs_AgentLog); ok {
|
||||
return x.AgentLog
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *EventsLogs) GetAgentEvent() *AgentEvent {
|
||||
if x != nil {
|
||||
if x, ok := x.Message.(*EventsLogs_AgentEvent); ok {
|
||||
return x.AgentEvent
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type isEventsLogs_Message interface {
|
||||
isEventsLogs_Message()
|
||||
}
|
||||
|
||||
type EventsLogs_AgentLog struct {
|
||||
AgentLog *AgentLog `protobuf:"bytes,1,opt,name=agent_log,json=agentLog,proto3,oneof"`
|
||||
}
|
||||
|
||||
type EventsLogs_AgentEvent struct {
|
||||
AgentEvent *AgentEvent `protobuf:"bytes,2,opt,name=agent_event,json=agentEvent,proto3,oneof"`
|
||||
}
|
||||
|
||||
func (*EventsLogs_AgentLog) isEventsLogs_Message() {}
|
||||
|
||||
func (*EventsLogs_AgentEvent) isEventsLogs_Message() {}
|
||||
|
||||
var File_agent_events_events_proto protoreflect.FileDescriptor
|
||||
|
||||
const file_agent_events_events_proto_rawDesc = "" +
|
||||
"\n" +
|
||||
"\x19agent/events/events.proto\x12\x06events\x1a\x1fgoogle/protobuf/timestamp.proto\"\xde\x01\n" +
|
||||
"\n" +
|
||||
"AgentEvent\x12\x1d\n" +
|
||||
"\n" +
|
||||
"event_type\x18\x01 \x01(\tR\teventType\x128\n" +
|
||||
"\ttimestamp\x18\x02 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12%\n" +
|
||||
"\x0ecomputation_id\x18\x03 \x01(\tR\rcomputationId\x12\x18\n" +
|
||||
"\adetails\x18\x04 \x01(\fR\adetails\x12\x1e\n" +
|
||||
"\n" +
|
||||
"originator\x18\x05 \x01(\tR\n" +
|
||||
"originator\x12\x16\n" +
|
||||
"\x06status\x18\x06 \x01(\tR\x06status\"\x9b\x01\n" +
|
||||
"\bAgentLog\x12\x18\n" +
|
||||
"\amessage\x18\x01 \x01(\tR\amessage\x12%\n" +
|
||||
"\x0ecomputation_id\x18\x02 \x01(\tR\rcomputationId\x12\x14\n" +
|
||||
"\x05level\x18\x03 \x01(\tR\x05level\x128\n" +
|
||||
"\ttimestamp\x18\x04 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\"\x7f\n" +
|
||||
"\n" +
|
||||
"EventsLogs\x12/\n" +
|
||||
"\tagent_log\x18\x01 \x01(\v2\x10.events.AgentLogH\x00R\bagentLog\x125\n" +
|
||||
"\vagent_event\x18\x02 \x01(\v2\x12.events.AgentEventH\x00R\n" +
|
||||
"agentEventB\t\n" +
|
||||
"\amessageB\n" +
|
||||
"Z\b./eventsb\x06proto3"
|
||||
|
||||
var (
|
||||
file_agent_events_events_proto_rawDescOnce sync.Once
|
||||
file_agent_events_events_proto_rawDescData []byte
|
||||
)
|
||||
|
||||
func file_agent_events_events_proto_rawDescGZIP() []byte {
|
||||
file_agent_events_events_proto_rawDescOnce.Do(func() {
|
||||
file_agent_events_events_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_events_events_proto_rawDesc), len(file_agent_events_events_proto_rawDesc)))
|
||||
})
|
||||
return file_agent_events_events_proto_rawDescData
|
||||
}
|
||||
|
||||
var file_agent_events_events_proto_msgTypes = make([]protoimpl.MessageInfo, 3)
|
||||
var file_agent_events_events_proto_goTypes = []any{
|
||||
(*AgentEvent)(nil), // 0: events.AgentEvent
|
||||
(*AgentLog)(nil), // 1: events.AgentLog
|
||||
(*EventsLogs)(nil), // 2: events.EventsLogs
|
||||
(*timestamppb.Timestamp)(nil), // 3: google.protobuf.Timestamp
|
||||
}
|
||||
var file_agent_events_events_proto_depIdxs = []int32{
|
||||
3, // 0: events.AgentEvent.timestamp:type_name -> google.protobuf.Timestamp
|
||||
3, // 1: events.AgentLog.timestamp:type_name -> google.protobuf.Timestamp
|
||||
1, // 2: events.EventsLogs.agent_log:type_name -> events.AgentLog
|
||||
0, // 3: events.EventsLogs.agent_event:type_name -> events.AgentEvent
|
||||
4, // [4:4] is the sub-list for method output_type
|
||||
4, // [4:4] is the sub-list for method input_type
|
||||
4, // [4:4] is the sub-list for extension type_name
|
||||
4, // [4:4] is the sub-list for extension extendee
|
||||
0, // [0:4] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_agent_events_events_proto_init() }
|
||||
func file_agent_events_events_proto_init() {
|
||||
if File_agent_events_events_proto != nil {
|
||||
return
|
||||
}
|
||||
file_agent_events_events_proto_msgTypes[2].OneofWrappers = []any{
|
||||
(*EventsLogs_AgentLog)(nil),
|
||||
(*EventsLogs_AgentEvent)(nil),
|
||||
}
|
||||
type x struct{}
|
||||
out := protoimpl.TypeBuilder{
|
||||
File: protoimpl.DescBuilder{
|
||||
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
|
||||
RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_events_events_proto_rawDesc), len(file_agent_events_events_proto_rawDesc)),
|
||||
NumEnums: 0,
|
||||
NumMessages: 3,
|
||||
NumExtensions: 0,
|
||||
NumServices: 0,
|
||||
},
|
||||
GoTypes: file_agent_events_events_proto_goTypes,
|
||||
DependencyIndexes: file_agent_events_events_proto_depIdxs,
|
||||
MessageInfos: file_agent_events_events_proto_msgTypes,
|
||||
}.Build()
|
||||
File_agent_events_events_proto = out.File
|
||||
file_agent_events_events_proto_goTypes = nil
|
||||
file_agent_events_events_proto_depIdxs = nil
|
||||
}
|
||||
@@ -1,33 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package events;
|
||||
|
||||
import "google/protobuf/timestamp.proto";
|
||||
|
||||
option go_package = "./events";
|
||||
|
||||
message AgentEvent {
|
||||
string event_type = 1;
|
||||
google.protobuf.Timestamp timestamp = 2;
|
||||
string computation_id = 3;
|
||||
bytes details = 4;
|
||||
string originator = 5;
|
||||
string status = 6;
|
||||
}
|
||||
|
||||
message AgentLog {
|
||||
string message = 1;
|
||||
string computation_id = 2;
|
||||
string level = 3;
|
||||
google.protobuf.Timestamp timestamp = 4;
|
||||
}
|
||||
|
||||
message EventsLogs {
|
||||
oneof message {
|
||||
AgentLog agent_log = 1;
|
||||
AgentEvent agent_event = 2;
|
||||
}
|
||||
}
|
||||
@@ -1,38 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package events
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
)
|
||||
|
||||
func TestSendEventSuccess(t *testing.T) {
|
||||
queue := make(chan *cvms.ClientStreamMessage, 1)
|
||||
svc, err := New("test_service", queue)
|
||||
assert.NoError(t, err)
|
||||
|
||||
details := json.RawMessage(`{"key": "value"}`)
|
||||
|
||||
go func() {
|
||||
msg := <-queue
|
||||
assert.NotNil(t, msg)
|
||||
assert.NotNil(t, msg.GetAgentEvent())
|
||||
assert.Equal(t, "test_event", msg.GetAgentEvent().EventType)
|
||||
assert.Equal(t, "testid", msg.GetAgentEvent().ComputationId)
|
||||
assert.Equal(t, "test_service", msg.GetAgentEvent().Originator)
|
||||
assert.Equal(t, "success", msg.GetAgentEvent().Status)
|
||||
|
||||
now := time.Now()
|
||||
eventTimestamp := msg.GetAgentEvent().GetTimestamp().AsTime()
|
||||
assert.WithinDuration(t, now, eventTimestamp, 1*time.Second)
|
||||
}()
|
||||
|
||||
svc.SendEvent("testid", "test_event", "success", details)
|
||||
|
||||
time.Sleep(1 * time.Second)
|
||||
}
|
||||
@@ -1,99 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
|
||||
// The first argument is typically a *testing.T value.
|
||||
func NewService(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *Service {
|
||||
mock := &Service{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
|
||||
// Service is an autogenerated mock type for the Service type
|
||||
type Service struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type Service_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *Service) EXPECT() *Service_Expecter {
|
||||
return &Service_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// SendEvent provides a mock function for the type Service
|
||||
func (_mock *Service) SendEvent(cmpID string, event string, status string, details json.RawMessage) {
|
||||
_mock.Called(cmpID, event, status, details)
|
||||
return
|
||||
}
|
||||
|
||||
// Service_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
|
||||
type Service_SendEvent_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// SendEvent is a helper method to define mock.On call
|
||||
// - cmpID string
|
||||
// - event string
|
||||
// - status string
|
||||
// - details json.RawMessage
|
||||
func (_e *Service_Expecter) SendEvent(cmpID interface{}, event interface{}, status interface{}, details interface{}) *Service_SendEvent_Call {
|
||||
return &Service_SendEvent_Call{Call: _e.mock.On("SendEvent", cmpID, event, status, details)}
|
||||
}
|
||||
|
||||
func (_c *Service_SendEvent_Call) Run(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 string
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(string)
|
||||
}
|
||||
var arg1 string
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].(string)
|
||||
}
|
||||
var arg2 string
|
||||
if args[2] != nil {
|
||||
arg2 = args[2].(string)
|
||||
}
|
||||
var arg3 json.RawMessage
|
||||
if args[3] != nil {
|
||||
arg3 = args[3].(json.RawMessage)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
arg2,
|
||||
arg3,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_SendEvent_Call) Return() *Service_SendEvent_Call {
|
||||
_c.Call.Return()
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_SendEvent_Call) RunAndReturn(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
|
||||
_c.Run(run)
|
||||
return _c
|
||||
}
|
||||
@@ -1,261 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.11
|
||||
// protoc v7.35.0
|
||||
// source: agent/log/log.proto
|
||||
|
||||
package log
|
||||
|
||||
import (
|
||||
protoreflect "google.golang.org/protobuf/reflect/protoreflect"
|
||||
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
|
||||
emptypb "google.golang.org/protobuf/types/known/emptypb"
|
||||
timestamppb "google.golang.org/protobuf/types/known/timestamppb"
|
||||
reflect "reflect"
|
||||
sync "sync"
|
||||
unsafe "unsafe"
|
||||
)
|
||||
|
||||
const (
|
||||
// Verify that this generated code is sufficiently up-to-date.
|
||||
_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
|
||||
// Verify that runtime/protoimpl is sufficiently up-to-date.
|
||||
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
|
||||
)
|
||||
|
||||
type LogEntry struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Message string `protobuf:"bytes,1,opt,name=message,proto3" json:"message,omitempty"`
|
||||
ComputationId string `protobuf:"bytes,2,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
|
||||
Level string `protobuf:"bytes,3,opt,name=level,proto3" json:"level,omitempty"`
|
||||
Timestamp *timestamppb.Timestamp `protobuf:"bytes,4,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *LogEntry) Reset() {
|
||||
*x = LogEntry{}
|
||||
mi := &file_agent_log_log_proto_msgTypes[0]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *LogEntry) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*LogEntry) ProtoMessage() {}
|
||||
|
||||
func (x *LogEntry) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_log_log_proto_msgTypes[0]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use LogEntry.ProtoReflect.Descriptor instead.
|
||||
func (*LogEntry) Descriptor() ([]byte, []int) {
|
||||
return file_agent_log_log_proto_rawDescGZIP(), []int{0}
|
||||
}
|
||||
|
||||
func (x *LogEntry) GetMessage() string {
|
||||
if x != nil {
|
||||
return x.Message
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *LogEntry) GetComputationId() string {
|
||||
if x != nil {
|
||||
return x.ComputationId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *LogEntry) GetLevel() string {
|
||||
if x != nil {
|
||||
return x.Level
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *LogEntry) GetTimestamp() *timestamppb.Timestamp {
|
||||
if x != nil {
|
||||
return x.Timestamp
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type EventEntry struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
EventType string `protobuf:"bytes,1,opt,name=event_type,json=eventType,proto3" json:"event_type,omitempty"`
|
||||
Timestamp *timestamppb.Timestamp `protobuf:"bytes,2,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
|
||||
ComputationId string `protobuf:"bytes,3,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
|
||||
Details []byte `protobuf:"bytes,4,opt,name=details,proto3" json:"details,omitempty"` // JSON payload
|
||||
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 *EventEntry) Reset() {
|
||||
*x = EventEntry{}
|
||||
mi := &file_agent_log_log_proto_msgTypes[1]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *EventEntry) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*EventEntry) ProtoMessage() {}
|
||||
|
||||
func (x *EventEntry) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_log_log_proto_msgTypes[1]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use EventEntry.ProtoReflect.Descriptor instead.
|
||||
func (*EventEntry) Descriptor() ([]byte, []int) {
|
||||
return file_agent_log_log_proto_rawDescGZIP(), []int{1}
|
||||
}
|
||||
|
||||
func (x *EventEntry) GetEventType() string {
|
||||
if x != nil {
|
||||
return x.EventType
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *EventEntry) GetTimestamp() *timestamppb.Timestamp {
|
||||
if x != nil {
|
||||
return x.Timestamp
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *EventEntry) GetComputationId() string {
|
||||
if x != nil {
|
||||
return x.ComputationId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *EventEntry) GetDetails() []byte {
|
||||
if x != nil {
|
||||
return x.Details
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *EventEntry) GetOriginator() string {
|
||||
if x != nil {
|
||||
return x.Originator
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *EventEntry) GetStatus() string {
|
||||
if x != nil {
|
||||
return x.Status
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
var File_agent_log_log_proto protoreflect.FileDescriptor
|
||||
|
||||
const file_agent_log_log_proto_rawDesc = "" +
|
||||
"\n" +
|
||||
"\x13agent/log/log.proto\x12\x03log\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1bgoogle/protobuf/empty.proto\"\x9b\x01\n" +
|
||||
"\bLogEntry\x12\x18\n" +
|
||||
"\amessage\x18\x01 \x01(\tR\amessage\x12%\n" +
|
||||
"\x0ecomputation_id\x18\x02 \x01(\tR\rcomputationId\x12\x14\n" +
|
||||
"\x05level\x18\x03 \x01(\tR\x05level\x128\n" +
|
||||
"\ttimestamp\x18\x04 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\"\xde\x01\n" +
|
||||
"\n" +
|
||||
"EventEntry\x12\x1d\n" +
|
||||
"\n" +
|
||||
"event_type\x18\x01 \x01(\tR\teventType\x128\n" +
|
||||
"\ttimestamp\x18\x02 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12%\n" +
|
||||
"\x0ecomputation_id\x18\x03 \x01(\tR\rcomputationId\x12\x18\n" +
|
||||
"\adetails\x18\x04 \x01(\fR\adetails\x12\x1e\n" +
|
||||
"\n" +
|
||||
"originator\x18\x05 \x01(\tR\n" +
|
||||
"originator\x12\x16\n" +
|
||||
"\x06status\x18\x06 \x01(\tR\x06status2v\n" +
|
||||
"\fLogCollector\x120\n" +
|
||||
"\aSendLog\x12\r.log.LogEntry\x1a\x16.google.protobuf.Empty\x124\n" +
|
||||
"\tSendEvent\x12\x0f.log.EventEntry\x1a\x16.google.protobuf.EmptyB\aZ\x05./logb\x06proto3"
|
||||
|
||||
var (
|
||||
file_agent_log_log_proto_rawDescOnce sync.Once
|
||||
file_agent_log_log_proto_rawDescData []byte
|
||||
)
|
||||
|
||||
func file_agent_log_log_proto_rawDescGZIP() []byte {
|
||||
file_agent_log_log_proto_rawDescOnce.Do(func() {
|
||||
file_agent_log_log_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_log_log_proto_rawDesc), len(file_agent_log_log_proto_rawDesc)))
|
||||
})
|
||||
return file_agent_log_log_proto_rawDescData
|
||||
}
|
||||
|
||||
var file_agent_log_log_proto_msgTypes = make([]protoimpl.MessageInfo, 2)
|
||||
var file_agent_log_log_proto_goTypes = []any{
|
||||
(*LogEntry)(nil), // 0: log.LogEntry
|
||||
(*EventEntry)(nil), // 1: log.EventEntry
|
||||
(*timestamppb.Timestamp)(nil), // 2: google.protobuf.Timestamp
|
||||
(*emptypb.Empty)(nil), // 3: google.protobuf.Empty
|
||||
}
|
||||
var file_agent_log_log_proto_depIdxs = []int32{
|
||||
2, // 0: log.LogEntry.timestamp:type_name -> google.protobuf.Timestamp
|
||||
2, // 1: log.EventEntry.timestamp:type_name -> google.protobuf.Timestamp
|
||||
0, // 2: log.LogCollector.SendLog:input_type -> log.LogEntry
|
||||
1, // 3: log.LogCollector.SendEvent:input_type -> log.EventEntry
|
||||
3, // 4: log.LogCollector.SendLog:output_type -> google.protobuf.Empty
|
||||
3, // 5: log.LogCollector.SendEvent:output_type -> google.protobuf.Empty
|
||||
4, // [4:6] is the sub-list for method output_type
|
||||
2, // [2:4] is the sub-list for method input_type
|
||||
2, // [2:2] is the sub-list for extension type_name
|
||||
2, // [2:2] is the sub-list for extension extendee
|
||||
0, // [0:2] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_agent_log_log_proto_init() }
|
||||
func file_agent_log_log_proto_init() {
|
||||
if File_agent_log_log_proto != nil {
|
||||
return
|
||||
}
|
||||
type x struct{}
|
||||
out := protoimpl.TypeBuilder{
|
||||
File: protoimpl.DescBuilder{
|
||||
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
|
||||
RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_log_log_proto_rawDesc), len(file_agent_log_log_proto_rawDesc)),
|
||||
NumEnums: 0,
|
||||
NumMessages: 2,
|
||||
NumExtensions: 0,
|
||||
NumServices: 1,
|
||||
},
|
||||
GoTypes: file_agent_log_log_proto_goTypes,
|
||||
DependencyIndexes: file_agent_log_log_proto_depIdxs,
|
||||
MessageInfos: file_agent_log_log_proto_msgTypes,
|
||||
}.Build()
|
||||
File_agent_log_log_proto = out.File
|
||||
file_agent_log_log_proto_goTypes = nil
|
||||
file_agent_log_log_proto_depIdxs = nil
|
||||
}
|
||||
@@ -1,32 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package log;
|
||||
|
||||
option go_package = "./log";
|
||||
|
||||
import "google/protobuf/timestamp.proto";
|
||||
import "google/protobuf/empty.proto";
|
||||
|
||||
service LogCollector {
|
||||
rpc SendLog(LogEntry) returns (google.protobuf.Empty);
|
||||
rpc SendEvent(EventEntry) returns (google.protobuf.Empty);
|
||||
}
|
||||
|
||||
message LogEntry {
|
||||
string message = 1;
|
||||
string computation_id = 2;
|
||||
string level = 3;
|
||||
google.protobuf.Timestamp timestamp = 4;
|
||||
}
|
||||
|
||||
message EventEntry {
|
||||
string event_type = 1;
|
||||
google.protobuf.Timestamp timestamp = 2;
|
||||
string computation_id = 3;
|
||||
bytes details = 4; // JSON payload
|
||||
string originator = 5;
|
||||
string status = 6;
|
||||
}
|
||||
@@ -1,163 +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.6.0
|
||||
// - protoc v7.35.0
|
||||
// source: agent/log/log.proto
|
||||
|
||||
package log
|
||||
|
||||
import (
|
||||
context "context"
|
||||
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
|
||||
// is compatible with the grpc package it is being compiled against.
|
||||
// Requires gRPC-Go v1.64.0 or later.
|
||||
const _ = grpc.SupportPackageIsVersion9
|
||||
|
||||
const (
|
||||
LogCollector_SendLog_FullMethodName = "/log.LogCollector/SendLog"
|
||||
LogCollector_SendEvent_FullMethodName = "/log.LogCollector/SendEvent"
|
||||
)
|
||||
|
||||
// LogCollectorClient is the client API for LogCollector 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 LogCollectorClient interface {
|
||||
SendLog(ctx context.Context, in *LogEntry, opts ...grpc.CallOption) (*emptypb.Empty, error)
|
||||
SendEvent(ctx context.Context, in *EventEntry, opts ...grpc.CallOption) (*emptypb.Empty, error)
|
||||
}
|
||||
|
||||
type logCollectorClient struct {
|
||||
cc grpc.ClientConnInterface
|
||||
}
|
||||
|
||||
func NewLogCollectorClient(cc grpc.ClientConnInterface) LogCollectorClient {
|
||||
return &logCollectorClient{cc}
|
||||
}
|
||||
|
||||
func (c *logCollectorClient) SendLog(ctx context.Context, in *LogEntry, opts ...grpc.CallOption) (*emptypb.Empty, error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(emptypb.Empty)
|
||||
err := c.cc.Invoke(ctx, LogCollector_SendLog_FullMethodName, in, out, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (c *logCollectorClient) SendEvent(ctx context.Context, in *EventEntry, opts ...grpc.CallOption) (*emptypb.Empty, error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(emptypb.Empty)
|
||||
err := c.cc.Invoke(ctx, LogCollector_SendEvent_FullMethodName, in, out, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// LogCollectorServer is the server API for LogCollector service.
|
||||
// All implementations must embed UnimplementedLogCollectorServer
|
||||
// for forward compatibility.
|
||||
type LogCollectorServer interface {
|
||||
SendLog(context.Context, *LogEntry) (*emptypb.Empty, error)
|
||||
SendEvent(context.Context, *EventEntry) (*emptypb.Empty, error)
|
||||
mustEmbedUnimplementedLogCollectorServer()
|
||||
}
|
||||
|
||||
// UnimplementedLogCollectorServer 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 UnimplementedLogCollectorServer struct{}
|
||||
|
||||
func (UnimplementedLogCollectorServer) SendLog(context.Context, *LogEntry) (*emptypb.Empty, error) {
|
||||
return nil, status.Error(codes.Unimplemented, "method SendLog not implemented")
|
||||
}
|
||||
func (UnimplementedLogCollectorServer) SendEvent(context.Context, *EventEntry) (*emptypb.Empty, error) {
|
||||
return nil, status.Error(codes.Unimplemented, "method SendEvent not implemented")
|
||||
}
|
||||
func (UnimplementedLogCollectorServer) mustEmbedUnimplementedLogCollectorServer() {}
|
||||
func (UnimplementedLogCollectorServer) testEmbeddedByValue() {}
|
||||
|
||||
// UnsafeLogCollectorServer may be embedded to opt out of forward compatibility for this service.
|
||||
// Use of this interface is not recommended, as added methods to LogCollectorServer will
|
||||
// result in compilation errors.
|
||||
type UnsafeLogCollectorServer interface {
|
||||
mustEmbedUnimplementedLogCollectorServer()
|
||||
}
|
||||
|
||||
func RegisterLogCollectorServer(s grpc.ServiceRegistrar, srv LogCollectorServer) {
|
||||
// If the following call panics, it indicates UnimplementedLogCollectorServer 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(&LogCollector_ServiceDesc, srv)
|
||||
}
|
||||
|
||||
func _LogCollector_SendLog_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(LogEntry)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(LogCollectorServer).SendLog(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: LogCollector_SendLog_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(LogCollectorServer).SendLog(ctx, req.(*LogEntry))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
func _LogCollector_SendEvent_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(EventEntry)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(LogCollectorServer).SendEvent(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: LogCollector_SendEvent_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(LogCollectorServer).SendEvent(ctx, req.(*EventEntry))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
// LogCollector_ServiceDesc is the grpc.ServiceDesc for LogCollector service.
|
||||
// It's only intended for direct use with grpc.RegisterService,
|
||||
// and not to be introspected or modified (even as a copy)
|
||||
var LogCollector_ServiceDesc = grpc.ServiceDesc{
|
||||
ServiceName: "log.LogCollector",
|
||||
HandlerType: (*LogCollectorServer)(nil),
|
||||
Methods: []grpc.MethodDesc{
|
||||
{
|
||||
MethodName: "SendLog",
|
||||
Handler: _LogCollector_SendLog_Handler,
|
||||
},
|
||||
{
|
||||
MethodName: "SendEvent",
|
||||
Handler: _LogCollector_SendEvent_Handler,
|
||||
},
|
||||
},
|
||||
Streams: []grpc.StreamDesc{},
|
||||
Metadata: "agent/log/log.proto",
|
||||
}
|
||||
@@ -1,59 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"github.com/ultravioletrs/cocos/agent/log"
|
||||
"google.golang.org/protobuf/types/known/emptypb"
|
||||
)
|
||||
|
||||
var _ log.LogCollectorServer = (*LogForwarder)(nil)
|
||||
|
||||
type LogForwarder struct {
|
||||
log.UnimplementedLogCollectorServer
|
||||
cvmsClient cvms.ServiceClient
|
||||
logger *slog.Logger
|
||||
logQueue chan *cvms.ClientStreamMessage
|
||||
}
|
||||
|
||||
func New(logger *slog.Logger, cvmsClient cvms.ServiceClient, queue chan *cvms.ClientStreamMessage) *LogForwarder {
|
||||
return &LogForwarder{
|
||||
cvmsClient: cvmsClient,
|
||||
logger: logger,
|
||||
logQueue: queue,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *LogForwarder) SendLog(ctx context.Context, req *log.LogEntry) (*emptypb.Empty, error) {
|
||||
s.logQueue <- &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_AgentLog{
|
||||
AgentLog: &cvms.AgentLog{
|
||||
Message: req.Message,
|
||||
ComputationId: req.ComputationId,
|
||||
Level: req.Level,
|
||||
Timestamp: req.Timestamp,
|
||||
},
|
||||
},
|
||||
}
|
||||
return &emptypb.Empty{}, nil
|
||||
}
|
||||
|
||||
func (s *LogForwarder) SendEvent(ctx context.Context, req *log.EventEntry) (*emptypb.Empty, error) {
|
||||
s.logQueue <- &cvms.ClientStreamMessage{
|
||||
Message: &cvms.ClientStreamMessage_AgentEvent{
|
||||
AgentEvent: &cvms.AgentEvent{
|
||||
EventType: req.EventType,
|
||||
Timestamp: req.Timestamp,
|
||||
ComputationId: req.ComputationId,
|
||||
Details: req.Details,
|
||||
Originator: req.Originator,
|
||||
Status: req.Status,
|
||||
},
|
||||
},
|
||||
}
|
||||
return &emptypb.Empty{}, nil
|
||||
}
|
||||
@@ -1,303 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
"github.com/ultravioletrs/cocos/agent/log"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
)
|
||||
|
||||
// TestNewLogForwarder tests the creation of a new log forwarder.
|
||||
func TestNewLogForwarder(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
require.NotNil(t, lf)
|
||||
assert.NotNil(t, lf.logger)
|
||||
assert.Nil(t, lf.cvmsClient)
|
||||
assert.NotNil(t, lf.logQueue)
|
||||
}
|
||||
|
||||
// TestSendLog tests sending a log entry.
|
||||
func TestSendLog(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
req := &log.LogEntry{
|
||||
Message: "Test log message",
|
||||
ComputationId: "computation-1",
|
||||
Level: "INFO",
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
}
|
||||
|
||||
resp, err := lf.SendLog(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
|
||||
// Verify message was queued
|
||||
select {
|
||||
case msg := <-queue:
|
||||
require.NotNil(t, msg)
|
||||
agentLog := msg.GetAgentLog()
|
||||
assert.NotNil(t, agentLog)
|
||||
assert.Equal(t, "Test log message", agentLog.Message)
|
||||
assert.Equal(t, "computation-1", agentLog.ComputationId)
|
||||
assert.Equal(t, "INFO", agentLog.Level)
|
||||
default:
|
||||
t.Fatal("No message in queue")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSendEvent tests sending an event entry.
|
||||
func TestSendEvent(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
details, err := json.Marshal(map[string]string{"key": "value"})
|
||||
require.NoError(t, err)
|
||||
|
||||
req := &log.EventEntry{
|
||||
EventType: "COMPUTATION_STARTED",
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
ComputationId: "computation-1",
|
||||
Details: details,
|
||||
Originator: "runner",
|
||||
Status: "SUCCESS",
|
||||
}
|
||||
|
||||
resp, err := lf.SendEvent(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
|
||||
// Verify message was queued
|
||||
select {
|
||||
case msg := <-queue:
|
||||
require.NotNil(t, msg)
|
||||
agentEvent := msg.GetAgentEvent()
|
||||
assert.NotNil(t, agentEvent)
|
||||
assert.Equal(t, "COMPUTATION_STARTED", agentEvent.EventType)
|
||||
assert.Equal(t, "computation-1", agentEvent.ComputationId)
|
||||
assert.Equal(t, "runner", agentEvent.Originator)
|
||||
assert.Equal(t, "SUCCESS", agentEvent.Status)
|
||||
default:
|
||||
t.Fatal("No message in queue")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSendMultipleLogs tests sending multiple log entries.
|
||||
func TestSendMultipleLogs(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 100)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
req := &log.LogEntry{
|
||||
Message: "Log message",
|
||||
ComputationId: "computation-1",
|
||||
Level: "INFO",
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
}
|
||||
|
||||
resp, err := lf.SendLog(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
}
|
||||
|
||||
assert.Equal(t, 5, len(queue))
|
||||
}
|
||||
|
||||
// TestSendEventWithVariousTypes tests sending events with different types.
|
||||
func TestSendEventWithVariousTypes(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 100)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
eventTypes := []string{"STARTED", "RUNNING", "COMPLETED", "FAILED"}
|
||||
for _, eventType := range eventTypes {
|
||||
details, err := json.Marshal(map[string]string{"type": eventType})
|
||||
require.NoError(t, err)
|
||||
req := &log.EventEntry{
|
||||
EventType: eventType,
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
ComputationId: "computation-1",
|
||||
Details: details,
|
||||
Originator: "runner",
|
||||
Status: "OK",
|
||||
}
|
||||
|
||||
resp, err := lf.SendEvent(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
}
|
||||
|
||||
assert.Equal(t, 4, len(queue))
|
||||
}
|
||||
|
||||
// TestSendLogWithEmptyMessage tests sending log with empty message.
|
||||
func TestSendLogWithEmptyMessage(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
req := &log.LogEntry{
|
||||
Message: "",
|
||||
ComputationId: "computation-1",
|
||||
Level: "INFO",
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
}
|
||||
|
||||
resp, err := lf.SendLog(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
|
||||
select {
|
||||
case msg := <-queue:
|
||||
agentLog := msg.GetAgentLog()
|
||||
assert.Equal(t, "", agentLog.Message)
|
||||
default:
|
||||
t.Fatal("No message in queue")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSendEventWithNilDetails tests sending event with nil details.
|
||||
func TestSendEventWithNilDetails(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 10)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
req := &log.EventEntry{
|
||||
EventType: "TEST_EVENT",
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
ComputationId: "computation-1",
|
||||
Details: nil,
|
||||
Originator: "test",
|
||||
Status: "OK",
|
||||
}
|
||||
|
||||
resp, err := lf.SendEvent(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
|
||||
select {
|
||||
case msg := <-queue:
|
||||
agentEvent := msg.GetAgentEvent()
|
||||
assert.Nil(t, agentEvent.Details)
|
||||
default:
|
||||
t.Fatal("No message in queue")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSendLogWithVariousLevels tests sending logs with various severity levels.
|
||||
func TestSendLogWithVariousLevels(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 100)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
levels := []string{"DEBUG", "INFO", "WARN", "ERROR"}
|
||||
for _, level := range levels {
|
||||
req := &log.LogEntry{
|
||||
Message: "Test " + level,
|
||||
ComputationId: "computation-1",
|
||||
Level: level,
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
}
|
||||
|
||||
resp, err := lf.SendLog(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
}
|
||||
|
||||
assert.Equal(t, 4, len(queue))
|
||||
}
|
||||
|
||||
// TestSendLogWithDifferentComputationIds tests sending logs with different computation IDs.
|
||||
func TestSendLogWithDifferentComputationIds(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 100)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
req := &log.LogEntry{
|
||||
Message: "Message",
|
||||
ComputationId: "computation-" + string(rune(48+i)),
|
||||
Level: "INFO",
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
}
|
||||
|
||||
resp, err := lf.SendLog(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
}
|
||||
|
||||
assert.Equal(t, 3, len(queue))
|
||||
}
|
||||
|
||||
// TestQueueBehavior tests that queue is properly used.
|
||||
func TestQueueBehavior(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 1)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
req := &log.LogEntry{
|
||||
Message: "Test",
|
||||
ComputationId: "computation-1",
|
||||
Level: "INFO",
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
}
|
||||
|
||||
resp, err := lf.SendLog(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
|
||||
assert.Equal(t, 1, len(queue))
|
||||
}
|
||||
|
||||
// TestConcurrentSendLog tests concurrent log sending.
|
||||
func TestConcurrentSendLog(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
queue := make(chan *cvms.ClientStreamMessage, 100)
|
||||
|
||||
lf := New(logger, nil, queue)
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
go func(id int) {
|
||||
req := &log.LogEntry{
|
||||
Message: "Concurrent log",
|
||||
ComputationId: "computation-1",
|
||||
Level: "INFO",
|
||||
Timestamp: timestamppb.New(time.Now()),
|
||||
}
|
||||
|
||||
_, err := lf.SendLog(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
}(i)
|
||||
}
|
||||
|
||||
// Give goroutines time to complete
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Should have received all messages
|
||||
assert.True(t, len(queue) > 0)
|
||||
}
|
||||
@@ -1,39 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
)
|
||||
|
||||
type MockAttestationClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockAttestationClient) GetAttestation(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
|
||||
args := m.Called(ctx, reportData, nonce, attType)
|
||||
return args.Get(0).([]byte), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockAttestationClient) GetRawEvidence(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
|
||||
args := m.Called(ctx, reportData, nonce, attType)
|
||||
return args.Get(0).([]byte), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockAttestationClient) GetAzureToken(ctx context.Context, nonce [32]byte) ([]byte, error) {
|
||||
args := m.Called(ctx, nonce)
|
||||
return args.Get(0).([]byte), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockAttestationClient) GetKbsToken(ctx context.Context) ([]byte, error) {
|
||||
args := m.Called(ctx)
|
||||
return args.Get(0).([]byte), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockAttestationClient) Close() error {
|
||||
args := m.Called()
|
||||
return args.Error(0)
|
||||
}
|
||||
@@ -1,442 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
// NewAgentService_AlgoClient creates a new instance of AgentService_AlgoClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
|
||||
// The first argument is typically a *testing.T value.
|
||||
func NewAgentService_AlgoClient(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *AgentService_AlgoClient {
|
||||
mock := &AgentService_AlgoClient{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient is an autogenerated mock type for the AgentService_AlgoClient type
|
||||
type AgentService_AlgoClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type AgentService_AlgoClient_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *AgentService_AlgoClient) EXPECT() *AgentService_AlgoClient_Expecter {
|
||||
return &AgentService_AlgoClient_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// CloseAndRecv provides a mock function for the type AgentService_AlgoClient
|
||||
func (_mock *AgentService_AlgoClient) CloseAndRecv() (*agent.AlgoResponse, error) {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for CloseAndRecv")
|
||||
}
|
||||
|
||||
var r0 *agent.AlgoResponse
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func() (*agent.AlgoResponse, error)); ok {
|
||||
return returnFunc()
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func() *agent.AlgoResponse); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*agent.AlgoResponse)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func() error); ok {
|
||||
r1 = returnFunc()
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
|
||||
type AgentService_AlgoClient_CloseAndRecv_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// CloseAndRecv is a helper method to define mock.On call
|
||||
func (_e *AgentService_AlgoClient_Expecter) CloseAndRecv() *AgentService_AlgoClient_CloseAndRecv_Call {
|
||||
return &AgentService_AlgoClient_CloseAndRecv_Call{Call: _e.mock.On("CloseAndRecv")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_CloseAndRecv_Call) Run(run func()) *AgentService_AlgoClient_CloseAndRecv_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_CloseAndRecv_Call) Return(algoResponse *agent.AlgoResponse, err error) *AgentService_AlgoClient_CloseAndRecv_Call {
|
||||
_c.Call.Return(algoResponse, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_CloseAndRecv_Call) RunAndReturn(run func() (*agent.AlgoResponse, error)) *AgentService_AlgoClient_CloseAndRecv_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// CloseSend provides a mock function for the type AgentService_AlgoClient
|
||||
func (_mock *AgentService_AlgoClient) CloseSend() error {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for CloseSend")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func() error); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
|
||||
type AgentService_AlgoClient_CloseSend_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// CloseSend is a helper method to define mock.On call
|
||||
func (_e *AgentService_AlgoClient_Expecter) CloseSend() *AgentService_AlgoClient_CloseSend_Call {
|
||||
return &AgentService_AlgoClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_CloseSend_Call) Run(run func()) *AgentService_AlgoClient_CloseSend_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_CloseSend_Call) Return(err error) *AgentService_AlgoClient_CloseSend_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_AlgoClient_CloseSend_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Context provides a mock function for the type AgentService_AlgoClient
|
||||
func (_mock *AgentService_AlgoClient) Context() context.Context {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Context")
|
||||
}
|
||||
|
||||
var r0 context.Context
|
||||
if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(context.Context)
|
||||
}
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
|
||||
type AgentService_AlgoClient_Context_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Context is a helper method to define mock.On call
|
||||
func (_e *AgentService_AlgoClient_Expecter) Context() *AgentService_AlgoClient_Context_Call {
|
||||
return &AgentService_AlgoClient_Context_Call{Call: _e.mock.On("Context")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Context_Call) Run(run func()) *AgentService_AlgoClient_Context_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Context_Call) Return(context1 context.Context) *AgentService_AlgoClient_Context_Call {
|
||||
_c.Call.Return(context1)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_AlgoClient_Context_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Header provides a mock function for the type AgentService_AlgoClient
|
||||
func (_mock *AgentService_AlgoClient) Header() (metadata.MD, error) {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Header")
|
||||
}
|
||||
|
||||
var r0 metadata.MD
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
|
||||
return returnFunc()
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(metadata.MD)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func() error); ok {
|
||||
r1 = returnFunc()
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
|
||||
type AgentService_AlgoClient_Header_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Header is a helper method to define mock.On call
|
||||
func (_e *AgentService_AlgoClient_Expecter) Header() *AgentService_AlgoClient_Header_Call {
|
||||
return &AgentService_AlgoClient_Header_Call{Call: _e.mock.On("Header")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Header_Call) Run(run func()) *AgentService_AlgoClient_Header_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_AlgoClient_Header_Call {
|
||||
_c.Call.Return(mD, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_AlgoClient_Header_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// RecvMsg provides a mock function for the type AgentService_AlgoClient
|
||||
func (_mock *AgentService_AlgoClient) RecvMsg(m any) error {
|
||||
ret := _mock.Called(m)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for RecvMsg")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(any) error); ok {
|
||||
r0 = returnFunc(m)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
|
||||
type AgentService_AlgoClient_RecvMsg_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// RecvMsg is a helper method to define mock.On call
|
||||
// - m any
|
||||
func (_e *AgentService_AlgoClient_Expecter) RecvMsg(m interface{}) *AgentService_AlgoClient_RecvMsg_Call {
|
||||
return &AgentService_AlgoClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_RecvMsg_Call) Run(run func(m any)) *AgentService_AlgoClient_RecvMsg_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 any
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(any)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_RecvMsg_Call) Return(err error) *AgentService_AlgoClient_RecvMsg_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_AlgoClient_RecvMsg_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Send provides a mock function for the type AgentService_AlgoClient
|
||||
func (_mock *AgentService_AlgoClient) Send(algoRequest *agent.AlgoRequest) error {
|
||||
ret := _mock.Called(algoRequest)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Send")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(*agent.AlgoRequest) error); ok {
|
||||
r0 = returnFunc(algoRequest)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
|
||||
type AgentService_AlgoClient_Send_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Send is a helper method to define mock.On call
|
||||
// - algoRequest *agent.AlgoRequest
|
||||
func (_e *AgentService_AlgoClient_Expecter) Send(algoRequest interface{}) *AgentService_AlgoClient_Send_Call {
|
||||
return &AgentService_AlgoClient_Send_Call{Call: _e.mock.On("Send", algoRequest)}
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Send_Call) Run(run func(algoRequest *agent.AlgoRequest)) *AgentService_AlgoClient_Send_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 *agent.AlgoRequest
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(*agent.AlgoRequest)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Send_Call) Return(err error) *AgentService_AlgoClient_Send_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Send_Call) RunAndReturn(run func(algoRequest *agent.AlgoRequest) error) *AgentService_AlgoClient_Send_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SendMsg provides a mock function for the type AgentService_AlgoClient
|
||||
func (_mock *AgentService_AlgoClient) SendMsg(m any) error {
|
||||
ret := _mock.Called(m)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for SendMsg")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(any) error); ok {
|
||||
r0 = returnFunc(m)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
|
||||
type AgentService_AlgoClient_SendMsg_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// SendMsg is a helper method to define mock.On call
|
||||
// - m any
|
||||
func (_e *AgentService_AlgoClient_Expecter) SendMsg(m interface{}) *AgentService_AlgoClient_SendMsg_Call {
|
||||
return &AgentService_AlgoClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_SendMsg_Call) Run(run func(m any)) *AgentService_AlgoClient_SendMsg_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 any
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(any)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_SendMsg_Call) Return(err error) *AgentService_AlgoClient_SendMsg_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_AlgoClient_SendMsg_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Trailer provides a mock function for the type AgentService_AlgoClient
|
||||
func (_mock *AgentService_AlgoClient) Trailer() metadata.MD {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Trailer")
|
||||
}
|
||||
|
||||
var r0 metadata.MD
|
||||
if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(metadata.MD)
|
||||
}
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_AlgoClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
|
||||
type AgentService_AlgoClient_Trailer_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Trailer is a helper method to define mock.On call
|
||||
func (_e *AgentService_AlgoClient_Expecter) Trailer() *AgentService_AlgoClient_Trailer_Call {
|
||||
return &AgentService_AlgoClient_Trailer_Call{Call: _e.mock.On("Trailer")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Trailer_Call) Run(run func()) *AgentService_AlgoClient_Trailer_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Trailer_Call) Return(mD metadata.MD) *AgentService_AlgoClient_Trailer_Call {
|
||||
_c.Call.Return(mD)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_AlgoClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_AlgoClient_Trailer_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -1,442 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
// NewAgentService_DataClient creates a new instance of AgentService_DataClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
|
||||
// The first argument is typically a *testing.T value.
|
||||
func NewAgentService_DataClient(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *AgentService_DataClient {
|
||||
mock := &AgentService_DataClient{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
|
||||
// AgentService_DataClient is an autogenerated mock type for the AgentService_DataClient type
|
||||
type AgentService_DataClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type AgentService_DataClient_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *AgentService_DataClient) EXPECT() *AgentService_DataClient_Expecter {
|
||||
return &AgentService_DataClient_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// CloseAndRecv provides a mock function for the type AgentService_DataClient
|
||||
func (_mock *AgentService_DataClient) CloseAndRecv() (*agent.DataResponse, error) {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for CloseAndRecv")
|
||||
}
|
||||
|
||||
var r0 *agent.DataResponse
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func() (*agent.DataResponse, error)); ok {
|
||||
return returnFunc()
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func() *agent.DataResponse); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*agent.DataResponse)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func() error); ok {
|
||||
r1 = returnFunc()
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// AgentService_DataClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
|
||||
type AgentService_DataClient_CloseAndRecv_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// CloseAndRecv is a helper method to define mock.On call
|
||||
func (_e *AgentService_DataClient_Expecter) CloseAndRecv() *AgentService_DataClient_CloseAndRecv_Call {
|
||||
return &AgentService_DataClient_CloseAndRecv_Call{Call: _e.mock.On("CloseAndRecv")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_CloseAndRecv_Call) Run(run func()) *AgentService_DataClient_CloseAndRecv_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_CloseAndRecv_Call) Return(dataResponse *agent.DataResponse, err error) *AgentService_DataClient_CloseAndRecv_Call {
|
||||
_c.Call.Return(dataResponse, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_CloseAndRecv_Call) RunAndReturn(run func() (*agent.DataResponse, error)) *AgentService_DataClient_CloseAndRecv_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// CloseSend provides a mock function for the type AgentService_DataClient
|
||||
func (_mock *AgentService_DataClient) CloseSend() error {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for CloseSend")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func() error); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_DataClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
|
||||
type AgentService_DataClient_CloseSend_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// CloseSend is a helper method to define mock.On call
|
||||
func (_e *AgentService_DataClient_Expecter) CloseSend() *AgentService_DataClient_CloseSend_Call {
|
||||
return &AgentService_DataClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_CloseSend_Call) Run(run func()) *AgentService_DataClient_CloseSend_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_CloseSend_Call) Return(err error) *AgentService_DataClient_CloseSend_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_DataClient_CloseSend_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Context provides a mock function for the type AgentService_DataClient
|
||||
func (_mock *AgentService_DataClient) Context() context.Context {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Context")
|
||||
}
|
||||
|
||||
var r0 context.Context
|
||||
if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(context.Context)
|
||||
}
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_DataClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
|
||||
type AgentService_DataClient_Context_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Context is a helper method to define mock.On call
|
||||
func (_e *AgentService_DataClient_Expecter) Context() *AgentService_DataClient_Context_Call {
|
||||
return &AgentService_DataClient_Context_Call{Call: _e.mock.On("Context")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Context_Call) Run(run func()) *AgentService_DataClient_Context_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Context_Call) Return(context1 context.Context) *AgentService_DataClient_Context_Call {
|
||||
_c.Call.Return(context1)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_DataClient_Context_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Header provides a mock function for the type AgentService_DataClient
|
||||
func (_mock *AgentService_DataClient) Header() (metadata.MD, error) {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Header")
|
||||
}
|
||||
|
||||
var r0 metadata.MD
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
|
||||
return returnFunc()
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(metadata.MD)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func() error); ok {
|
||||
r1 = returnFunc()
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// AgentService_DataClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
|
||||
type AgentService_DataClient_Header_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Header is a helper method to define mock.On call
|
||||
func (_e *AgentService_DataClient_Expecter) Header() *AgentService_DataClient_Header_Call {
|
||||
return &AgentService_DataClient_Header_Call{Call: _e.mock.On("Header")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Header_Call) Run(run func()) *AgentService_DataClient_Header_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_DataClient_Header_Call {
|
||||
_c.Call.Return(mD, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_DataClient_Header_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// RecvMsg provides a mock function for the type AgentService_DataClient
|
||||
func (_mock *AgentService_DataClient) RecvMsg(m any) error {
|
||||
ret := _mock.Called(m)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for RecvMsg")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(any) error); ok {
|
||||
r0 = returnFunc(m)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_DataClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
|
||||
type AgentService_DataClient_RecvMsg_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// RecvMsg is a helper method to define mock.On call
|
||||
// - m any
|
||||
func (_e *AgentService_DataClient_Expecter) RecvMsg(m interface{}) *AgentService_DataClient_RecvMsg_Call {
|
||||
return &AgentService_DataClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_RecvMsg_Call) Run(run func(m any)) *AgentService_DataClient_RecvMsg_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 any
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(any)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_RecvMsg_Call) Return(err error) *AgentService_DataClient_RecvMsg_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_DataClient_RecvMsg_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Send provides a mock function for the type AgentService_DataClient
|
||||
func (_mock *AgentService_DataClient) Send(dataRequest *agent.DataRequest) error {
|
||||
ret := _mock.Called(dataRequest)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Send")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(*agent.DataRequest) error); ok {
|
||||
r0 = returnFunc(dataRequest)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_DataClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
|
||||
type AgentService_DataClient_Send_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Send is a helper method to define mock.On call
|
||||
// - dataRequest *agent.DataRequest
|
||||
func (_e *AgentService_DataClient_Expecter) Send(dataRequest interface{}) *AgentService_DataClient_Send_Call {
|
||||
return &AgentService_DataClient_Send_Call{Call: _e.mock.On("Send", dataRequest)}
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Send_Call) Run(run func(dataRequest *agent.DataRequest)) *AgentService_DataClient_Send_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 *agent.DataRequest
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(*agent.DataRequest)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Send_Call) Return(err error) *AgentService_DataClient_Send_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Send_Call) RunAndReturn(run func(dataRequest *agent.DataRequest) error) *AgentService_DataClient_Send_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SendMsg provides a mock function for the type AgentService_DataClient
|
||||
func (_mock *AgentService_DataClient) SendMsg(m any) error {
|
||||
ret := _mock.Called(m)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for SendMsg")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(any) error); ok {
|
||||
r0 = returnFunc(m)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_DataClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
|
||||
type AgentService_DataClient_SendMsg_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// SendMsg is a helper method to define mock.On call
|
||||
// - m any
|
||||
func (_e *AgentService_DataClient_Expecter) SendMsg(m interface{}) *AgentService_DataClient_SendMsg_Call {
|
||||
return &AgentService_DataClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_SendMsg_Call) Run(run func(m any)) *AgentService_DataClient_SendMsg_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 any
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(any)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_SendMsg_Call) Return(err error) *AgentService_DataClient_SendMsg_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_DataClient_SendMsg_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Trailer provides a mock function for the type AgentService_DataClient
|
||||
func (_mock *AgentService_DataClient) Trailer() metadata.MD {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Trailer")
|
||||
}
|
||||
|
||||
var r0 metadata.MD
|
||||
if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(metadata.MD)
|
||||
}
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_DataClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
|
||||
type AgentService_DataClient_Trailer_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Trailer is a helper method to define mock.On call
|
||||
func (_e *AgentService_DataClient_Expecter) Trailer() *AgentService_DataClient_Trailer_Call {
|
||||
return &AgentService_DataClient_Trailer_Call{Call: _e.mock.On("Trailer")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Trailer_Call) Run(run func()) *AgentService_DataClient_Trailer_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Trailer_Call) Return(mD metadata.MD) *AgentService_DataClient_Trailer_Call {
|
||||
_c.Call.Return(mD)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_DataClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_DataClient_Trailer_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -1,391 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
// NewAgentService_IMAMeasurementsClient creates a new instance of AgentService_IMAMeasurementsClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
|
||||
// The first argument is typically a *testing.T value.
|
||||
func NewAgentService_IMAMeasurementsClient(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *AgentService_IMAMeasurementsClient {
|
||||
mock := &AgentService_IMAMeasurementsClient{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
|
||||
// AgentService_IMAMeasurementsClient is an autogenerated mock type for the AgentService_IMAMeasurementsClient type
|
||||
type AgentService_IMAMeasurementsClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type AgentService_IMAMeasurementsClient_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *AgentService_IMAMeasurementsClient) EXPECT() *AgentService_IMAMeasurementsClient_Expecter {
|
||||
return &AgentService_IMAMeasurementsClient_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// CloseSend provides a mock function for the type AgentService_IMAMeasurementsClient
|
||||
func (_mock *AgentService_IMAMeasurementsClient) CloseSend() error {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for CloseSend")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func() error); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_IMAMeasurementsClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
|
||||
type AgentService_IMAMeasurementsClient_CloseSend_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// CloseSend is a helper method to define mock.On call
|
||||
func (_e *AgentService_IMAMeasurementsClient_Expecter) CloseSend() *AgentService_IMAMeasurementsClient_CloseSend_Call {
|
||||
return &AgentService_IMAMeasurementsClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) Run(run func()) *AgentService_IMAMeasurementsClient_CloseSend_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) Return(err error) *AgentService_IMAMeasurementsClient_CloseSend_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_IMAMeasurementsClient_CloseSend_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Context provides a mock function for the type AgentService_IMAMeasurementsClient
|
||||
func (_mock *AgentService_IMAMeasurementsClient) Context() context.Context {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Context")
|
||||
}
|
||||
|
||||
var r0 context.Context
|
||||
if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(context.Context)
|
||||
}
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_IMAMeasurementsClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
|
||||
type AgentService_IMAMeasurementsClient_Context_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Context is a helper method to define mock.On call
|
||||
func (_e *AgentService_IMAMeasurementsClient_Expecter) Context() *AgentService_IMAMeasurementsClient_Context_Call {
|
||||
return &AgentService_IMAMeasurementsClient_Context_Call{Call: _e.mock.On("Context")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Context_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Context_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Context_Call) Return(context1 context.Context) *AgentService_IMAMeasurementsClient_Context_Call {
|
||||
_c.Call.Return(context1)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_IMAMeasurementsClient_Context_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Header provides a mock function for the type AgentService_IMAMeasurementsClient
|
||||
func (_mock *AgentService_IMAMeasurementsClient) Header() (metadata.MD, error) {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Header")
|
||||
}
|
||||
|
||||
var r0 metadata.MD
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
|
||||
return returnFunc()
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(metadata.MD)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func() error); ok {
|
||||
r1 = returnFunc()
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// AgentService_IMAMeasurementsClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
|
||||
type AgentService_IMAMeasurementsClient_Header_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Header is a helper method to define mock.On call
|
||||
func (_e *AgentService_IMAMeasurementsClient_Expecter) Header() *AgentService_IMAMeasurementsClient_Header_Call {
|
||||
return &AgentService_IMAMeasurementsClient_Header_Call{Call: _e.mock.On("Header")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Header_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Header_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_IMAMeasurementsClient_Header_Call {
|
||||
_c.Call.Return(mD, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_IMAMeasurementsClient_Header_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Recv provides a mock function for the type AgentService_IMAMeasurementsClient
|
||||
func (_mock *AgentService_IMAMeasurementsClient) Recv() (*agent.IMAMeasurementsResponse, error) {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Recv")
|
||||
}
|
||||
|
||||
var r0 *agent.IMAMeasurementsResponse
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func() (*agent.IMAMeasurementsResponse, error)); ok {
|
||||
return returnFunc()
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func() *agent.IMAMeasurementsResponse); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*agent.IMAMeasurementsResponse)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func() error); ok {
|
||||
r1 = returnFunc()
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// AgentService_IMAMeasurementsClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv'
|
||||
type AgentService_IMAMeasurementsClient_Recv_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Recv is a helper method to define mock.On call
|
||||
func (_e *AgentService_IMAMeasurementsClient_Expecter) Recv() *AgentService_IMAMeasurementsClient_Recv_Call {
|
||||
return &AgentService_IMAMeasurementsClient_Recv_Call{Call: _e.mock.On("Recv")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Recv_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Recv_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Recv_Call) Return(iMAMeasurementsResponse *agent.IMAMeasurementsResponse, err error) *AgentService_IMAMeasurementsClient_Recv_Call {
|
||||
_c.Call.Return(iMAMeasurementsResponse, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Recv_Call) RunAndReturn(run func() (*agent.IMAMeasurementsResponse, error)) *AgentService_IMAMeasurementsClient_Recv_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// RecvMsg provides a mock function for the type AgentService_IMAMeasurementsClient
|
||||
func (_mock *AgentService_IMAMeasurementsClient) RecvMsg(m any) error {
|
||||
ret := _mock.Called(m)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for RecvMsg")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(any) error); ok {
|
||||
r0 = returnFunc(m)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_IMAMeasurementsClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
|
||||
type AgentService_IMAMeasurementsClient_RecvMsg_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// RecvMsg is a helper method to define mock.On call
|
||||
// - m any
|
||||
func (_e *AgentService_IMAMeasurementsClient_Expecter) RecvMsg(m interface{}) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
|
||||
return &AgentService_IMAMeasurementsClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) Run(run func(m any)) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 any
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(any)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) Return(err error) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SendMsg provides a mock function for the type AgentService_IMAMeasurementsClient
|
||||
func (_mock *AgentService_IMAMeasurementsClient) SendMsg(m any) error {
|
||||
ret := _mock.Called(m)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for SendMsg")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(any) error); ok {
|
||||
r0 = returnFunc(m)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_IMAMeasurementsClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
|
||||
type AgentService_IMAMeasurementsClient_SendMsg_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// SendMsg is a helper method to define mock.On call
|
||||
// - m any
|
||||
func (_e *AgentService_IMAMeasurementsClient_Expecter) SendMsg(m interface{}) *AgentService_IMAMeasurementsClient_SendMsg_Call {
|
||||
return &AgentService_IMAMeasurementsClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) Run(run func(m any)) *AgentService_IMAMeasurementsClient_SendMsg_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 any
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(any)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) Return(err error) *AgentService_IMAMeasurementsClient_SendMsg_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_IMAMeasurementsClient_SendMsg_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Trailer provides a mock function for the type AgentService_IMAMeasurementsClient
|
||||
func (_mock *AgentService_IMAMeasurementsClient) Trailer() metadata.MD {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Trailer")
|
||||
}
|
||||
|
||||
var r0 metadata.MD
|
||||
if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(metadata.MD)
|
||||
}
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// AgentService_IMAMeasurementsClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
|
||||
type AgentService_IMAMeasurementsClient_Trailer_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Trailer is a helper method to define mock.On call
|
||||
func (_e *AgentService_IMAMeasurementsClient_Expecter) Trailer() *AgentService_IMAMeasurementsClient_Trailer_Call {
|
||||
return &AgentService_IMAMeasurementsClient_Trailer_Call{Call: _e.mock.On("Trailer")}
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Trailer_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) Return(mD metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call {
|
||||
_c.Call.Return(mD)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -1,589 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
)
|
||||
|
||||
// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
|
||||
// The first argument is typically a *testing.T value.
|
||||
func NewService(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *Service {
|
||||
mock := &Service{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
|
||||
// Service is an autogenerated mock type for the Service type
|
||||
type Service struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type Service_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *Service) EXPECT() *Service_Expecter {
|
||||
return &Service_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// Algo provides a mock function for the type Service
|
||||
func (_mock *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
|
||||
ret := _mock.Called(ctx, algorithm)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Algo")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Algorithm) error); ok {
|
||||
r0 = returnFunc(ctx, algorithm)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// Service_Algo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Algo'
|
||||
type Service_Algo_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Algo is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - algorithm agent.Algorithm
|
||||
func (_e *Service_Expecter) Algo(ctx interface{}, algorithm interface{}) *Service_Algo_Call {
|
||||
return &Service_Algo_Call{Call: _e.mock.On("Algo", ctx, algorithm)}
|
||||
}
|
||||
|
||||
func (_c *Service_Algo_Call) Run(run func(ctx context.Context, algorithm agent.Algorithm)) *Service_Algo_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
var arg1 agent.Algorithm
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].(agent.Algorithm)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Algo_Call) Return(err error) *Service_Algo_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Algo_Call) RunAndReturn(run func(ctx context.Context, algorithm agent.Algorithm) error) *Service_Algo_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Attestation provides a mock function for the type Service
|
||||
func (_mock *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
|
||||
ret := _mock.Called(ctx, reportData, nonce, attType)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Attestation")
|
||||
}
|
||||
|
||||
var r0 []byte
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)); ok {
|
||||
return returnFunc(ctx, reportData, nonce, attType)
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) []byte); ok {
|
||||
r0 = returnFunc(ctx, reportData, nonce, attType)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).([]byte)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) error); ok {
|
||||
r1 = returnFunc(ctx, reportData, nonce, attType)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// Service_Attestation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Attestation'
|
||||
type Service_Attestation_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Attestation is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - reportData [64]byte
|
||||
// - nonce [32]byte
|
||||
// - attType attestation.PlatformType
|
||||
func (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{}, nonce interface{}, attType interface{}) *Service_Attestation_Call {
|
||||
return &Service_Attestation_Call{Call: _e.mock.On("Attestation", ctx, reportData, nonce, attType)}
|
||||
}
|
||||
|
||||
func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType)) *Service_Attestation_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
var arg1 [64]byte
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].([64]byte)
|
||||
}
|
||||
var arg2 [32]byte
|
||||
if args[2] != nil {
|
||||
arg2 = args[2].([32]byte)
|
||||
}
|
||||
var arg3 attestation.PlatformType
|
||||
if args[3] != nil {
|
||||
arg3 = args[3].(attestation.PlatformType)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
arg2,
|
||||
arg3,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Attestation_Call) Return(bytes []byte, err error) *Service_Attestation_Call {
|
||||
_c.Call.Return(bytes, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Attestation_Call) RunAndReturn(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error)) *Service_Attestation_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// AzureAttestationToken provides a mock function for the type Service
|
||||
func (_mock *Service) AzureAttestationToken(ctx context.Context, nonce [32]byte) ([]byte, error) {
|
||||
ret := _mock.Called(ctx, nonce)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for AzureAttestationToken")
|
||||
}
|
||||
|
||||
var r0 []byte
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, [32]byte) ([]byte, error)); ok {
|
||||
return returnFunc(ctx, nonce)
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, [32]byte) []byte); ok {
|
||||
r0 = returnFunc(ctx, nonce)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).([]byte)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func(context.Context, [32]byte) error); ok {
|
||||
r1 = returnFunc(ctx, nonce)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// Service_AzureAttestationToken_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AzureAttestationToken'
|
||||
type Service_AzureAttestationToken_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// AzureAttestationToken is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - nonce [32]byte
|
||||
func (_e *Service_Expecter) AzureAttestationToken(ctx interface{}, nonce interface{}) *Service_AzureAttestationToken_Call {
|
||||
return &Service_AzureAttestationToken_Call{Call: _e.mock.On("AzureAttestationToken", ctx, nonce)}
|
||||
}
|
||||
|
||||
func (_c *Service_AzureAttestationToken_Call) Run(run func(ctx context.Context, nonce [32]byte)) *Service_AzureAttestationToken_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
var arg1 [32]byte
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].([32]byte)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_AzureAttestationToken_Call) Return(bytes []byte, err error) *Service_AzureAttestationToken_Call {
|
||||
_c.Call.Return(bytes, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_AzureAttestationToken_Call) RunAndReturn(run func(ctx context.Context, nonce [32]byte) ([]byte, error)) *Service_AzureAttestationToken_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Data provides a mock function for the type Service
|
||||
func (_mock *Service) Data(ctx context.Context, dataset agent.Dataset) error {
|
||||
ret := _mock.Called(ctx, dataset)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Data")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Dataset) error); ok {
|
||||
r0 = returnFunc(ctx, dataset)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// Service_Data_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Data'
|
||||
type Service_Data_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Data is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
// - dataset agent.Dataset
|
||||
func (_e *Service_Expecter) Data(ctx interface{}, dataset interface{}) *Service_Data_Call {
|
||||
return &Service_Data_Call{Call: _e.mock.On("Data", ctx, dataset)}
|
||||
}
|
||||
|
||||
func (_c *Service_Data_Call) Run(run func(ctx context.Context, dataset agent.Dataset)) *Service_Data_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
var arg1 agent.Dataset
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].(agent.Dataset)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Data_Call) Return(err error) *Service_Data_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Data_Call) RunAndReturn(run func(ctx context.Context, dataset agent.Dataset) error) *Service_Data_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// IMAMeasurements provides a mock function for the type Service
|
||||
func (_mock *Service) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
|
||||
ret := _mock.Called(ctx)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for IMAMeasurements")
|
||||
}
|
||||
|
||||
var r0 []byte
|
||||
var r1 []byte
|
||||
var r2 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context) ([]byte, []byte, error)); ok {
|
||||
return returnFunc(ctx)
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context) []byte); ok {
|
||||
r0 = returnFunc(ctx)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).([]byte)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func(context.Context) []byte); ok {
|
||||
r1 = returnFunc(ctx)
|
||||
} else {
|
||||
if ret.Get(1) != nil {
|
||||
r1 = ret.Get(1).([]byte)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(2).(func(context.Context) error); ok {
|
||||
r2 = returnFunc(ctx)
|
||||
} else {
|
||||
r2 = ret.Error(2)
|
||||
}
|
||||
return r0, r1, r2
|
||||
}
|
||||
|
||||
// Service_IMAMeasurements_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IMAMeasurements'
|
||||
type Service_IMAMeasurements_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// IMAMeasurements is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
func (_e *Service_Expecter) IMAMeasurements(ctx interface{}) *Service_IMAMeasurements_Call {
|
||||
return &Service_IMAMeasurements_Call{Call: _e.mock.On("IMAMeasurements", ctx)}
|
||||
}
|
||||
|
||||
func (_c *Service_IMAMeasurements_Call) Run(run func(ctx context.Context)) *Service_IMAMeasurements_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_IMAMeasurements_Call) Return(bytes []byte, bytes1 []byte, err error) *Service_IMAMeasurements_Call {
|
||||
_c.Call.Return(bytes, bytes1, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_IMAMeasurements_Call) RunAndReturn(run func(ctx context.Context) ([]byte, []byte, error)) *Service_IMAMeasurements_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// InitComputation provides a mock function for the type Service
|
||||
func (_mock *Service) InitComputation(ctx context.Context, cmp agent.Computation) error {
|
||||
ret := _mock.Called(ctx, cmp)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for InitComputation")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Computation) error); ok {
|
||||
r0 = returnFunc(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) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
var arg1 agent.Computation
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].(agent.Computation)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_InitComputation_Call) Return(err error) *Service_InitComputation_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_InitComputation_Call) RunAndReturn(run func(ctx context.Context, cmp agent.Computation) error) *Service_InitComputation_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Result provides a mock function for the type Service
|
||||
func (_mock *Service) Result(ctx context.Context) ([]byte, error) {
|
||||
ret := _mock.Called(ctx)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Result")
|
||||
}
|
||||
|
||||
var r0 []byte
|
||||
var r1 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
|
||||
return returnFunc(ctx)
|
||||
}
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context) []byte); ok {
|
||||
r0 = returnFunc(ctx)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).([]byte)
|
||||
}
|
||||
}
|
||||
if returnFunc, ok := ret.Get(1).(func(context.Context) error); ok {
|
||||
r1 = returnFunc(ctx)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// Service_Result_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Result'
|
||||
type Service_Result_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Result is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
func (_e *Service_Expecter) Result(ctx interface{}) *Service_Result_Call {
|
||||
return &Service_Result_Call{Call: _e.mock.On("Result", ctx)}
|
||||
}
|
||||
|
||||
func (_c *Service_Result_Call) Run(run func(ctx context.Context)) *Service_Result_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Result_Call) Return(bytes []byte, err error) *Service_Result_Call {
|
||||
_c.Call.Return(bytes, err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_Result_Call) RunAndReturn(run func(ctx context.Context) ([]byte, error)) *Service_Result_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// State provides a mock function for the type Service
|
||||
func (_mock *Service) State() string {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for State")
|
||||
}
|
||||
|
||||
var r0 string
|
||||
if returnFunc, ok := ret.Get(0).(func() string); ok {
|
||||
r0 = returnFunc()
|
||||
} 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(s string) *Service_State_Call {
|
||||
_c.Call.Return(s)
|
||||
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 for the type Service
|
||||
func (_mock *Service) StopComputation(ctx context.Context) error {
|
||||
ret := _mock.Called(ctx)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for StopComputation")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context) error); ok {
|
||||
r0 = returnFunc(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) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_StopComputation_Call) Return(err error) *Service_StopComputation_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *Service_StopComputation_Call) RunAndReturn(run func(ctx context.Context) error) *Service_StopComputation_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -1,200 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/pkg/resource"
|
||||
)
|
||||
|
||||
type MockDownloader struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockDownloader) Download(ctx context.Context, url string, destPath string) error {
|
||||
args := m.Called(ctx, url, destPath)
|
||||
if args.Error(0) == nil {
|
||||
// Simulate writing to destPath if it's a success
|
||||
content := "mock content"
|
||||
if len(args) > 1 {
|
||||
if c, ok := args.Get(1).(string); ok {
|
||||
content = c
|
||||
}
|
||||
}
|
||||
_ = os.MkdirAll(filepath.Dir(destPath), 0o755)
|
||||
_ = os.WriteFile(destPath, []byte(content), 0o644)
|
||||
}
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockDownloader) Type() string {
|
||||
return m.Called().String(0)
|
||||
}
|
||||
|
||||
func TestDownloadAndDecryptGenericResource(t *testing.T) {
|
||||
registry := resource.NewRegistry()
|
||||
mockDownloader := new(MockDownloader)
|
||||
mockDownloader.On("Type").Return(resource.SourceTypeHTTP)
|
||||
registry.Register(mockDownloader)
|
||||
|
||||
attestationClient := new(MockAttestationClient)
|
||||
attestationClient.On("GetKbsToken", mock.Anything).Return([]byte("mockToken"), nil).Maybe()
|
||||
|
||||
svc := &agentService{
|
||||
logger: slog.Default(),
|
||||
resourceRegistry: registry,
|
||||
attestationClient: attestationClient,
|
||||
computation: Computation{
|
||||
Algorithm: &Algorithm{
|
||||
KBS: &KBSConfig{
|
||||
Enabled: true,
|
||||
URL: "http://mock-kbs",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("Successful download without encryption", func(t *testing.T) {
|
||||
source := &ResourceSource{
|
||||
URL: "http://example.com/resource",
|
||||
}
|
||||
destPath := filepath.Join(os.TempDir(), "cocos-resources", "algo", "resource")
|
||||
mockDownloader.On("Download", ctx, source.URL, destPath).Return(nil, "some data").Once()
|
||||
|
||||
res, err := svc.downloadAndDecryptGenericResource(ctx, source, resource.SourceTypeHTTP, "", "algo")
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, []byte("some data"), res.Data)
|
||||
mockDownloader.AssertExpectations(t)
|
||||
})
|
||||
|
||||
t.Run("Successful download with encryption", func(t *testing.T) {
|
||||
key := make([]byte, 32)
|
||||
_, _ = io.ReadFull(rand.Reader, key)
|
||||
|
||||
plaintext := []byte("secret data")
|
||||
block, _ := aes.NewCipher(key)
|
||||
gcm, _ := cipher.NewGCM(block)
|
||||
nonce := make([]byte, gcm.NonceSize())
|
||||
_, _ = io.ReadFull(rand.Reader, nonce)
|
||||
ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)
|
||||
|
||||
// Mock KBS
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(key)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
svc.computation.Algorithm.KBS.URL = ts.URL
|
||||
|
||||
source := &ResourceSource{
|
||||
URL: "http://example.com/encrypted",
|
||||
Encrypted: true,
|
||||
KBSResourcePath: "keys/1",
|
||||
}
|
||||
destPath := filepath.Join(os.TempDir(), "cocos-resources", "data", "encrypted")
|
||||
mockDownloader.On("Download", ctx, source.URL, destPath).Return(nil, string(ciphertext)).Once()
|
||||
|
||||
res, err := svc.downloadAndDecryptGenericResource(ctx, source, resource.SourceTypeHTTP, svc.computation.Algorithm.KBS.URL, "data")
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, plaintext, res.Data)
|
||||
mockDownloader.AssertExpectations(t)
|
||||
})
|
||||
|
||||
t.Run("Registry not initialized", func(t *testing.T) {
|
||||
badSvc := &agentService{logger: slog.Default()}
|
||||
_, err := badSvc.downloadAndDecryptGenericResource(ctx, &ResourceSource{}, "http", "", "algo")
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "resource registry not initialized")
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetKeyFromKBS(t *testing.T) {
|
||||
attestationClient := new(MockAttestationClient)
|
||||
attestationClient.On("GetKbsToken", mock.Anything).Return([]byte("mockToken"), nil).Maybe()
|
||||
|
||||
svc := &agentService{
|
||||
logger: slog.Default(),
|
||||
attestationClient: attestationClient,
|
||||
computation: Computation{
|
||||
Algorithm: &Algorithm{
|
||||
KBS: &KBSConfig{
|
||||
Enabled: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("KBS disabled", func(t *testing.T) {
|
||||
svc.computation.Algorithm.KBS.Enabled = false
|
||||
_, err := svc.getKeyFromKBS(ctx, "", "path")
|
||||
assert.Error(t, err)
|
||||
})
|
||||
|
||||
t.Run("Successful fetch", func(t *testing.T) {
|
||||
svc.computation.Algorithm.KBS.Enabled = true
|
||||
key := []byte("this is a 32-byte key!!!!!!!!!!!")
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Contains(t, r.URL.Path, "resource/path")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(key)
|
||||
}))
|
||||
defer ts.Close()
|
||||
svc.computation.Algorithm.KBS.URL = ts.URL
|
||||
|
||||
fetched, err := svc.getKeyFromKBS(ctx, svc.computation.Algorithm.KBS.URL, "path")
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, key, fetched)
|
||||
})
|
||||
|
||||
t.Run("KBS error", func(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
defer ts.Close()
|
||||
svc.computation.Algorithm.KBS.URL = ts.URL
|
||||
|
||||
_, err := svc.getKeyFromKBS(ctx, svc.computation.Algorithm.KBS.URL, "path")
|
||||
assert.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestInferSourceTypeDetailed(t *testing.T) {
|
||||
tests := []struct {
|
||||
url string
|
||||
expected string
|
||||
}{
|
||||
{"s3://bucket/key", resource.SourceTypeS3},
|
||||
{"gs://bucket/key", resource.SourceTypeGCS},
|
||||
{"https://example.com/file", resource.SourceTypeHTTPS},
|
||||
{"http://example.com/file", resource.SourceTypeHTTP},
|
||||
{"docker://ubuntu", resource.SourceTypeOCIImage},
|
||||
{"oci:/path/to/dir", resource.SourceTypeOCIImage},
|
||||
{"ubuntu:latest", resource.SourceTypeOCIImage},
|
||||
{"myregistry.io/myimage:tag", resource.SourceTypeOCIImage},
|
||||
{"invalid-url-no-slash", ""},
|
||||
{"", ""},
|
||||
{"ftp://server/file", ""},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
assert.Equal(t, tt.expected, inferSourceType(tt.url), "URL: %s", tt.url)
|
||||
}
|
||||
}
|
||||
@@ -1,38 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package events
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
logpb "github.com/ultravioletrs/cocos/agent/log"
|
||||
logclient "github.com/ultravioletrs/cocos/pkg/clients/grpc/log"
|
||||
)
|
||||
|
||||
type adapter struct {
|
||||
client logclient.Client
|
||||
svc string
|
||||
}
|
||||
|
||||
func NewAdapter(client logclient.Client, svc string) events.Service {
|
||||
return &adapter{
|
||||
client: client,
|
||||
svc: svc,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *adapter) SendEvent(cmpID, event, status string, details json.RawMessage) {
|
||||
err := a.client.SendEvent(context.Background(), &logpb.EventEntry{
|
||||
EventType: event,
|
||||
ComputationId: cmpID,
|
||||
Details: details,
|
||||
Originator: a.svc,
|
||||
Status: status,
|
||||
})
|
||||
if err != nil {
|
||||
slog.Error("failed to send event to log-forwarder", "error", err)
|
||||
}
|
||||
}
|
||||
@@ -1,138 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package events
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
logpb "github.com/ultravioletrs/cocos/agent/log"
|
||||
)
|
||||
|
||||
const testServiceName = "test-service"
|
||||
|
||||
// mockLogClient is a mock implementation of the log client.
|
||||
type mockLogClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *mockLogClient) SendLog(ctx context.Context, entry *logpb.LogEntry) error {
|
||||
args := m.Called(ctx, entry)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *mockLogClient) SendEvent(ctx context.Context, entry *logpb.EventEntry) error {
|
||||
args := m.Called(ctx, entry)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *mockLogClient) Close() error {
|
||||
args := m.Called()
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
// TestNewAdapter tests creating a new adapter.
|
||||
func TestNewAdapter(t *testing.T) {
|
||||
mockClient := new(mockLogClient)
|
||||
svc := testServiceName
|
||||
|
||||
adapter := NewAdapter(mockClient, svc)
|
||||
|
||||
assert.NotNil(t, adapter)
|
||||
}
|
||||
|
||||
// TestSendEvent tests sending an event successfully.
|
||||
func TestSendEvent(t *testing.T) {
|
||||
mockClient := new(mockLogClient)
|
||||
svc := testServiceName
|
||||
adapter := NewAdapter(mockClient, svc)
|
||||
|
||||
cmpID := "test-computation-id"
|
||||
event := "computation.started"
|
||||
status := "success"
|
||||
details := json.RawMessage(`{"key": "value"}`)
|
||||
|
||||
expectedEntry := &logpb.EventEntry{
|
||||
EventType: event,
|
||||
ComputationId: cmpID,
|
||||
Details: details,
|
||||
Originator: svc,
|
||||
Status: status,
|
||||
}
|
||||
|
||||
mockClient.On("SendEvent", mock.Anything, expectedEntry).Return(nil)
|
||||
|
||||
adapter.SendEvent(cmpID, event, status, details)
|
||||
|
||||
mockClient.AssertExpectations(t)
|
||||
mockClient.AssertCalled(t, "SendEvent", mock.Anything, expectedEntry)
|
||||
}
|
||||
|
||||
// TestSendEventWithError tests sending an event when client returns an error.
|
||||
func TestSendEventWithError(t *testing.T) {
|
||||
mockClient := new(mockLogClient)
|
||||
svc := testServiceName
|
||||
adapter := NewAdapter(mockClient, svc)
|
||||
|
||||
cmpID := "test-computation-id"
|
||||
event := "computation.failed"
|
||||
status := "error"
|
||||
details := json.RawMessage(`{"error": "something went wrong"}`)
|
||||
|
||||
mockClient.On("SendEvent", mock.Anything, mock.Anything).Return(assert.AnError)
|
||||
|
||||
// This should not panic even when error occurs
|
||||
adapter.SendEvent(cmpID, event, status, details)
|
||||
|
||||
mockClient.AssertExpectations(t)
|
||||
mockClient.AssertCalled(t, "SendEvent", mock.Anything, mock.Anything)
|
||||
}
|
||||
|
||||
// TestSendEventWithNilDetails tests sending an event with nil details.
|
||||
func TestSendEventWithNilDetails(t *testing.T) {
|
||||
mockClient := new(mockLogClient)
|
||||
svc := "runner-service"
|
||||
adapter := NewAdapter(mockClient, svc)
|
||||
|
||||
cmpID := "comp-123"
|
||||
event := "test.event"
|
||||
status := "pending"
|
||||
|
||||
expectedEntry := &logpb.EventEntry{
|
||||
EventType: event,
|
||||
ComputationId: cmpID,
|
||||
Details: nil,
|
||||
Originator: svc,
|
||||
Status: status,
|
||||
}
|
||||
|
||||
mockClient.On("SendEvent", mock.Anything, expectedEntry).Return(nil)
|
||||
|
||||
adapter.SendEvent(cmpID, event, status, nil)
|
||||
|
||||
mockClient.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// TestSendEventWithEmptyStrings tests sending an event with empty strings.
|
||||
func TestSendEventWithEmptyStrings(t *testing.T) {
|
||||
mockClient := new(mockLogClient)
|
||||
svc := testServiceName
|
||||
adapter := NewAdapter(mockClient, svc)
|
||||
|
||||
expectedEntry := &logpb.EventEntry{
|
||||
EventType: "",
|
||||
ComputationId: "",
|
||||
Details: nil,
|
||||
Originator: svc,
|
||||
Status: "",
|
||||
}
|
||||
|
||||
mockClient.On("SendEvent", mock.Anything, expectedEntry).Return(nil)
|
||||
|
||||
adapter.SendEvent("", "", "", nil)
|
||||
|
||||
mockClient.AssertExpectations(t)
|
||||
}
|
||||
@@ -1,341 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.36.11
|
||||
// protoc v7.35.0
|
||||
// source: agent/runner/runner.proto
|
||||
|
||||
package runner
|
||||
|
||||
import (
|
||||
protoreflect "google.golang.org/protobuf/reflect/protoreflect"
|
||||
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
|
||||
emptypb "google.golang.org/protobuf/types/known/emptypb"
|
||||
reflect "reflect"
|
||||
sync "sync"
|
||||
unsafe "unsafe"
|
||||
)
|
||||
|
||||
const (
|
||||
// Verify that this generated code is sufficiently up-to-date.
|
||||
_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
|
||||
// Verify that runtime/protoimpl is sufficiently up-to-date.
|
||||
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
|
||||
)
|
||||
|
||||
type RunRequest struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
ComputationId string `protobuf:"bytes,1,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
|
||||
AlgoType string `protobuf:"bytes,2,opt,name=algo_type,json=algoType,proto3" json:"algo_type,omitempty"` // "binary", "python", "wasm", "docker"
|
||||
Algorithm []byte `protobuf:"bytes,3,opt,name=algorithm,proto3" json:"algorithm,omitempty"` // The algorithm binary/script content
|
||||
Requirements []byte `protobuf:"bytes,4,opt,name=requirements,proto3" json:"requirements,omitempty"` // Python requirements.txt content
|
||||
Args []string `protobuf:"bytes,5,rep,name=args,proto3" json:"args,omitempty"`
|
||||
Datasets []*Dataset `protobuf:"bytes,6,rep,name=datasets,proto3" json:"datasets,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *RunRequest) Reset() {
|
||||
*x = RunRequest{}
|
||||
mi := &file_agent_runner_runner_proto_msgTypes[0]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *RunRequest) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*RunRequest) ProtoMessage() {}
|
||||
|
||||
func (x *RunRequest) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_runner_runner_proto_msgTypes[0]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use RunRequest.ProtoReflect.Descriptor instead.
|
||||
func (*RunRequest) Descriptor() ([]byte, []int) {
|
||||
return file_agent_runner_runner_proto_rawDescGZIP(), []int{0}
|
||||
}
|
||||
|
||||
func (x *RunRequest) GetComputationId() string {
|
||||
if x != nil {
|
||||
return x.ComputationId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunRequest) GetAlgoType() string {
|
||||
if x != nil {
|
||||
return x.AlgoType
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunRequest) GetAlgorithm() []byte {
|
||||
if x != nil {
|
||||
return x.Algorithm
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *RunRequest) GetRequirements() []byte {
|
||||
if x != nil {
|
||||
return x.Requirements
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *RunRequest) GetArgs() []string {
|
||||
if x != nil {
|
||||
return x.Args
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *RunRequest) GetDatasets() []*Dataset {
|
||||
if x != nil {
|
||||
return x.Datasets
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type Dataset struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Filename string `protobuf:"bytes,1,opt,name=filename,proto3" json:"filename,omitempty"`
|
||||
Hash []byte `protobuf:"bytes,2,opt,name=hash,proto3" json:"hash,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Dataset) Reset() {
|
||||
*x = Dataset{}
|
||||
mi := &file_agent_runner_runner_proto_msgTypes[1]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *Dataset) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*Dataset) ProtoMessage() {}
|
||||
|
||||
func (x *Dataset) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_runner_runner_proto_msgTypes[1]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use Dataset.ProtoReflect.Descriptor instead.
|
||||
func (*Dataset) Descriptor() ([]byte, []int) {
|
||||
return file_agent_runner_runner_proto_rawDescGZIP(), []int{1}
|
||||
}
|
||||
|
||||
func (x *Dataset) GetFilename() string {
|
||||
if x != nil {
|
||||
return x.Filename
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *Dataset) GetHash() []byte {
|
||||
if x != nil {
|
||||
return x.Hash
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type RunResponse struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
ComputationId string `protobuf:"bytes,1,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
|
||||
Error string `protobuf:"bytes,2,opt,name=error,proto3" json:"error,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *RunResponse) Reset() {
|
||||
*x = RunResponse{}
|
||||
mi := &file_agent_runner_runner_proto_msgTypes[2]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *RunResponse) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*RunResponse) ProtoMessage() {}
|
||||
|
||||
func (x *RunResponse) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_runner_runner_proto_msgTypes[2]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use RunResponse.ProtoReflect.Descriptor instead.
|
||||
func (*RunResponse) Descriptor() ([]byte, []int) {
|
||||
return file_agent_runner_runner_proto_rawDescGZIP(), []int{2}
|
||||
}
|
||||
|
||||
func (x *RunResponse) GetComputationId() string {
|
||||
if x != nil {
|
||||
return x.ComputationId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunResponse) GetError() string {
|
||||
if x != nil {
|
||||
return x.Error
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type StopRequest struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
ComputationId string `protobuf:"bytes,1,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *StopRequest) Reset() {
|
||||
*x = StopRequest{}
|
||||
mi := &file_agent_runner_runner_proto_msgTypes[3]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *StopRequest) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*StopRequest) ProtoMessage() {}
|
||||
|
||||
func (x *StopRequest) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_agent_runner_runner_proto_msgTypes[3]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use StopRequest.ProtoReflect.Descriptor instead.
|
||||
func (*StopRequest) Descriptor() ([]byte, []int) {
|
||||
return file_agent_runner_runner_proto_rawDescGZIP(), []int{3}
|
||||
}
|
||||
|
||||
func (x *StopRequest) GetComputationId() string {
|
||||
if x != nil {
|
||||
return x.ComputationId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
var File_agent_runner_runner_proto protoreflect.FileDescriptor
|
||||
|
||||
const file_agent_runner_runner_proto_rawDesc = "" +
|
||||
"\n" +
|
||||
"\x19agent/runner/runner.proto\x12\x06runner\x1a\x1bgoogle/protobuf/empty.proto\"\xd3\x01\n" +
|
||||
"\n" +
|
||||
"RunRequest\x12%\n" +
|
||||
"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\x12\x1b\n" +
|
||||
"\talgo_type\x18\x02 \x01(\tR\balgoType\x12\x1c\n" +
|
||||
"\talgorithm\x18\x03 \x01(\fR\talgorithm\x12\"\n" +
|
||||
"\frequirements\x18\x04 \x01(\fR\frequirements\x12\x12\n" +
|
||||
"\x04args\x18\x05 \x03(\tR\x04args\x12+\n" +
|
||||
"\bdatasets\x18\x06 \x03(\v2\x0f.runner.DatasetR\bdatasets\"9\n" +
|
||||
"\aDataset\x12\x1a\n" +
|
||||
"\bfilename\x18\x01 \x01(\tR\bfilename\x12\x12\n" +
|
||||
"\x04hash\x18\x02 \x01(\fR\x04hash\"J\n" +
|
||||
"\vRunResponse\x12%\n" +
|
||||
"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\x12\x14\n" +
|
||||
"\x05error\x18\x02 \x01(\tR\x05error\"4\n" +
|
||||
"\vStopRequest\x12%\n" +
|
||||
"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId2x\n" +
|
||||
"\x11ComputationRunner\x12.\n" +
|
||||
"\x03Run\x12\x12.runner.RunRequest\x1a\x13.runner.RunResponse\x123\n" +
|
||||
"\x04Stop\x12\x13.runner.StopRequest\x1a\x16.google.protobuf.EmptyB\n" +
|
||||
"Z\b./runnerb\x06proto3"
|
||||
|
||||
var (
|
||||
file_agent_runner_runner_proto_rawDescOnce sync.Once
|
||||
file_agent_runner_runner_proto_rawDescData []byte
|
||||
)
|
||||
|
||||
func file_agent_runner_runner_proto_rawDescGZIP() []byte {
|
||||
file_agent_runner_runner_proto_rawDescOnce.Do(func() {
|
||||
file_agent_runner_runner_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_runner_runner_proto_rawDesc), len(file_agent_runner_runner_proto_rawDesc)))
|
||||
})
|
||||
return file_agent_runner_runner_proto_rawDescData
|
||||
}
|
||||
|
||||
var file_agent_runner_runner_proto_msgTypes = make([]protoimpl.MessageInfo, 4)
|
||||
var file_agent_runner_runner_proto_goTypes = []any{
|
||||
(*RunRequest)(nil), // 0: runner.RunRequest
|
||||
(*Dataset)(nil), // 1: runner.Dataset
|
||||
(*RunResponse)(nil), // 2: runner.RunResponse
|
||||
(*StopRequest)(nil), // 3: runner.StopRequest
|
||||
(*emptypb.Empty)(nil), // 4: google.protobuf.Empty
|
||||
}
|
||||
var file_agent_runner_runner_proto_depIdxs = []int32{
|
||||
1, // 0: runner.RunRequest.datasets:type_name -> runner.Dataset
|
||||
0, // 1: runner.ComputationRunner.Run:input_type -> runner.RunRequest
|
||||
3, // 2: runner.ComputationRunner.Stop:input_type -> runner.StopRequest
|
||||
2, // 3: runner.ComputationRunner.Run:output_type -> runner.RunResponse
|
||||
4, // 4: runner.ComputationRunner.Stop:output_type -> google.protobuf.Empty
|
||||
3, // [3:5] is the sub-list for method output_type
|
||||
1, // [1:3] is the sub-list for method input_type
|
||||
1, // [1:1] is the sub-list for extension type_name
|
||||
1, // [1:1] is the sub-list for extension extendee
|
||||
0, // [0:1] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_agent_runner_runner_proto_init() }
|
||||
func file_agent_runner_runner_proto_init() {
|
||||
if File_agent_runner_runner_proto != nil {
|
||||
return
|
||||
}
|
||||
type x struct{}
|
||||
out := protoimpl.TypeBuilder{
|
||||
File: protoimpl.DescBuilder{
|
||||
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
|
||||
RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_runner_runner_proto_rawDesc), len(file_agent_runner_runner_proto_rawDesc)),
|
||||
NumEnums: 0,
|
||||
NumMessages: 4,
|
||||
NumExtensions: 0,
|
||||
NumServices: 1,
|
||||
},
|
||||
GoTypes: file_agent_runner_runner_proto_goTypes,
|
||||
DependencyIndexes: file_agent_runner_runner_proto_depIdxs,
|
||||
MessageInfos: file_agent_runner_runner_proto_msgTypes,
|
||||
}.Build()
|
||||
File_agent_runner_runner_proto = out.File
|
||||
file_agent_runner_runner_proto_goTypes = nil
|
||||
file_agent_runner_runner_proto_depIdxs = nil
|
||||
}
|
||||
@@ -1,38 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package runner;
|
||||
|
||||
option go_package = "./runner";
|
||||
|
||||
import "google/protobuf/empty.proto";
|
||||
|
||||
service ComputationRunner {
|
||||
rpc Run(RunRequest) returns (RunResponse);
|
||||
rpc Stop(StopRequest) returns (google.protobuf.Empty);
|
||||
}
|
||||
|
||||
message RunRequest {
|
||||
string computation_id = 1;
|
||||
string algo_type = 2; // "binary", "python", "wasm", "docker"
|
||||
bytes algorithm = 3; // The algorithm binary/script content
|
||||
bytes requirements = 4; // Python requirements.txt content
|
||||
repeated string args = 5;
|
||||
repeated Dataset datasets = 6;
|
||||
}
|
||||
|
||||
message Dataset {
|
||||
string filename = 1;
|
||||
bytes hash = 2;
|
||||
}
|
||||
|
||||
message RunResponse {
|
||||
string computation_id = 1;
|
||||
string error = 2;
|
||||
}
|
||||
|
||||
message StopRequest {
|
||||
string computation_id = 1;
|
||||
}
|
||||
@@ -1,163 +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.6.0
|
||||
// - protoc v7.35.0
|
||||
// source: agent/runner/runner.proto
|
||||
|
||||
package runner
|
||||
|
||||
import (
|
||||
context "context"
|
||||
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
|
||||
// is compatible with the grpc package it is being compiled against.
|
||||
// Requires gRPC-Go v1.64.0 or later.
|
||||
const _ = grpc.SupportPackageIsVersion9
|
||||
|
||||
const (
|
||||
ComputationRunner_Run_FullMethodName = "/runner.ComputationRunner/Run"
|
||||
ComputationRunner_Stop_FullMethodName = "/runner.ComputationRunner/Stop"
|
||||
)
|
||||
|
||||
// ComputationRunnerClient is the client API for ComputationRunner 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 ComputationRunnerClient interface {
|
||||
Run(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (*RunResponse, error)
|
||||
Stop(ctx context.Context, in *StopRequest, opts ...grpc.CallOption) (*emptypb.Empty, error)
|
||||
}
|
||||
|
||||
type computationRunnerClient struct {
|
||||
cc grpc.ClientConnInterface
|
||||
}
|
||||
|
||||
func NewComputationRunnerClient(cc grpc.ClientConnInterface) ComputationRunnerClient {
|
||||
return &computationRunnerClient{cc}
|
||||
}
|
||||
|
||||
func (c *computationRunnerClient) Run(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (*RunResponse, error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(RunResponse)
|
||||
err := c.cc.Invoke(ctx, ComputationRunner_Run_FullMethodName, in, out, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (c *computationRunnerClient) Stop(ctx context.Context, in *StopRequest, opts ...grpc.CallOption) (*emptypb.Empty, error) {
|
||||
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
|
||||
out := new(emptypb.Empty)
|
||||
err := c.cc.Invoke(ctx, ComputationRunner_Stop_FullMethodName, in, out, cOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ComputationRunnerServer is the server API for ComputationRunner service.
|
||||
// All implementations must embed UnimplementedComputationRunnerServer
|
||||
// for forward compatibility.
|
||||
type ComputationRunnerServer interface {
|
||||
Run(context.Context, *RunRequest) (*RunResponse, error)
|
||||
Stop(context.Context, *StopRequest) (*emptypb.Empty, error)
|
||||
mustEmbedUnimplementedComputationRunnerServer()
|
||||
}
|
||||
|
||||
// UnimplementedComputationRunnerServer 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 UnimplementedComputationRunnerServer struct{}
|
||||
|
||||
func (UnimplementedComputationRunnerServer) Run(context.Context, *RunRequest) (*RunResponse, error) {
|
||||
return nil, status.Error(codes.Unimplemented, "method Run not implemented")
|
||||
}
|
||||
func (UnimplementedComputationRunnerServer) Stop(context.Context, *StopRequest) (*emptypb.Empty, error) {
|
||||
return nil, status.Error(codes.Unimplemented, "method Stop not implemented")
|
||||
}
|
||||
func (UnimplementedComputationRunnerServer) mustEmbedUnimplementedComputationRunnerServer() {}
|
||||
func (UnimplementedComputationRunnerServer) testEmbeddedByValue() {}
|
||||
|
||||
// UnsafeComputationRunnerServer may be embedded to opt out of forward compatibility for this service.
|
||||
// Use of this interface is not recommended, as added methods to ComputationRunnerServer will
|
||||
// result in compilation errors.
|
||||
type UnsafeComputationRunnerServer interface {
|
||||
mustEmbedUnimplementedComputationRunnerServer()
|
||||
}
|
||||
|
||||
func RegisterComputationRunnerServer(s grpc.ServiceRegistrar, srv ComputationRunnerServer) {
|
||||
// If the following call panics, it indicates UnimplementedComputationRunnerServer 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(&ComputationRunner_ServiceDesc, srv)
|
||||
}
|
||||
|
||||
func _ComputationRunner_Run_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(RunRequest)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(ComputationRunnerServer).Run(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: ComputationRunner_Run_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(ComputationRunnerServer).Run(ctx, req.(*RunRequest))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
func _ComputationRunner_Stop_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
|
||||
in := new(StopRequest)
|
||||
if err := dec(in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if interceptor == nil {
|
||||
return srv.(ComputationRunnerServer).Stop(ctx, in)
|
||||
}
|
||||
info := &grpc.UnaryServerInfo{
|
||||
Server: srv,
|
||||
FullMethod: ComputationRunner_Stop_FullMethodName,
|
||||
}
|
||||
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
|
||||
return srv.(ComputationRunnerServer).Stop(ctx, req.(*StopRequest))
|
||||
}
|
||||
return interceptor(ctx, in, info, handler)
|
||||
}
|
||||
|
||||
// ComputationRunner_ServiceDesc is the grpc.ServiceDesc for ComputationRunner service.
|
||||
// It's only intended for direct use with grpc.RegisterService,
|
||||
// and not to be introspected or modified (even as a copy)
|
||||
var ComputationRunner_ServiceDesc = grpc.ServiceDesc{
|
||||
ServiceName: "runner.ComputationRunner",
|
||||
HandlerType: (*ComputationRunnerServer)(nil),
|
||||
Methods: []grpc.MethodDesc{
|
||||
{
|
||||
MethodName: "Run",
|
||||
Handler: _ComputationRunner_Run_Handler,
|
||||
},
|
||||
{
|
||||
MethodName: "Stop",
|
||||
Handler: _ComputationRunner_Stop_Handler,
|
||||
},
|
||||
},
|
||||
Streams: []grpc.StreamDesc{},
|
||||
Metadata: "agent/runner/runner.proto",
|
||||
}
|
||||
@@ -1,151 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/binary"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/docker"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/python"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/wasm"
|
||||
"github.com/ultravioletrs/cocos/agent/events"
|
||||
pb "github.com/ultravioletrs/cocos/agent/runner"
|
||||
"google.golang.org/protobuf/types/known/emptypb"
|
||||
)
|
||||
|
||||
const (
|
||||
algoFilePermission = 0o700
|
||||
)
|
||||
|
||||
var _ pb.ComputationRunnerServer = (*RunnerService)(nil)
|
||||
|
||||
type RunnerService struct {
|
||||
pb.UnimplementedComputationRunnerServer
|
||||
logger *slog.Logger
|
||||
eventSvc events.Service
|
||||
currentAlgo algorithm.Algorithm
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func New(logger *slog.Logger, eventSvc events.Service) *RunnerService {
|
||||
return &RunnerService{
|
||||
logger: logger,
|
||||
eventSvc: eventSvc,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *RunnerService) Run(ctx context.Context, req *pb.RunRequest) (*pb.RunResponse, error) {
|
||||
s.mu.Lock()
|
||||
if s.currentAlgo != nil {
|
||||
s.mu.Unlock()
|
||||
return &pb.RunResponse{
|
||||
ComputationId: req.ComputationId,
|
||||
Error: "computation already running",
|
||||
}, nil
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
s.currentAlgo = nil
|
||||
s.mu.Unlock()
|
||||
}()
|
||||
|
||||
currentDir, err := os.Getwd()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting current directory: %v", err)
|
||||
}
|
||||
|
||||
// Write Algo File
|
||||
algoPath := filepath.Join(currentDir, "algo")
|
||||
f, err := os.Create(algoPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating algorithm file: %v", err)
|
||||
}
|
||||
if _, err := f.Write(req.Algorithm); err != nil {
|
||||
return nil, fmt.Errorf("error writing algorithm to file: %v", err)
|
||||
}
|
||||
if err := os.Chmod(algoPath, algoFilePermission); err != nil {
|
||||
return nil, fmt.Errorf("error changing file permissions: %v", err)
|
||||
}
|
||||
if err := f.Close(); err != nil {
|
||||
return nil, fmt.Errorf("error closing file: %v", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := os.Remove(algoPath); err != nil {
|
||||
s.logger.Warn("error removing algorithm file", "error", err)
|
||||
}
|
||||
}()
|
||||
|
||||
var algo algorithm.Algorithm
|
||||
|
||||
switch req.AlgoType {
|
||||
case string(algorithm.AlgoTypeBin):
|
||||
algo = binary.NewAlgorithm(s.logger, s.eventSvc, algoPath, req.Args, req.ComputationId)
|
||||
case string(algorithm.AlgoTypePython):
|
||||
var requirementsFile string
|
||||
if len(req.Requirements) > 0 {
|
||||
fr, err := os.CreateTemp("", "requirements.txt")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating requirments file: %v", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := os.Remove(fr.Name()); err != nil {
|
||||
s.logger.Warn("error removing requirements file", "error", err)
|
||||
}
|
||||
}()
|
||||
if _, err := fr.Write(req.Requirements); err != nil {
|
||||
return nil, fmt.Errorf("error writing requirements to file: %v", err)
|
||||
}
|
||||
if err := fr.Close(); err != nil {
|
||||
return nil, fmt.Errorf("error closing file: %v", err)
|
||||
}
|
||||
requirementsFile = fr.Name()
|
||||
}
|
||||
// Assuming default python runtime if not specified in request (proto doesn't have runtime field yet)
|
||||
// We can add it or assume.
|
||||
runtime := python.PyRuntime
|
||||
algo = python.NewAlgorithm(s.logger, s.eventSvc, runtime, requirementsFile, algoPath, req.Args, req.ComputationId)
|
||||
case string(algorithm.AlgoTypeWasm):
|
||||
algo = wasm.NewAlgorithm(s.logger, s.eventSvc, req.Args, algoPath, req.ComputationId)
|
||||
case string(algorithm.AlgoTypeDocker):
|
||||
algo = docker.NewAlgorithm(s.logger, s.eventSvc, algoPath, req.ComputationId)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported algorithm type: %s", req.AlgoType)
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
s.currentAlgo = algo
|
||||
s.mu.Unlock()
|
||||
|
||||
if err := algo.Run(); err != nil {
|
||||
s.logger.Error("computation failed", "error", err)
|
||||
return &pb.RunResponse{
|
||||
ComputationId: req.ComputationId,
|
||||
Error: err.Error(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
return &pb.RunResponse{
|
||||
ComputationId: req.ComputationId,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *RunnerService) Stop(ctx context.Context, req *pb.StopRequest) (*emptypb.Empty, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.currentAlgo != nil {
|
||||
if err := s.currentAlgo.Stop(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return &emptypb.Empty{}, nil
|
||||
}
|
||||
@@ -1,382 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
pb "github.com/ultravioletrs/cocos/agent/runner"
|
||||
)
|
||||
|
||||
// MockEventService is a mock implementation of events.Service.
|
||||
type MockEventService struct {
|
||||
events []interface{}
|
||||
}
|
||||
|
||||
func (m *MockEventService) SendEvent(cmpID, event, status string, details json.RawMessage) {
|
||||
m.events = append(m.events, map[string]interface{}{
|
||||
"cmpID": cmpID,
|
||||
"event": event,
|
||||
"status": status,
|
||||
"details": details,
|
||||
})
|
||||
}
|
||||
|
||||
// TestNewRunnerService tests the creation of a new runner service.
|
||||
func TestNewRunnerService(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
|
||||
rs := New(logger, eventSvc)
|
||||
require.NotNil(t, rs)
|
||||
assert.NotNil(t, rs.logger)
|
||||
assert.NotNil(t, rs.eventSvc)
|
||||
assert.Nil(t, rs.currentAlgo)
|
||||
}
|
||||
|
||||
// TestRunWithBinaryAlgorithm tests running a binary algorithm.
|
||||
func TestRunWithBinaryAlgorithm(t *testing.T) {
|
||||
origDir, _ := os.Getwd()
|
||||
tmpDir := t.TempDir()
|
||||
require.NoError(t, os.Chdir(tmpDir))
|
||||
defer func() { require.NoError(t, os.Chdir(origDir)) }()
|
||||
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-1",
|
||||
AlgoType: "bin",
|
||||
Algorithm: []byte("#!/bin/bash\necho 'test'"),
|
||||
Args: []string{"arg1", "arg2"},
|
||||
}
|
||||
|
||||
resp, err := rs.Run(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
assert.Empty(t, resp.Error)
|
||||
assert.Equal(t, "test-1", resp.ComputationId)
|
||||
}
|
||||
|
||||
// TestRunWithPythonAlgorithm tests running a Python algorithm.
|
||||
func TestRunWithPythonAlgorithm(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-python",
|
||||
AlgoType: "python",
|
||||
Algorithm: []byte("print('hello')"),
|
||||
Args: []string{},
|
||||
Requirements: []byte("numpy==2.2.0"),
|
||||
}
|
||||
|
||||
resp, err := rs.Run(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
assert.Empty(t, resp.Error)
|
||||
assert.Equal(t, "test-python", resp.ComputationId)
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove("algo")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRunWithPythonAlgorithmNoRequirements tests running Python without requirements.
|
||||
func TestRunWithPythonAlgorithmNoRequirements(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-python-noreq",
|
||||
AlgoType: "python",
|
||||
Algorithm: []byte("print('hello')"),
|
||||
Args: []string{},
|
||||
}
|
||||
|
||||
resp, err := rs.Run(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
assert.Empty(t, resp.Error)
|
||||
assert.Equal(t, "test-python-noreq", resp.ComputationId)
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove("algo")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRunWithWasmAlgorithm tests running a WASM algorithm.
|
||||
func TestRunWithWasmAlgorithm(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-wasm",
|
||||
AlgoType: "wasm",
|
||||
Algorithm: []byte{0x00, 0x61, 0x73, 0x6d},
|
||||
Args: []string{},
|
||||
}
|
||||
|
||||
resp, err := rs.Run(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
if resp.Error != "" {
|
||||
assert.Contains(t, resp.Error, "wasmedge")
|
||||
t.Skip("wasmedge not found, skipping test")
|
||||
}
|
||||
assert.Equal(t, "test-wasm", resp.ComputationId)
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove("algo")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRunWithDockerAlgorithm tests running a Docker algorithm.
|
||||
func TestRunWithDockerAlgorithm(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-docker",
|
||||
AlgoType: "docker",
|
||||
Algorithm: []byte("FROM ubuntu:latest\nRUN echo 'test'"),
|
||||
Args: []string{},
|
||||
}
|
||||
|
||||
resp, err := rs.Run(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
if resp.Error != "" {
|
||||
assert.Contains(t, resp.Error, "Docker")
|
||||
t.Skip("Docker issue, skipping test")
|
||||
}
|
||||
assert.Equal(t, "test-docker", resp.ComputationId)
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove("algo")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRunWithUnsupportedAlgorithmType tests running with unsupported algorithm type.
|
||||
func TestRunWithUnsupportedAlgorithmType(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-unsupported",
|
||||
AlgoType: "unsupported",
|
||||
Algorithm: []byte("test"),
|
||||
Args: []string{},
|
||||
}
|
||||
|
||||
resp, err := rs.Run(context.Background(), req)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, resp)
|
||||
}
|
||||
|
||||
// TestRunAlreadyRunning tests running computation when one is already running.
|
||||
func TestRunAlreadyRunning(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
// Use a long-running bash script
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-running",
|
||||
AlgoType: "bin",
|
||||
Algorithm: []byte("#!/bin/bash\nsleep 30"),
|
||||
Args: []string{},
|
||||
}
|
||||
|
||||
// Start first computation (will run for 30 seconds)
|
||||
go func() {
|
||||
_, _ = rs.Run(context.Background(), req)
|
||||
}()
|
||||
|
||||
// Give it time to start
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Try to run another immediately - should fail
|
||||
resp, err := rs.Run(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
assert.Equal(t, "computation already running", resp.Error)
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove("algo")
|
||||
})
|
||||
}
|
||||
|
||||
// TestStopWhenRunning tests stopping a running computation.
|
||||
func TestStopWhenRunning(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-stop",
|
||||
AlgoType: "bin",
|
||||
Algorithm: []byte("#!/bin/bash\nsleep 10"),
|
||||
Args: []string{},
|
||||
}
|
||||
|
||||
go func() {
|
||||
_, _ = rs.Run(context.Background(), req)
|
||||
}()
|
||||
|
||||
// Give it time to start
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
stopReq := &pb.StopRequest{
|
||||
ComputationId: "test-stop",
|
||||
}
|
||||
|
||||
stopResp, err := rs.Stop(context.Background(), stopReq)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, stopResp)
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove("algo")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRunErrors tests error paths in Run.
|
||||
func TestRunErrors(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
t.Run("create algo file failure", func(t *testing.T) {
|
||||
// Create a directory named "algo" to make os.Create("algo") fail
|
||||
err := os.Mkdir("algo", 0o755)
|
||||
require.NoError(t, err)
|
||||
defer os.RemoveAll("algo")
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-err",
|
||||
AlgoType: "bin",
|
||||
Algorithm: []byte("test"),
|
||||
}
|
||||
_, err = rs.Run(context.Background(), req)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "error creating algorithm file")
|
||||
})
|
||||
|
||||
t.Run("getwd failure", func(t *testing.T) {
|
||||
origDir, _ := os.Getwd()
|
||||
tmpDir := t.TempDir()
|
||||
err := os.Chdir(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Remove the current working directory to trigger Getwd failure
|
||||
err = os.RemoveAll(tmpDir)
|
||||
require.NoError(t, err)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-err-getwd",
|
||||
AlgoType: "bin",
|
||||
Algorithm: []byte("test"),
|
||||
}
|
||||
_, err = rs.Run(context.Background(), req)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "error getting current directory")
|
||||
|
||||
// Restore working directory
|
||||
_ = os.Chdir(origDir)
|
||||
})
|
||||
|
||||
t.Run("requirements file creation failure", func(t *testing.T) {
|
||||
// This one is harder because it uses os.CreateTemp("", "requirements.txt")
|
||||
// We can't easily make this fail without reaching into the system's temp dir.
|
||||
// Skipping for now as it's a very unlikely edge case.
|
||||
})
|
||||
|
||||
t.Run("chmod failure", func(t *testing.T) {
|
||||
// We can't easily mock os.Chmod, but we can try to make the file unmodifiable
|
||||
// On Linux, we can set the immutable attribute, but that requires root.
|
||||
// Alternatively, we can try to use a directory with permissions that prevent chmod?
|
||||
// No, chmod usually works if you own the file.
|
||||
})
|
||||
|
||||
t.Run("write algorithm failure", func(t *testing.T) {
|
||||
// This is also hard without mocking os.File.Write or reaching internal limits.
|
||||
})
|
||||
}
|
||||
|
||||
// TestConcurrentRun tests that concurrent runs are properly serialized.
|
||||
func TestConcurrentRun(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-concurrent",
|
||||
AlgoType: "bin",
|
||||
Algorithm: []byte("#!/bin/bash\nsleep 15"),
|
||||
Args: []string{},
|
||||
}
|
||||
|
||||
// Start first run in goroutine (will run for 15 seconds)
|
||||
go func() {
|
||||
_, _ = rs.Run(context.Background(), req)
|
||||
}()
|
||||
|
||||
// Give it time to actually start
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Concurrent attempt should fail
|
||||
resp2, err := rs.Run(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "computation already running", resp2.Error)
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove("algo")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRunWithMultipleArgs tests running with multiple arguments.
|
||||
func TestRunWithMultipleArgs(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
req := &pb.RunRequest{
|
||||
ComputationId: "test-multi-args",
|
||||
AlgoType: "bin",
|
||||
Algorithm: []byte("#!/bin/bash\necho $@"),
|
||||
Args: []string{"arg1", "arg2", "arg3", "arg4"},
|
||||
}
|
||||
|
||||
resp, err := rs.Run(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
assert.Empty(t, resp.Error)
|
||||
assert.Equal(t, "test-multi-args", resp.ComputationId)
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove("algo")
|
||||
})
|
||||
}
|
||||
|
||||
func TestStopFailure(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
eventSvc := &MockEventService{}
|
||||
rs := New(logger, eventSvc)
|
||||
|
||||
// Mock an algorithm that fails on Stop
|
||||
rs.currentAlgo = &MockAlgorithmStopFail{}
|
||||
|
||||
_, err := rs.Stop(context.Background(), &pb.StopRequest{})
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
type MockAlgorithmStopFail struct{}
|
||||
|
||||
func (m *MockAlgorithmStopFail) Run() error { return nil }
|
||||
func (m *MockAlgorithmStopFail) Stop() error { return fmt.Errorf("stop failed") }
|
||||
+150
-1270
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+117
@@ -0,0 +1,117 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
)
|
||||
|
||||
//go:generate stringer -type=state
|
||||
type state int
|
||||
|
||||
const (
|
||||
idle state = iota
|
||||
receivingManifests
|
||||
receivingAlgorithms
|
||||
receivingData
|
||||
running
|
||||
resultsReady
|
||||
complete
|
||||
)
|
||||
|
||||
type event int
|
||||
|
||||
const (
|
||||
start event = iota
|
||||
manifestsReceived
|
||||
algorithmsReceived
|
||||
dataReceived
|
||||
runComplete
|
||||
resultsConsumed
|
||||
)
|
||||
|
||||
// StateMachine represents the state machine.
|
||||
type StateMachine struct {
|
||||
mu sync.Mutex
|
||||
State state
|
||||
EventChan chan event
|
||||
Transitions map[state]map[event]state
|
||||
StateFunctions map[state]func()
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
// NewStateMachine creates a new StateMachine.
|
||||
func NewStateMachine(logger *slog.Logger) *StateMachine {
|
||||
sm := &StateMachine{
|
||||
State: idle,
|
||||
EventChan: make(chan event),
|
||||
Transitions: make(map[state]map[event]state),
|
||||
StateFunctions: make(map[state]func()),
|
||||
logger: logger,
|
||||
}
|
||||
|
||||
sm.Transitions[idle] = make(map[event]state)
|
||||
sm.Transitions[idle][start] = receivingManifests
|
||||
|
||||
sm.Transitions[receivingManifests] = make(map[event]state)
|
||||
sm.Transitions[receivingManifests][manifestsReceived] = receivingAlgorithms
|
||||
|
||||
sm.Transitions[receivingAlgorithms] = make(map[event]state)
|
||||
sm.Transitions[receivingAlgorithms][algorithmsReceived] = receivingData
|
||||
|
||||
sm.Transitions[receivingData] = make(map[event]state)
|
||||
sm.Transitions[receivingData][dataReceived] = running
|
||||
|
||||
sm.Transitions[running] = make(map[event]state)
|
||||
sm.Transitions[running][runComplete] = resultsReady
|
||||
|
||||
sm.Transitions[resultsReady] = make(map[event]state)
|
||||
sm.Transitions[resultsReady][resultsConsumed] = complete
|
||||
|
||||
return sm
|
||||
}
|
||||
|
||||
// Start the state machine.
|
||||
func (sm *StateMachine) Start(ctx context.Context) {
|
||||
for {
|
||||
select {
|
||||
case event := <-sm.EventChan:
|
||||
nextState, valid := sm.Transitions[sm.GetState()][event]
|
||||
if valid {
|
||||
sm.mu.Lock()
|
||||
sm.State = nextState
|
||||
sm.mu.Unlock()
|
||||
sm.logger.Debug(fmt.Sprintf("Transition: %v -> %v\n", sm.GetState(), nextState))
|
||||
} else {
|
||||
sm.logger.Error(fmt.Sprintf("Invalid transition: %v -> ???\n", sm.GetState()))
|
||||
}
|
||||
stateFunc, exists := sm.StateFunctions[sm.GetState()]
|
||||
if exists {
|
||||
go stateFunc()
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SendEvent sends an event to the state machine.
|
||||
func (sm *StateMachine) SendEvent(event event) {
|
||||
sm.EventChan <- event
|
||||
}
|
||||
|
||||
func (sm *StateMachine) GetState() state {
|
||||
sm.mu.Lock()
|
||||
state := sm.State
|
||||
sm.mu.Unlock()
|
||||
return state
|
||||
}
|
||||
|
||||
func (sm *StateMachine) SetState(state state) {
|
||||
sm.mu.Lock()
|
||||
sm.State = state
|
||||
sm.mu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
// Code generated by "stringer -type=state"; DO NOT EDIT.
|
||||
|
||||
package agent
|
||||
|
||||
import "strconv"
|
||||
|
||||
func _() {
|
||||
// An "invalid array index" compiler error signifies that the constant values have changed.
|
||||
// Re-run the stringer command to generate them again.
|
||||
var x [1]struct{}
|
||||
_ = x[idle-0]
|
||||
_ = x[receivingManifests-1]
|
||||
_ = x[receivingAlgorithms-2]
|
||||
_ = x[receivingData-3]
|
||||
_ = x[running-4]
|
||||
_ = x[resultsReady-5]
|
||||
_ = x[complete-6]
|
||||
}
|
||||
|
||||
const _state_name = "idlereceivingManifestsreceivingAlgorithmsreceivingDatarunningresultsReadycomplete"
|
||||
|
||||
var _state_index = [...]uint8{0, 4, 22, 41, 54, 61, 73, 81}
|
||||
|
||||
func (i state) String() string {
|
||||
if i < 0 || i >= state(len(_state_index)-1) {
|
||||
return "state(" + strconv.FormatInt(int64(i), 10) + ")"
|
||||
}
|
||||
return _state_name[_state_index[i]:_state_index[i+1]]
|
||||
}
|
||||
+41
-236
@@ -4,255 +4,60 @@ package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
sync "sync"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/ultravioletrs/cocos/agent/statemachine"
|
||||
mglog "github.com/absmach/magistrala/logger"
|
||||
)
|
||||
|
||||
type MockState int
|
||||
|
||||
type MockEvent int
|
||||
|
||||
func (s MockState) String() string {
|
||||
return []string{"State1", "State2", "State3"}[s]
|
||||
}
|
||||
|
||||
func (e MockEvent) String() string {
|
||||
return []string{"Event1", "Event2", "Event3"}[e]
|
||||
}
|
||||
|
||||
const (
|
||||
State1 MockState = iota
|
||||
State2
|
||||
State3
|
||||
)
|
||||
|
||||
const (
|
||||
Event1 MockEvent = iota
|
||||
Event2
|
||||
Event3
|
||||
)
|
||||
|
||||
func TestNewStateMachine(t *testing.T) {
|
||||
sm := statemachine.NewStateMachine(State1)
|
||||
if sm == nil {
|
||||
t.Fatal("NewStateMachine returned nil")
|
||||
}
|
||||
if sm.GetState() != State1 {
|
||||
t.Errorf("Initial state not set correctly, got %v, want %v", sm.GetState(), State1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAddTransition(t *testing.T) {
|
||||
sm := statemachine.NewStateMachine(State1)
|
||||
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
go func() {
|
||||
if err := sm.Start(ctx); err != context.Canceled {
|
||||
}
|
||||
}()
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
sm.SendEvent(Event1)
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
if sm.GetState() != State2 {
|
||||
t.Errorf("Transition not applied correctly, got state %v, want %v", sm.GetState(), State2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAction(t *testing.T) {
|
||||
sm := statemachine.NewStateMachine(State1)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
|
||||
sm.SetAction(State2, func(s statemachine.State) {
|
||||
defer wg.Done()
|
||||
})
|
||||
|
||||
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
go func() {
|
||||
if err := sm.Start(ctx); err != context.Canceled {
|
||||
t.Errorf("Start returned error: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
sm.SendEvent(Event1)
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
wg.Wait()
|
||||
|
||||
if ctx.Err() != nil {
|
||||
t.Error("Action was not called within the expected time")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInvalidTransition(t *testing.T) {
|
||||
sm := statemachine.NewStateMachine(State1)
|
||||
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
errChan := make(chan error)
|
||||
go func() {
|
||||
errChan <- sm.Start(ctx)
|
||||
}()
|
||||
|
||||
sm.SendEvent(Event2)
|
||||
|
||||
select {
|
||||
case err := <-errChan:
|
||||
if err == nil {
|
||||
t.Errorf("Expected invalid transition error, got: %v", err)
|
||||
}
|
||||
case <-time.After(150 * time.Millisecond):
|
||||
t.Error("Timeout waiting for invalid transition error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultipleTransitions(t *testing.T) {
|
||||
sm := statemachine.NewStateMachine(State1)
|
||||
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
|
||||
sm.AddTransition(statemachine.Transition{From: State2, Event: Event2, To: State3})
|
||||
sm.AddTransition(statemachine.Transition{From: State3, Event: Event3, To: State1})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
go func() {
|
||||
if err := sm.Start(ctx); err != context.Canceled {
|
||||
}
|
||||
}()
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
transitions := []struct {
|
||||
event MockEvent
|
||||
want MockState
|
||||
func TestStateMachineTransitions(t *testing.T) {
|
||||
testCases := []struct {
|
||||
fromState state
|
||||
event event
|
||||
expected state
|
||||
}{
|
||||
{Event1, State2},
|
||||
{Event2, State3},
|
||||
{Event3, State1},
|
||||
{idle, start, receivingManifests},
|
||||
{receivingManifests, manifestsReceived, receivingAlgorithms},
|
||||
{receivingAlgorithms, algorithmsReceived, receivingData},
|
||||
{receivingData, dataReceived, running},
|
||||
{running, runComplete, resultsReady},
|
||||
{resultsReady, resultsConsumed, complete},
|
||||
}
|
||||
|
||||
for _, tt := range transitions {
|
||||
sm.SendEvent(tt.event)
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
for _, testCase := range testCases {
|
||||
t.Run(fmt.Sprintf("Transition from %v to %v", testCase.fromState, testCase.expected), func(t *testing.T) {
|
||||
sm := NewStateMachine(mglog.NewMock())
|
||||
done := make(chan struct{})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() {
|
||||
sm.Start(ctx)
|
||||
close(done)
|
||||
}()
|
||||
sm.SetState(testCase.fromState)
|
||||
|
||||
if sm.GetState() != tt.want {
|
||||
t.Errorf("After event %v, got state %v, want %v", tt.event, sm.GetState(), tt.want)
|
||||
}
|
||||
sm.SendEvent(testCase.event)
|
||||
|
||||
if sm.GetState() != testCase.expected {
|
||||
t.Errorf("Expected state %v after the event, but got %v", testCase.expected, sm.GetState())
|
||||
}
|
||||
close(sm.EventChan)
|
||||
cancel()
|
||||
<-done
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrency(t *testing.T) {
|
||||
sm := statemachine.NewStateMachine(State1)
|
||||
sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
|
||||
sm.AddTransition(statemachine.Transition{From: State2, Event: Event2, To: State1})
|
||||
func TestStateMachineInvalidTransition(t *testing.T) {
|
||||
sm := NewStateMachine(mglog.NewMock())
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go sm.Start(ctx)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
|
||||
defer cancel()
|
||||
sm.SetState(idle)
|
||||
|
||||
go func() {
|
||||
if err := sm.Start(ctx); err == nil {
|
||||
t.Errorf("Expected context error, got nil")
|
||||
}
|
||||
}()
|
||||
sm.SendEvent(dataReceived)
|
||||
|
||||
for i := 0; i < 100; i++ {
|
||||
go func() {
|
||||
sm.SendEvent(Event1)
|
||||
sm.SendEvent(Event2)
|
||||
}()
|
||||
}
|
||||
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
|
||||
finalState := sm.GetState()
|
||||
if finalState != State1 && finalState != State2 {
|
||||
t.Errorf("Unexpected final state: %v", finalState)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentEventString(t *testing.T) {
|
||||
tests := []struct {
|
||||
event AgentEvent
|
||||
want string
|
||||
}{
|
||||
{Start, "Start"},
|
||||
{ManifestReceived, "ManifestReceived"},
|
||||
{AlgorithmReceived, "AlgorithmReceived"},
|
||||
{DataReceived, "DataReceived"},
|
||||
{RunComplete, "RunComplete"},
|
||||
{ResultsConsumed, "ResultsConsumed"},
|
||||
{RunFailed, "RunFailed"},
|
||||
{AgentEvent(-1), "AgentEvent(-1)"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
if got := tt.event.String(); got != tt.want {
|
||||
t.Errorf("AgentEvent.String() = %v, want %v", got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentStateString(t *testing.T) {
|
||||
tests := []struct {
|
||||
state AgentState
|
||||
want string
|
||||
}{
|
||||
{Idle, "Idle"},
|
||||
{ReceivingManifest, "ReceivingManifest"},
|
||||
{ReceivingAlgorithm, "ReceivingAlgorithm"},
|
||||
{ReceivingData, "ReceivingData"},
|
||||
{Running, "Running"},
|
||||
{ConsumingResults, "ConsumingResults"},
|
||||
{Complete, "Complete"},
|
||||
{Failed, "Failed"},
|
||||
{AgentState(-1), "AgentState(-1)"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
if got := tt.state.String(); got != tt.want {
|
||||
t.Errorf("AgentState.String() = %v, want %v", got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusString(t *testing.T) {
|
||||
tests := []struct {
|
||||
status Status
|
||||
want string
|
||||
}{
|
||||
{IdleState, "IdleState"},
|
||||
{InProgress, "InProgress"},
|
||||
{Ready, "Ready"},
|
||||
{Completed, "Completed"},
|
||||
{Terminated, "Terminated"},
|
||||
{Warning, "Warning"},
|
||||
{Starting, "Starting"},
|
||||
{Status(uint8(8)), "Status(8)"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
if got := tt.status.String(); got != tt.want {
|
||||
t.Errorf("Status.String() = %v, want %v", got, tt.want)
|
||||
}
|
||||
if sm.State != idle {
|
||||
t.Errorf("State should not change on an invalid event, but got %v", sm.State)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
|
||||
@@ -1,305 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Code generated by mockery; DO NOT EDIT.
|
||||
// github.com/vektra/mockery
|
||||
// template: testify
|
||||
|
||||
package mocks
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
mock "github.com/stretchr/testify/mock"
|
||||
"github.com/ultravioletrs/cocos/agent/statemachine"
|
||||
)
|
||||
|
||||
// NewStateMachine creates a new instance of StateMachine. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
|
||||
// The first argument is typically a *testing.T value.
|
||||
func NewStateMachine(t interface {
|
||||
mock.TestingT
|
||||
Cleanup(func())
|
||||
}) *StateMachine {
|
||||
mock := &StateMachine{}
|
||||
mock.Mock.Test(t)
|
||||
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
|
||||
// StateMachine is an autogenerated mock type for the StateMachine type
|
||||
type StateMachine struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
type StateMachine_Expecter struct {
|
||||
mock *mock.Mock
|
||||
}
|
||||
|
||||
func (_m *StateMachine) EXPECT() *StateMachine_Expecter {
|
||||
return &StateMachine_Expecter{mock: &_m.Mock}
|
||||
}
|
||||
|
||||
// AddTransition provides a mock function for the type StateMachine
|
||||
func (_mock *StateMachine) AddTransition(t statemachine.Transition) {
|
||||
_mock.Called(t)
|
||||
return
|
||||
}
|
||||
|
||||
// StateMachine_AddTransition_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AddTransition'
|
||||
type StateMachine_AddTransition_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// AddTransition is a helper method to define mock.On call
|
||||
// - t statemachine.Transition
|
||||
func (_e *StateMachine_Expecter) AddTransition(t interface{}) *StateMachine_AddTransition_Call {
|
||||
return &StateMachine_AddTransition_Call{Call: _e.mock.On("AddTransition", t)}
|
||||
}
|
||||
|
||||
func (_c *StateMachine_AddTransition_Call) Run(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 statemachine.Transition
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(statemachine.Transition)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_AddTransition_Call) Return() *StateMachine_AddTransition_Call {
|
||||
_c.Call.Return()
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_AddTransition_Call) RunAndReturn(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
|
||||
_c.Run(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// GetState provides a mock function for the type StateMachine
|
||||
func (_mock *StateMachine) GetState() statemachine.State {
|
||||
ret := _mock.Called()
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for GetState")
|
||||
}
|
||||
|
||||
var r0 statemachine.State
|
||||
if returnFunc, ok := ret.Get(0).(func() statemachine.State); ok {
|
||||
r0 = returnFunc()
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(statemachine.State)
|
||||
}
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// StateMachine_GetState_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetState'
|
||||
type StateMachine_GetState_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// GetState is a helper method to define mock.On call
|
||||
func (_e *StateMachine_Expecter) GetState() *StateMachine_GetState_Call {
|
||||
return &StateMachine_GetState_Call{Call: _e.mock.On("GetState")}
|
||||
}
|
||||
|
||||
func (_c *StateMachine_GetState_Call) Run(run func()) *StateMachine_GetState_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
run()
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_GetState_Call) Return(state statemachine.State) *StateMachine_GetState_Call {
|
||||
_c.Call.Return(state)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_GetState_Call) RunAndReturn(run func() statemachine.State) *StateMachine_GetState_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Reset provides a mock function for the type StateMachine
|
||||
func (_mock *StateMachine) Reset(initialState statemachine.State) {
|
||||
_mock.Called(initialState)
|
||||
return
|
||||
}
|
||||
|
||||
// StateMachine_Reset_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Reset'
|
||||
type StateMachine_Reset_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Reset is a helper method to define mock.On call
|
||||
// - initialState statemachine.State
|
||||
func (_e *StateMachine_Expecter) Reset(initialState interface{}) *StateMachine_Reset_Call {
|
||||
return &StateMachine_Reset_Call{Call: _e.mock.On("Reset", initialState)}
|
||||
}
|
||||
|
||||
func (_c *StateMachine_Reset_Call) Run(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 statemachine.State
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(statemachine.State)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_Reset_Call) Return() *StateMachine_Reset_Call {
|
||||
_c.Call.Return()
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_Reset_Call) RunAndReturn(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
|
||||
_c.Run(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SendEvent provides a mock function for the type StateMachine
|
||||
func (_mock *StateMachine) SendEvent(event statemachine.Event) {
|
||||
_mock.Called(event)
|
||||
return
|
||||
}
|
||||
|
||||
// StateMachine_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
|
||||
type StateMachine_SendEvent_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// SendEvent is a helper method to define mock.On call
|
||||
// - event statemachine.Event
|
||||
func (_e *StateMachine_Expecter) SendEvent(event interface{}) *StateMachine_SendEvent_Call {
|
||||
return &StateMachine_SendEvent_Call{Call: _e.mock.On("SendEvent", event)}
|
||||
}
|
||||
|
||||
func (_c *StateMachine_SendEvent_Call) Run(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 statemachine.Event
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(statemachine.Event)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_SendEvent_Call) Return() *StateMachine_SendEvent_Call {
|
||||
_c.Call.Return()
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_SendEvent_Call) RunAndReturn(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
|
||||
_c.Run(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetAction provides a mock function for the type StateMachine
|
||||
func (_mock *StateMachine) SetAction(state statemachine.State, action statemachine.Action) {
|
||||
_mock.Called(state, action)
|
||||
return
|
||||
}
|
||||
|
||||
// StateMachine_SetAction_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetAction'
|
||||
type StateMachine_SetAction_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// SetAction is a helper method to define mock.On call
|
||||
// - state statemachine.State
|
||||
// - action statemachine.Action
|
||||
func (_e *StateMachine_Expecter) SetAction(state interface{}, action interface{}) *StateMachine_SetAction_Call {
|
||||
return &StateMachine_SetAction_Call{Call: _e.mock.On("SetAction", state, action)}
|
||||
}
|
||||
|
||||
func (_c *StateMachine_SetAction_Call) Run(run func(state statemachine.State, action statemachine.Action)) *StateMachine_SetAction_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 statemachine.State
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(statemachine.State)
|
||||
}
|
||||
var arg1 statemachine.Action
|
||||
if args[1] != nil {
|
||||
arg1 = args[1].(statemachine.Action)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
arg1,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_SetAction_Call) Return() *StateMachine_SetAction_Call {
|
||||
_c.Call.Return()
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_SetAction_Call) RunAndReturn(run func(state statemachine.State, action statemachine.Action)) *StateMachine_SetAction_Call {
|
||||
_c.Run(run)
|
||||
return _c
|
||||
}
|
||||
|
||||
// Start provides a mock function for the type StateMachine
|
||||
func (_mock *StateMachine) Start(ctx context.Context) error {
|
||||
ret := _mock.Called(ctx)
|
||||
|
||||
if len(ret) == 0 {
|
||||
panic("no return value specified for Start")
|
||||
}
|
||||
|
||||
var r0 error
|
||||
if returnFunc, ok := ret.Get(0).(func(context.Context) error); ok {
|
||||
r0 = returnFunc(ctx)
|
||||
} else {
|
||||
r0 = ret.Error(0)
|
||||
}
|
||||
return r0
|
||||
}
|
||||
|
||||
// StateMachine_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
|
||||
type StateMachine_Start_Call struct {
|
||||
*mock.Call
|
||||
}
|
||||
|
||||
// Start is a helper method to define mock.On call
|
||||
// - ctx context.Context
|
||||
func (_e *StateMachine_Expecter) Start(ctx interface{}) *StateMachine_Start_Call {
|
||||
return &StateMachine_Start_Call{Call: _e.mock.On("Start", ctx)}
|
||||
}
|
||||
|
||||
func (_c *StateMachine_Start_Call) Run(run func(ctx context.Context)) *StateMachine_Start_Call {
|
||||
_c.Call.Run(func(args mock.Arguments) {
|
||||
var arg0 context.Context
|
||||
if args[0] != nil {
|
||||
arg0 = args[0].(context.Context)
|
||||
}
|
||||
run(
|
||||
arg0,
|
||||
)
|
||||
})
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_Start_Call) Return(err error) *StateMachine_Start_Call {
|
||||
_c.Call.Return(err)
|
||||
return _c
|
||||
}
|
||||
|
||||
func (_c *StateMachine_Start_Call) RunAndReturn(run func(ctx context.Context) error) *StateMachine_Start_Call {
|
||||
_c.Call.Return(run)
|
||||
return _c
|
||||
}
|
||||
@@ -1,147 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package statemachine
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type State interface {
|
||||
String() string
|
||||
}
|
||||
|
||||
type Event interface {
|
||||
String() string
|
||||
}
|
||||
|
||||
type Action func(State)
|
||||
|
||||
type Transition struct {
|
||||
From State
|
||||
Event Event
|
||||
To State
|
||||
}
|
||||
|
||||
type StateMachine interface {
|
||||
AddTransition(t Transition)
|
||||
SetAction(state State, action Action)
|
||||
GetState() State
|
||||
SendEvent(event Event)
|
||||
Start(ctx context.Context) error
|
||||
Reset(initialState State)
|
||||
}
|
||||
|
||||
type stateMachine struct {
|
||||
mu sync.Mutex
|
||||
currentState State
|
||||
transitions map[State]map[Event]State
|
||||
actions map[State]Action
|
||||
eventChan chan Event
|
||||
resetChan chan struct{}
|
||||
}
|
||||
|
||||
func NewStateMachine(initialState State) StateMachine {
|
||||
return &stateMachine{
|
||||
currentState: initialState,
|
||||
transitions: make(map[State]map[Event]State),
|
||||
actions: make(map[State]Action),
|
||||
eventChan: make(chan Event),
|
||||
resetChan: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *stateMachine) AddTransition(t Transition) {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
|
||||
if _, ok := sm.transitions[t.From]; !ok {
|
||||
sm.transitions[t.From] = make(map[Event]State)
|
||||
}
|
||||
sm.transitions[t.From][t.Event] = t.To
|
||||
}
|
||||
|
||||
func (sm *stateMachine) SetAction(state State, action Action) {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
|
||||
sm.actions[state] = action
|
||||
}
|
||||
|
||||
func (sm *stateMachine) GetState() State {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
return sm.currentState
|
||||
}
|
||||
|
||||
func (sm *stateMachine) SendEvent(event Event) {
|
||||
sm.mu.Lock()
|
||||
eventChan := sm.eventChan
|
||||
sm.mu.Unlock()
|
||||
|
||||
select {
|
||||
case eventChan <- event:
|
||||
default:
|
||||
// Channel might be closed or full, ignore the event
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *stateMachine) Start(ctx context.Context) error {
|
||||
for {
|
||||
sm.mu.Lock()
|
||||
eventChan := sm.eventChan
|
||||
resetChan := sm.resetChan
|
||||
sm.mu.Unlock()
|
||||
|
||||
select {
|
||||
case event := <-eventChan:
|
||||
if err := sm.handleEvent(event); err != nil {
|
||||
return err
|
||||
}
|
||||
case <-resetChan:
|
||||
continue
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *stateMachine) Reset(initialState State) {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
|
||||
// Reset current state to initial state
|
||||
sm.currentState = initialState
|
||||
|
||||
// Close the existing event channel to stop processing events
|
||||
close(sm.eventChan)
|
||||
|
||||
// Close the reset channel to signal Start() to restart
|
||||
close(sm.resetChan)
|
||||
|
||||
sm.eventChan = make(chan Event)
|
||||
sm.resetChan = make(chan struct{})
|
||||
}
|
||||
|
||||
func (sm *stateMachine) handleEvent(event Event) error {
|
||||
sm.mu.Lock()
|
||||
currentState := sm.currentState
|
||||
nextState, valid := sm.transitions[currentState][event]
|
||||
sm.mu.Unlock()
|
||||
|
||||
if !valid {
|
||||
return fmt.Errorf("invalid transition: %v -> %v", currentState, event)
|
||||
}
|
||||
|
||||
sm.mu.Lock()
|
||||
sm.currentState = nextState
|
||||
action := sm.actions[nextState]
|
||||
sm.mu.Unlock()
|
||||
|
||||
if action != nil {
|
||||
go action(nextState)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,607 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package statemachine
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type testState string
|
||||
|
||||
func (s testState) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
type testEvent string
|
||||
|
||||
func (e testEvent) String() string {
|
||||
return string(e)
|
||||
}
|
||||
|
||||
const (
|
||||
StateIdle testState = "idle"
|
||||
StateRunning testState = "running"
|
||||
StatePaused testState = "paused"
|
||||
StateStopped testState = "stopped"
|
||||
StateError testState = "error"
|
||||
)
|
||||
|
||||
const (
|
||||
EventStart testEvent = "start"
|
||||
EventPause testEvent = "pause"
|
||||
EventStop testEvent = "stop"
|
||||
EventReset testEvent = "reset"
|
||||
EventError testEvent = "error"
|
||||
)
|
||||
|
||||
func TestNewStateMachine(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialState State
|
||||
want State
|
||||
}{
|
||||
{
|
||||
name: "create with idle state",
|
||||
initialState: StateIdle,
|
||||
want: StateIdle,
|
||||
},
|
||||
{
|
||||
name: "create with running state",
|
||||
initialState: StateRunning,
|
||||
want: StateRunning,
|
||||
},
|
||||
{
|
||||
name: "create with custom state",
|
||||
initialState: testState("custom"),
|
||||
want: testState("custom"),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sm := NewStateMachine(tt.initialState)
|
||||
if got := sm.GetState(); got != tt.want {
|
||||
t.Errorf("NewStateMachine() initial state = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateMachine_AddTransition(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
transitions []Transition
|
||||
from State
|
||||
event Event
|
||||
expectTo State
|
||||
expectValid bool
|
||||
}{
|
||||
{
|
||||
name: "single transition",
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
},
|
||||
from: StateIdle,
|
||||
event: EventStart,
|
||||
expectTo: StateRunning,
|
||||
expectValid: true,
|
||||
},
|
||||
{
|
||||
name: "multiple transitions from same state",
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
{From: StateIdle, Event: EventError, To: StateError},
|
||||
},
|
||||
from: StateIdle,
|
||||
event: EventError,
|
||||
expectTo: StateError,
|
||||
expectValid: true,
|
||||
},
|
||||
{
|
||||
name: "overwrite existing transition",
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
{From: StateIdle, Event: EventStart, To: StatePaused}, // Overwrite
|
||||
},
|
||||
from: StateIdle,
|
||||
event: EventStart,
|
||||
expectTo: StatePaused,
|
||||
expectValid: true,
|
||||
},
|
||||
{
|
||||
name: "transition not found",
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
},
|
||||
from: StateRunning,
|
||||
event: EventPause,
|
||||
expectValid: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sm := NewStateMachine(StateIdle).(*stateMachine)
|
||||
|
||||
for _, transition := range tt.transitions {
|
||||
sm.AddTransition(transition)
|
||||
}
|
||||
|
||||
sm.mu.Lock()
|
||||
nextState, valid := sm.transitions[tt.from][tt.event]
|
||||
sm.mu.Unlock()
|
||||
|
||||
if valid != tt.expectValid {
|
||||
t.Errorf("Transition validity = %v, want %v", valid, tt.expectValid)
|
||||
}
|
||||
|
||||
if tt.expectValid && nextState != tt.expectTo {
|
||||
t.Errorf("Transition destination = %v, want %v", nextState, tt.expectTo)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateMachine_SetAction(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
state State
|
||||
action Action
|
||||
expectAction bool
|
||||
}{
|
||||
{
|
||||
name: "set action for state",
|
||||
state: StateRunning,
|
||||
action: func(s State) {
|
||||
},
|
||||
expectAction: true,
|
||||
},
|
||||
{
|
||||
name: "set nil action",
|
||||
state: StatePaused,
|
||||
action: nil,
|
||||
expectAction: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sm := NewStateMachine(StateIdle).(*stateMachine)
|
||||
sm.SetAction(tt.state, tt.action)
|
||||
|
||||
sm.mu.Lock()
|
||||
action := sm.actions[tt.state]
|
||||
sm.mu.Unlock()
|
||||
|
||||
if tt.expectAction && action == nil {
|
||||
t.Error("Expected action to be set, but it was nil")
|
||||
}
|
||||
if !tt.expectAction && action != nil {
|
||||
t.Error("Expected action to be nil, but it was set")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateMachine_GetState(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialState State
|
||||
transitions []Transition
|
||||
events []Event
|
||||
finalState State
|
||||
}{
|
||||
{
|
||||
name: "get initial state",
|
||||
initialState: StateIdle,
|
||||
finalState: StateIdle,
|
||||
},
|
||||
{
|
||||
name: "get state after transition",
|
||||
initialState: StateIdle,
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
},
|
||||
events: []Event{EventStart},
|
||||
finalState: StateRunning,
|
||||
},
|
||||
{
|
||||
name: "get state after multiple transitions",
|
||||
initialState: StateIdle,
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
{From: StateRunning, Event: EventPause, To: StatePaused},
|
||||
{From: StatePaused, Event: EventStart, To: StateRunning},
|
||||
},
|
||||
events: []Event{EventStart, EventPause, EventStart},
|
||||
finalState: StateRunning,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sm := NewStateMachine(tt.initialState)
|
||||
|
||||
for _, transition := range tt.transitions {
|
||||
sm.AddTransition(transition)
|
||||
}
|
||||
|
||||
smImpl := sm.(*stateMachine)
|
||||
for _, event := range tt.events {
|
||||
if err := smImpl.handleEvent(event); err != nil {
|
||||
t.Fatalf("Failed to handle event %v: %v", event, err)
|
||||
}
|
||||
}
|
||||
|
||||
if got := sm.GetState(); got != tt.finalState {
|
||||
t.Errorf("GetState() = %v, want %v", got, tt.finalState)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateMachine_Start(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialState State
|
||||
transitions []Transition
|
||||
events []Event
|
||||
cancelAfter time.Duration
|
||||
expectError bool
|
||||
expectedStates []State
|
||||
}{
|
||||
{
|
||||
name: "start and cancel immediately",
|
||||
initialState: StateIdle,
|
||||
cancelAfter: 10 * time.Millisecond,
|
||||
expectError: true, // context.Canceled
|
||||
},
|
||||
{
|
||||
name: "process events then cancel",
|
||||
initialState: StateIdle,
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
{From: StateRunning, Event: EventStop, To: StateStopped},
|
||||
},
|
||||
events: []Event{EventStart, EventStop},
|
||||
cancelAfter: 100 * time.Millisecond,
|
||||
expectError: true, // context.Canceled
|
||||
expectedStates: []State{StateRunning, StateStopped},
|
||||
},
|
||||
{
|
||||
name: "invalid transition error",
|
||||
initialState: StateIdle,
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
},
|
||||
events: []Event{EventPause}, // Invalid from StateIdle
|
||||
cancelAfter: 50 * time.Millisecond,
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sm := NewStateMachine(tt.initialState)
|
||||
|
||||
for _, transition := range tt.transitions {
|
||||
sm.AddTransition(transition)
|
||||
}
|
||||
|
||||
var states []State
|
||||
var mu sync.Mutex
|
||||
|
||||
for _, state := range tt.expectedStates {
|
||||
sm.SetAction(state, func(s State) {
|
||||
mu.Lock()
|
||||
states = append(states, s)
|
||||
mu.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
errChan := make(chan error, 1)
|
||||
go func() {
|
||||
errChan <- sm.Start(ctx)
|
||||
}()
|
||||
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
|
||||
for _, event := range tt.events {
|
||||
sm.SendEvent(event)
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
|
||||
time.Sleep(tt.cancelAfter)
|
||||
cancel()
|
||||
|
||||
err := <-errChan
|
||||
|
||||
if tt.expectError && err == nil {
|
||||
t.Error("Expected error but got none")
|
||||
}
|
||||
if !tt.expectError && err != nil {
|
||||
t.Errorf("Unexpected error: %v", err)
|
||||
}
|
||||
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
mu.Lock()
|
||||
if len(states) != len(tt.expectedStates) {
|
||||
t.Errorf("Expected %d state changes, got %d", len(tt.expectedStates), len(states))
|
||||
}
|
||||
for i, expectedState := range tt.expectedStates {
|
||||
if i < len(states) && states[i] != expectedState {
|
||||
t.Errorf("State change %d = %v, want %v", i, states[i], expectedState)
|
||||
}
|
||||
}
|
||||
mu.Unlock()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateMachine_Reset(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialState State
|
||||
resetState State
|
||||
setupTransitions []Transition
|
||||
eventsBeforeReset []Event
|
||||
eventsAfterReset []Event
|
||||
expectedState State
|
||||
}{
|
||||
{
|
||||
name: "reset to same state",
|
||||
initialState: StateIdle,
|
||||
resetState: StateIdle,
|
||||
expectedState: StateIdle,
|
||||
},
|
||||
{
|
||||
name: "reset to different state",
|
||||
initialState: StateIdle,
|
||||
resetState: StateRunning,
|
||||
expectedState: StateRunning,
|
||||
},
|
||||
{
|
||||
name: "reset after state changes",
|
||||
initialState: StateIdle,
|
||||
resetState: StateIdle,
|
||||
setupTransitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
},
|
||||
eventsBeforeReset: []Event{EventStart},
|
||||
expectedState: StateIdle,
|
||||
},
|
||||
{
|
||||
name: "reset and send new events",
|
||||
initialState: StateIdle,
|
||||
resetState: StateIdle,
|
||||
setupTransitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
{From: StateRunning, Event: EventStop, To: StateStopped},
|
||||
},
|
||||
eventsBeforeReset: []Event{EventStart},
|
||||
eventsAfterReset: []Event{EventStart},
|
||||
expectedState: StateIdle,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sm := NewStateMachine(tt.initialState)
|
||||
smImpl := sm.(*stateMachine)
|
||||
|
||||
for _, transition := range tt.setupTransitions {
|
||||
sm.AddTransition(transition)
|
||||
}
|
||||
|
||||
for _, event := range tt.eventsBeforeReset {
|
||||
if err := smImpl.handleEvent(event); err != nil {
|
||||
// Ignore errors for this test
|
||||
}
|
||||
}
|
||||
|
||||
sm.Reset(tt.resetState)
|
||||
|
||||
if got := sm.GetState(); got != tt.expectedState {
|
||||
t.Errorf("State after reset = %v, want %v", got, tt.expectedState)
|
||||
}
|
||||
|
||||
for _, event := range tt.eventsAfterReset {
|
||||
sm.SendEvent(event)
|
||||
}
|
||||
|
||||
// For events after reset, we can't easily check the channel length
|
||||
// due to the synchronization changes, so we just verify the reset worked
|
||||
if len(tt.eventsAfterReset) > 0 {
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateMachine_Reset_WithRunningStateMachine(t *testing.T) {
|
||||
sm := NewStateMachine(StateIdle)
|
||||
sm.AddTransition(Transition{From: StateIdle, Event: EventStart, To: StateRunning})
|
||||
sm.AddTransition(Transition{From: StateRunning, Event: EventStop, To: StateStopped})
|
||||
|
||||
var stateChanges []State
|
||||
var mu sync.Mutex
|
||||
|
||||
sm.SetAction(StateRunning, func(s State) {
|
||||
mu.Lock()
|
||||
stateChanges = append(stateChanges, s)
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
sm.SetAction(StateStopped, func(s State) {
|
||||
mu.Lock()
|
||||
stateChanges = append(stateChanges, s)
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
go func() {
|
||||
if err := sm.Start(ctx); err != nil {
|
||||
}
|
||||
}()
|
||||
|
||||
// Give it time to start
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
|
||||
// Send an event
|
||||
sm.SendEvent(EventStart)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
// Reset while running
|
||||
sm.Reset(StateIdle)
|
||||
|
||||
// Verify state was reset
|
||||
if got := sm.GetState(); got != StateIdle {
|
||||
t.Errorf("State after reset = %v, want %v", got, StateIdle)
|
||||
}
|
||||
|
||||
// Send another event after reset
|
||||
sm.SendEvent(EventStart)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
mu.Lock()
|
||||
changes := len(stateChanges)
|
||||
mu.Unlock()
|
||||
|
||||
// Should have at least processed the first event
|
||||
if changes < 1 {
|
||||
t.Errorf("Expected at least 1 state change, got %d", changes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateMachine_HandleEvent(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialState State
|
||||
transitions []Transition
|
||||
event Event
|
||||
expectedState State
|
||||
expectError bool
|
||||
expectActionCall bool
|
||||
}{
|
||||
{
|
||||
name: "valid transition",
|
||||
initialState: StateIdle,
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
},
|
||||
event: EventStart,
|
||||
expectedState: StateRunning,
|
||||
expectError: false,
|
||||
expectActionCall: true,
|
||||
},
|
||||
{
|
||||
name: "invalid transition",
|
||||
initialState: StateIdle,
|
||||
transitions: []Transition{
|
||||
{From: StateRunning, Event: EventPause, To: StatePaused},
|
||||
},
|
||||
event: EventStart,
|
||||
expectedState: StateIdle,
|
||||
expectError: true,
|
||||
expectActionCall: false,
|
||||
},
|
||||
{
|
||||
name: "transition with no action",
|
||||
initialState: StateIdle,
|
||||
transitions: []Transition{
|
||||
{From: StateIdle, Event: EventStart, To: StateRunning},
|
||||
},
|
||||
event: EventStart,
|
||||
expectedState: StateRunning,
|
||||
expectError: false,
|
||||
expectActionCall: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sm := NewStateMachine(tt.initialState).(*stateMachine)
|
||||
|
||||
for _, transition := range tt.transitions {
|
||||
sm.AddTransition(transition)
|
||||
}
|
||||
|
||||
var actionCalled bool
|
||||
var mu sync.Mutex
|
||||
|
||||
if tt.expectActionCall {
|
||||
sm.SetAction(tt.expectedState, func(s State) {
|
||||
mu.Lock()
|
||||
actionCalled = true
|
||||
mu.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
err := sm.handleEvent(tt.event)
|
||||
|
||||
if tt.expectError && err == nil {
|
||||
t.Error("Expected error but got none")
|
||||
}
|
||||
if !tt.expectError && err != nil {
|
||||
t.Errorf("Unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if sm.GetState() != tt.expectedState {
|
||||
t.Errorf("State after handleEvent = %v, want %v", sm.GetState(), tt.expectedState)
|
||||
}
|
||||
|
||||
if tt.expectActionCall {
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
mu.Lock()
|
||||
called := actionCalled
|
||||
mu.Unlock()
|
||||
if !called {
|
||||
t.Error("Expected action to be called but it wasn't")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateMachine_SendEvent_ThreadSafety(t *testing.T) {
|
||||
sm := NewStateMachine(StateIdle)
|
||||
sm.AddTransition(Transition{From: StateIdle, Event: EventStart, To: StateRunning})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
go func() {
|
||||
if err := sm.Start(ctx); err != nil {
|
||||
}
|
||||
}()
|
||||
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 10
|
||||
eventsPerGoroutine := 100
|
||||
|
||||
// Send events concurrently
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < eventsPerGoroutine; j++ {
|
||||
sm.SendEvent(EventStart)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
// If we reach here without panicking, the test passes
|
||||
}
|
||||
@@ -1,29 +0,0 @@
|
||||
// Code generated by "stringer -type=Status"; DO NOT EDIT.
|
||||
|
||||
package agent
|
||||
|
||||
import "strconv"
|
||||
|
||||
func _() {
|
||||
// An "invalid array index" compiler error signifies that the constant values have changed.
|
||||
// Re-run the stringer command to generate them again.
|
||||
var x [1]struct{}
|
||||
_ = x[IdleState-0]
|
||||
_ = x[InProgress-1]
|
||||
_ = x[Ready-2]
|
||||
_ = x[Completed-3]
|
||||
_ = x[Terminated-4]
|
||||
_ = x[Warning-5]
|
||||
_ = x[Starting-6]
|
||||
}
|
||||
|
||||
const _Status_name = "IdleStateInProgressReadyCompletedTerminatedWarningStarting"
|
||||
|
||||
var _Status_index = [...]uint8{0, 9, 19, 24, 33, 43, 50, 58}
|
||||
|
||||
func (i Status) String() string {
|
||||
if i >= Status(len(_Status_index)-1) {
|
||||
return "Status(" + strconv.FormatInt(int64(i), 10) + ")"
|
||||
}
|
||||
return _Status_name[_Status_index[i]:_Status_index[i+1]]
|
||||
}
|
||||
Binary file not shown.
@@ -1,9 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package cocosai
|
||||
|
||||
import _ "embed"
|
||||
|
||||
//go:embed attestation.bin
|
||||
var EmbeddedAttestation []byte
|
||||
@@ -1,2 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
@@ -1,196 +0,0 @@
|
||||
# CoRIM Generation CLI Commands
|
||||
|
||||
This document describes the CLI commands for generating CoRIM (Concise Reference Integrity Manifest) attestation policies.
|
||||
|
||||
## Overview
|
||||
|
||||
The `cocos-cli policy create-corim` command provides subcommands for generating CoRIM policies for different platforms:
|
||||
- **azure**: Generate from Azure Attestation Token
|
||||
- **gcp**: Generate from GCP endorsements
|
||||
- **snp**: Generate for AMD SEV-SNP (direct host generation)
|
||||
- **tdx**: Generate for Intel TDX (direct host generation)
|
||||
|
||||
## Commands
|
||||
|
||||
### Azure SEV-SNP
|
||||
|
||||
Generate CoRIM from an Azure Attestation Token (JWT).
|
||||
|
||||
```bash
|
||||
cocos-cli policy create-corim azure --token <path-to-token> [--product <product>]
|
||||
```
|
||||
|
||||
**Flags:**
|
||||
- `--token` (required): Path to file containing Azure Attestation Token (JWT)
|
||||
- `--product` (optional): Processor product name (default: "Milan")
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
cocos-cli policy create-corim azure \
|
||||
--token /path/to/token.jwt \
|
||||
--product Milan \
|
||||
> azure-policy.corim
|
||||
```
|
||||
|
||||
### GCP SEV-SNP
|
||||
|
||||
Generate CoRIM from GCP SEV-SNP measurement and endorsements.
|
||||
|
||||
```bash
|
||||
cocos-cli policy create-corim gcp --measurement <hex> [--vcpu <num>]
|
||||
```
|
||||
|
||||
**Flags:**
|
||||
- `--measurement` (required): 384-bit measurement hex string
|
||||
- `--vcpu` (optional): vCPU number (default: 0)
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
cocos-cli policy create-corim gcp \
|
||||
--measurement abc123... \
|
||||
--vcpu 0 \
|
||||
> gcp-policy.corim
|
||||
```
|
||||
|
||||
### SEV-SNP (Direct Host)
|
||||
|
||||
Generate CoRIM for AMD SEV-SNP platform directly on the host.
|
||||
|
||||
```bash
|
||||
cocos-cli policy create-corim snp [flags]
|
||||
```
|
||||
|
||||
**Flags:**
|
||||
- `--measurement` (optional): Measurement/Launch Digest (hex string, defaults to zero if not provided)
|
||||
- `--policy` (optional): SNP policy flags (default: 0)
|
||||
- `--svn` (optional): Security Version Number/TCB (default: 0)
|
||||
- `--product` (optional): Processor product name (default: "Milan")
|
||||
- `--host-data` (optional): Host data (hex string)
|
||||
- `--launch-tcb` (optional): Minimum launch TCB (default: 0)
|
||||
- `--output` (optional): Output file path (default: stdout)
|
||||
|
||||
**Examples:**
|
||||
|
||||
Generate with defaults (zeroed measurement):
|
||||
```bash
|
||||
cocos-cli policy create-corim snp \
|
||||
--product Milan \
|
||||
--output snp-policy.corim
|
||||
```
|
||||
|
||||
Generate with custom measurement:
|
||||
```bash
|
||||
cocos-cli policy create-corim snp \
|
||||
--measurement abc123def456... \
|
||||
--product Genoa \
|
||||
--svn 1 \
|
||||
--policy 0x30000 \
|
||||
--output snp-policy.corim
|
||||
```
|
||||
|
||||
Generate with host data and launch TCB:
|
||||
```bash
|
||||
cocos-cli policy create-corim snp \
|
||||
--measurement abc123... \
|
||||
--host-data deadbeef \
|
||||
--launch-tcb 1 \
|
||||
--output snp-policy.corim
|
||||
```
|
||||
|
||||
### TDX (Direct Host)
|
||||
|
||||
Generate CoRIM for Intel TDX platform directly on the host.
|
||||
|
||||
```bash
|
||||
cocos-cli policy create-corim tdx [flags]
|
||||
```
|
||||
|
||||
**Flags:**
|
||||
- `--measurement` (optional): MRTD measurement (hex string, uses default if not provided)
|
||||
- `--svn` (optional): Security Version Number (default: 0)
|
||||
- `--rtmrs` (optional): Comma-separated RTMRs (hex)
|
||||
- `--mr-seam` (optional): MRSEAM (hex)
|
||||
- `--output` (optional): Output file path (default: stdout)
|
||||
|
||||
**Examples:**
|
||||
|
||||
Generate with defaults (matches legacy script behavior):
|
||||
```bash
|
||||
cocos-cli policy create-corim tdx \
|
||||
--output tdx-policy.corim
|
||||
```
|
||||
|
||||
Generate with custom values:
|
||||
```bash
|
||||
cocos-cli policy create-corim tdx \
|
||||
--measurement abc123def456... \
|
||||
--rtmrs rtmr0,rtmr1,rtmr2,rtmr3 \
|
||||
--mr-seam 789abc... \
|
||||
--svn 2 \
|
||||
--output tdx-policy.corim
|
||||
```
|
||||
|
||||
## Signing CoRIMs
|
||||
|
||||
CoRIMs can be signed using a private key (COSE_Sign1). The generated output will be a COSE-wrapped CoRIM in CBOR format.
|
||||
|
||||
### Prerequisite: Generate Signing Key
|
||||
|
||||
You will need an EC private key (P-256) in PEM format. You can generate one using `openssl`:
|
||||
|
||||
```bash
|
||||
openssl ecparam -name prime256v1 -genkey -noout -out private-key.pem
|
||||
```
|
||||
|
||||
### Signing with CLI
|
||||
|
||||
Use the `--signing-key` flag to sign the CoRIM during generation.
|
||||
|
||||
**SNP Example:**
|
||||
```bash
|
||||
cocos-cli policy create-corim snp \
|
||||
--product Milan \
|
||||
--signing-key private-key.pem \
|
||||
--output signed-snp.corim
|
||||
```
|
||||
|
||||
**TDX Example:**
|
||||
```bash
|
||||
cocos-cli policy create-corim tdx \
|
||||
--signing-key private-key.pem \
|
||||
--output signed-tdx.corim
|
||||
```
|
||||
|
||||
### Verification
|
||||
|
||||
The output file is a standard COSE_Sign1 message containing the CoRIM. It can be verified using any tool that supports COSE and CoRIM verification, such as the [veraison/corim](https://github.com/veraison/corim) library.
|
||||
|
||||
## Output Format
|
||||
|
||||
All commands output CoRIM in CBOR (Concise Binary Object Representation) format. By default, output is written to stdout, allowing for piping:
|
||||
|
||||
```bash
|
||||
# Pipe to file
|
||||
cocos-cli policy create-corim snp --product Milan > policy.corim
|
||||
|
||||
# Pipe to another command
|
||||
cocos-cli policy create-corim tdx | base64
|
||||
|
||||
# Use --output flag
|
||||
cocos-cli policy create-corim snp --product Milan --output policy.corim
|
||||
```
|
||||
|
||||
## Integration with Manager
|
||||
|
||||
The manager service can dynamically generate CoRIM policies using the same underlying generator package. When `FetchAttestationPolicy` is called:
|
||||
|
||||
1. For SNP: Calculates IGVM measurement using the `igvmmeasure` binary
|
||||
2. Extracts host data and launch TCB from VM configuration
|
||||
3. Generates CoRIM using the `generator` package
|
||||
4. Returns CBOR-encoded CoRIM
|
||||
|
||||
## See Also
|
||||
|
||||
- [Generator Package Documentation](../pkg/attestation/generator/README.md)
|
||||
- [IGVM Measure Package Documentation](../pkg/attestation/igvmmeasure/README.md)
|
||||
- [Manager README](../manager/README.md)
|
||||
+14
-45
@@ -1,6 +1,6 @@
|
||||
# Agent CLI
|
||||
|
||||
This repository contains the command-line interface (CLI) tool for interacting with the Agent and manager service. The CLI allows you to perform various tasks such as running computations, uploading algorithm and datasets, and retrieving results.
|
||||
This repository contains the command-line interface (CLI) tool for interacting with the Agent and manager service. The CLI allows you to perform various tasks such as running computations, uploading algorithms and datasets, and retrieving results.
|
||||
|
||||
## Build
|
||||
|
||||
@@ -12,18 +12,26 @@ make cli
|
||||
|
||||
## Usage
|
||||
|
||||
#### Run Computation
|
||||
|
||||
To run a computation, use the following command:
|
||||
|
||||
```bash
|
||||
./build/cocos-cli manager run --computation '{"name": "my-computation"}'
|
||||
```
|
||||
|
||||
#### Get attestation
|
||||
Retrieves attestation information from the SEV guest and saves it to a file.
|
||||
To retrieve attestation from agent, use the following command:
|
||||
```bash
|
||||
./build/cocos-cli attestation get '<report_data>'
|
||||
./build/cocos-cli agent attestation get '<report_data>'
|
||||
```
|
||||
|
||||
#### Validate attestation
|
||||
Validates the retrieved attestation information against a specified policy and checks its authenticity.
|
||||
To validate and verify attestation from agent, use the following command:
|
||||
```bash
|
||||
./build/cocos-cli attestation validate '<attestation>' --report_data '<report_data>'
|
||||
./build/cocos-cli agent attestation validate '<attestation>' --report_data '<report_data>'
|
||||
```
|
||||
##### Flags
|
||||
- --config: Path to a JSON file containing the validation configuration. This can be used to override individual flags.
|
||||
@@ -62,60 +70,21 @@ To validate and verify attestation from agent, use the following command:
|
||||
To upload an algorithm, use the following command:
|
||||
|
||||
```bash
|
||||
./build/cocos-cli algo /path/to/algorithm <private_key_file_path>
|
||||
./build/cocos-cli agent algo /path/to/algorithm
|
||||
```
|
||||
|
||||
##### Flags
|
||||
- -a, --algorithm string Algorithm type to run (default "bin")
|
||||
- --args stringArray Arguments to pass to the algorithm
|
||||
- --python-runtime string Python runtime to use (default "python3")
|
||||
- -r, --requirements string Python requirements file
|
||||
|
||||
#### Upload Dataset
|
||||
|
||||
To upload a dataset, use the following command:
|
||||
|
||||
```bash
|
||||
./build/cocos-cli data /path/to/dataset.csv <private_key_file_path>
|
||||
./build/cocos-cli agent data /path/to/dataset.csv
|
||||
```
|
||||
|
||||
Users can also upload directories which will be compressed on transit. Once received by agent they will be stored as compressed files or decompressed if the user passed the decompression argument.
|
||||
|
||||
##### Flags
|
||||
- -d, --decompress Decompress the dataset on agent
|
||||
|
||||
|
||||
|
||||
#### Retrieve result
|
||||
|
||||
To retrieve the computation result, use the following command:
|
||||
|
||||
```bash
|
||||
./build/cocos-cli result <private_key_file_path>
|
||||
```
|
||||
|
||||
#### Checksum
|
||||
When defining the manifest dataset and algorithm checksums are required. This can be done as below:
|
||||
|
||||
```bash
|
||||
./build/cocos-cli checksum <path_to_dataset_or_algorithm>
|
||||
```
|
||||
|
||||
#### Measure IGVM file
|
||||
We assume that our current working directory is the root of the cocos repository, both on the host machine and in the VM.
|
||||
|
||||
`igvmmeasure` calculates the launch measurement for an IGVM file and can generate a signed version. It ensures integrity by precomputing the expected launch digest, which can be verified against the attestation report. The tool parses IGVM directives, outputs the measurement as a hex string, or creates a signed file for verification at guest launch.
|
||||
|
||||
##### Example
|
||||
We measure an IGVM file using our measure command, run:
|
||||
|
||||
```bash
|
||||
./build/cocos-cli igvmmeasure /path/to/igvm/file
|
||||
```
|
||||
|
||||
The tool will parse the directives in the IGVM file, calculate the launch measurement, and output the computed digest. If successful, it prints the measurement to standard output.
|
||||
|
||||
Here is a sample output
|
||||
```
|
||||
91c4929bec2d0ecf11a708e09f0a57d7d82208bcba2451564444a4b01c22d047995ca27f9053f86de4e8063e9f810548
|
||||
./build/cocos-cli agent result
|
||||
```
|
||||
@@ -1,180 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/rand"
|
||||
"crypto/rsa"
|
||||
"crypto/x509"
|
||||
"encoding/pem"
|
||||
"errors"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
|
||||
)
|
||||
|
||||
const algorithmFile = "test_algo_file.py"
|
||||
|
||||
func generateRSAPrivateKeyFile(fileName string) error {
|
||||
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
privKeyFile, err := os.Create(fileName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer privKeyFile.Close()
|
||||
|
||||
privateKeyPEM := &pem.Block{
|
||||
Type: rsaKeyType,
|
||||
Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
|
||||
}
|
||||
|
||||
err = pem.Encode(privKeyFile, privateKeyPEM)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAlgorithmCmd(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setupMock func(*mocks.SDK)
|
||||
setupFiles func() error
|
||||
args []string
|
||||
connectErr error
|
||||
expectedOutput string
|
||||
cleanup func()
|
||||
}{
|
||||
{
|
||||
name: "successful upload",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() error {
|
||||
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
return generateRSAPrivateKeyFile(privateKeyFile)
|
||||
},
|
||||
args: []string{algorithmFile, privateKeyFile},
|
||||
expectedOutput: "Successfully uploaded algorithm",
|
||||
cleanup: func() {
|
||||
os.Remove(privateKeyFile)
|
||||
os.Remove(algorithmFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing algorithm file",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
args: []string{"non_existent_algo_file.py", privateKeyFile},
|
||||
expectedOutput: "Error reading algorithm file",
|
||||
},
|
||||
{
|
||||
name: "missing private key file",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() error {
|
||||
return os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
|
||||
},
|
||||
args: []string{algorithmFile, "non_existent_private_key.pem"},
|
||||
expectedOutput: "Error reading private key file",
|
||||
cleanup: func() {
|
||||
os.Remove(algorithmFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "upload failure",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
|
||||
},
|
||||
setupFiles: func() error {
|
||||
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
return generateRSAPrivateKeyFile(privateKeyFile)
|
||||
},
|
||||
args: []string{algorithmFile, privateKeyFile},
|
||||
expectedOutput: "Failed to upload algorithm",
|
||||
cleanup: func() {
|
||||
os.Remove(privateKeyFile)
|
||||
os.Remove(algorithmFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "invalid private key",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() error {
|
||||
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
privKeyFile, err := os.Create(privateKeyFile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer privKeyFile.Close()
|
||||
_, err = privKeyFile.WriteString("invalid private key")
|
||||
return err
|
||||
},
|
||||
args: []string{algorithmFile, privateKeyFile},
|
||||
expectedOutput: "Error decoding private key",
|
||||
cleanup: func() {
|
||||
os.Remove(algorithmFile)
|
||||
os.Remove(privateKeyFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "connection error",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
},
|
||||
connectErr: errors.New("failed to connect to agent"),
|
||||
args: []string{algorithmFile, privateKeyFile},
|
||||
expectedOutput: "Failed to connect to agent",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
if tt.setupMock != nil {
|
||||
tt.setupMock(mockSDK)
|
||||
}
|
||||
|
||||
testCLI := CLI{
|
||||
agentSDK: mockSDK,
|
||||
connectErr: tt.connectErr,
|
||||
}
|
||||
|
||||
if tt.setupFiles != nil {
|
||||
err := tt.setupFiles()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
cmd := testCLI.NewAlgorithmCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs(tt.args)
|
||||
err := cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), tt.expectedOutput)
|
||||
|
||||
if tt.cleanup != nil {
|
||||
tt.cleanup()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+16
-70
@@ -3,94 +3,40 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/pem"
|
||||
"log"
|
||||
"os"
|
||||
|
||||
"github.com/fatih/color"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm"
|
||||
"github.com/ultravioletrs/cocos/agent/algorithm/python"
|
||||
"google.golang.org/grpc/metadata"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
)
|
||||
|
||||
var (
|
||||
pythonRuntime string
|
||||
algoType string
|
||||
requirementsFile string
|
||||
algoArgs []string
|
||||
)
|
||||
|
||||
func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
func (cli *CLI) NewAlgorithmsCmd() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "algo",
|
||||
Short: "Upload an algorithm binary",
|
||||
Example: "algo <algo_file> <private_key_file_path>",
|
||||
Args: cobra.ExactArgs(2),
|
||||
Example: "algo <algo_file> <id> <provider>",
|
||||
Args: cobra.ExactArgs(3),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
if cli.connectErr != nil {
|
||||
cli.printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
|
||||
return
|
||||
}
|
||||
|
||||
algorithmFile := args[0]
|
||||
|
||||
cmd.Println("Uploading algorithm file:", algorithmFile)
|
||||
log.Println("Uploading algorithm binary:", algorithmFile)
|
||||
|
||||
algorithm, err := os.Open(algorithmFile)
|
||||
algorithm, err := os.ReadFile(algorithmFile)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Error reading algorithm file: %v ❌ ", err)
|
||||
return
|
||||
log.Fatalf("Error reading algorithm file: %v", err)
|
||||
}
|
||||
|
||||
defer algorithm.Close()
|
||||
|
||||
var req *os.File
|
||||
if requirementsFile != "" {
|
||||
req, err = os.Open(requirementsFile)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Error reading requirments file: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
defer req.Close()
|
||||
algoReq := agent.Algorithm{
|
||||
Algorithm: algorithm,
|
||||
ID: args[1],
|
||||
Provider: args[2],
|
||||
}
|
||||
|
||||
privKeyFile, err := os.ReadFile(args[1])
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Error reading private key file: %v ❌ ", err)
|
||||
return
|
||||
if err := cli.agentSDK.Algo(cmd.Context(), algoReq); err != nil {
|
||||
log.Fatalf("Error uploading algorithm with ID %s and provider %s: %v", algoReq.ID, algoReq.Provider, err)
|
||||
}
|
||||
|
||||
pemBlock, _ := pem.Decode(privKeyFile)
|
||||
|
||||
privKey, err := decodeKey(pemBlock)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Error decoding private key: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
|
||||
ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
|
||||
|
||||
if err := cli.agentSDK.Algo(addAlgoMetadata(ctx), algorithm, req, privKey); err != nil {
|
||||
cli.printError(cmd, "Failed to upload algorithm due to error: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
|
||||
cmd.Println(color.New(color.FgGreen).Sprint("Successfully uploaded algorithm! ✔ "))
|
||||
log.Println("Successfully uploaded algorithm")
|
||||
},
|
||||
}
|
||||
|
||||
cmd.Flags().StringVarP(&algoType, "algorithm", "a", string(algorithm.AlgoTypeBin), "Algorithm type to run")
|
||||
cmd.Flags().StringVar(&pythonRuntime, "python-runtime", python.PyRuntime, "Python runtime to use")
|
||||
cmd.Flags().StringVarP(&requirementsFile, "requirements", "r", "", "Python requirements file")
|
||||
cmd.Flags().StringArrayVar(&algoArgs, "args", []string{}, "Arguments to pass to the algorithm")
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
func addAlgoMetadata(ctx context.Context) context.Context {
|
||||
ctx = algorithm.AlgorithmTypeToContext(ctx, algoType)
|
||||
ctx = algorithm.AlgorithmArgsToContext(ctx, algoArgs)
|
||||
ctx = python.PythonRunTimeToContext(ctx, pythonRuntime)
|
||||
return ctx
|
||||
}
|
||||
|
||||
+555
-288
@@ -3,52 +3,115 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/absmach/magistrala/pkg/errors"
|
||||
"github.com/fatih/color"
|
||||
"github.com/google/go-sev-guest/abi"
|
||||
tpmAttest "github.com/google/go-tpm-tools/proto/attest"
|
||||
"github.com/google/go-sev-guest/proto/check"
|
||||
"github.com/google/go-sev-guest/validate"
|
||||
"github.com/google/go-sev-guest/verify"
|
||||
"github.com/google/go-sev-guest/verify/trust"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/spf13/pflag"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation"
|
||||
"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
|
||||
"google.golang.org/protobuf/encoding/prototext"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"google.golang.org/protobuf/encoding/protojson"
|
||||
"google.golang.org/protobuf/types/known/wrapperspb"
|
||||
)
|
||||
|
||||
const (
|
||||
size8 = 8
|
||||
size16 = 16
|
||||
size32 = 32
|
||||
size48 = 48
|
||||
size64 = 64
|
||||
attestationFilePath = "attestation.bin"
|
||||
azureAttestResultFilePath = "azure_attest_result.json"
|
||||
azureAttestTokenFilePath = "azure_attest_token.jwt"
|
||||
attestationReportJson = "attestation.json"
|
||||
TEE = "tee"
|
||||
SNP = "snp"
|
||||
VTPM = "vtpm"
|
||||
SNPvTPM = "snp-vtpm"
|
||||
AzureToken = "azure-token"
|
||||
CCNone = "none"
|
||||
CCAzure = "azure"
|
||||
CCGCP = "gcp"
|
||||
TDX = "tdx"
|
||||
defaultMinimumTcb = 0
|
||||
defaultMinimumLaunchTcb = 0
|
||||
defaultMinimumGuestSvn = 0
|
||||
defaultGuestPolicy = 0x0000000000030000
|
||||
defaultMinimumBuild = 0
|
||||
defaultCheckCrl = false
|
||||
defaultTimeout = 2 * time.Minute
|
||||
defaultMaxRetryDelay = 30 * time.Second
|
||||
defaultRequireAuthor = false
|
||||
defaultRequireIdBlock = false
|
||||
defaultMinVersion = "0.0"
|
||||
size16 = 16
|
||||
size32 = 32
|
||||
size48 = 48
|
||||
size64 = 64
|
||||
attestationFilePath = "attestation.bin"
|
||||
exampleJSONConfig = `
|
||||
{
|
||||
"rootOfTrust":{
|
||||
"product":"test_product",
|
||||
"cabundlePaths":[
|
||||
"test_cabundlePaths"
|
||||
],
|
||||
"cabundles":[
|
||||
"test_Cabundles"
|
||||
],
|
||||
"checkCrl":true,
|
||||
"disallowNetwork":true
|
||||
},
|
||||
"policy":{
|
||||
"minimumGuestSvn":1,
|
||||
"policy":"1",
|
||||
"familyId":"AQIDBAUGBwgJCgsMDQ4PEA==",
|
||||
"imageId":"AQIDBAUGBwgJCgsMDQ4PEA==",
|
||||
"vmpl":0,
|
||||
"minimumTcb":"1",
|
||||
"minimumLaunchTcb":"1",
|
||||
"platformInfo":"1",
|
||||
"requireAuthorKey":true,
|
||||
"reportData":"J+60aXs8btm8VcGgaJYURGeNCu0FIyWMFXQ7ZUlJDC0FJGJizJsOzDIXgQ75UtPC+Zqe0A3dvnnf5VEeQ61RTg==",
|
||||
"measurement":"8s78ewoX7Xkfy1qsgVnkZwLDotD768Nqt6qTL5wtQOxHsLczipKM6bhDmWiHLdP4",
|
||||
"hostData":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
|
||||
"reportId":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
|
||||
"reportIdMa":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
|
||||
"chipId":"J+60aXs8btm8VcGgaJYURGeNCu0FIyWMFXQ7ZUlJDC0FJGJizJsOzDIXgQ75UtPC+Zqe0A3dvnnf5VEeQ61RTg==",
|
||||
"minimumBuild":1,
|
||||
"minimumVersion":"0.90",
|
||||
"permitProvisionalFirmware":true,
|
||||
"requireIdBlock":true,
|
||||
"trustedAuthorKeys":[
|
||||
"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
|
||||
],
|
||||
"trustedAuthorKeyHashes":[
|
||||
"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
|
||||
],
|
||||
"trustedIdKeys":[
|
||||
"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
|
||||
],
|
||||
"trustedIdKeyHashes":[
|
||||
"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
|
||||
],
|
||||
"product":{
|
||||
"name":"SEV_PRODUCT_MILAN",
|
||||
"stepping":1,
|
||||
"machineStepping":1
|
||||
}
|
||||
}
|
||||
}
|
||||
`
|
||||
)
|
||||
|
||||
var (
|
||||
errReportSize = errors.New("attestation contents too small")
|
||||
nonce []byte
|
||||
teeNonce []byte
|
||||
tokenNonce []byte
|
||||
getTextProtoAttestationReport bool
|
||||
getAzureTokenJWT 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}
|
||||
)
|
||||
|
||||
func (cli *CLI) NewAttestationCmd() *cobra.Command {
|
||||
@@ -56,9 +119,9 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
|
||||
Use: "attestation [command]",
|
||||
Short: "Get and validate attestations",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
cmd.Printf("Get and validate attestations\n\n")
|
||||
cmd.Printf("Usage:\n %s [command]\n\n", cmd.CommandPath())
|
||||
cmd.Printf("Available Commands:\n")
|
||||
fmt.Printf("Get and validate attestations\n\n")
|
||||
fmt.Printf("Usage:\n %s [command]\n\n", cmd.CommandPath())
|
||||
fmt.Printf("Available Commands:\n")
|
||||
|
||||
// Filter out "completion" command
|
||||
availableCommands := make([]*cobra.Command, 0)
|
||||
@@ -69,298 +132,502 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
|
||||
}
|
||||
|
||||
for _, subCmd := range availableCommands {
|
||||
cmd.Printf(" %-15s%s\n", subCmd.Name(), subCmd.Short)
|
||||
fmt.Printf(" %-15s%s\n", subCmd.Name(), subCmd.Short)
|
||||
}
|
||||
|
||||
cmd.Printf("\nFlags:\n")
|
||||
fmt.Printf("\nFlags:\n")
|
||||
cmd.Flags().VisitAll(func(flag *pflag.Flag) {
|
||||
cmd.Printf(" -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
|
||||
fmt.Printf(" -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
|
||||
})
|
||||
cmd.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
|
||||
fmt.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "get",
|
||||
Short: "Retrieve attestation information from agent. The argument of the command must be the type of the report (snp or vtpm or snp-vtpm or tdx).",
|
||||
ValidArgs: []cobra.Completion{SNP, VTPM, SNPvTPM, AzureToken, TDX},
|
||||
Example: fmt.Sprintf(`Based on attestation report type:
|
||||
get %s --tee <512 bit hex value>
|
||||
get %s --vtpm <256 bit hex value>
|
||||
get %s --tee <512 bit hex value> --vtpm <256 bit hex value>
|
||||
get %s --token <256 bit hex value>
|
||||
get %s --tee <512 bit hex value>`, SNP, VTPM, SNPvTPM, AzureToken, TDX),
|
||||
Args: cobra.ExactArgs(1),
|
||||
return &cobra.Command{
|
||||
Use: "get",
|
||||
Short: "Retrieve attestation information from agent. Report data expected in hex enoded string of length 64 bytes.",
|
||||
Example: "report <report_data>",
|
||||
Args: cobra.ExactArgs(1),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
if cli.connectErr != nil {
|
||||
cli.printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
|
||||
return
|
||||
}
|
||||
log.Println("Getting attestation")
|
||||
|
||||
if err := cobra.OnlyValidArgs(cmd, args); err != nil {
|
||||
cli.printError(cmd, "Bad attestation type: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
|
||||
attestationType := args[0]
|
||||
|
||||
attType := attestation.SNP
|
||||
switch attestationType {
|
||||
case SNP:
|
||||
cmd.Println("Fetching SEV-SNP attestation report")
|
||||
case VTPM:
|
||||
cmd.Println("Fetching vTPM report")
|
||||
attType = attestation.VTPM
|
||||
case SNPvTPM:
|
||||
cmd.Println("Fetching SEV-SNP and vTPM report")
|
||||
attType = attestation.SNPvTPM
|
||||
case AzureToken:
|
||||
cmd.Println("Fetching Azure token")
|
||||
case TDX:
|
||||
cmd.Println("Fetching TDX attestation report")
|
||||
attType = attestation.TDX
|
||||
}
|
||||
|
||||
if (attestationType == VTPM || attestationType == SNPvTPM) && len(nonce) == 0 {
|
||||
msg := color.New(color.FgRed).Sprint("vTPM nonce must be defined for vTPM attestation ❌ ")
|
||||
cmd.Println(msg)
|
||||
return
|
||||
}
|
||||
|
||||
if (attestationType == SNP || attestationType == SNPvTPM) && len(teeNonce) == 0 {
|
||||
msg := color.New(color.FgRed).Sprint("TEE nonce must be defined for SEV-SNP attestation ❌ ")
|
||||
cmd.Println(msg)
|
||||
return
|
||||
}
|
||||
|
||||
if (attestationType == AzureToken) && len(tokenNonce) == 0 {
|
||||
msg := color.New(color.FgRed).Sprint("Token nonce must be defined for Azure attestation ❌ ")
|
||||
cmd.Println(msg)
|
||||
return
|
||||
}
|
||||
|
||||
var fixedReportData [vtpm.SEVNonce]byte
|
||||
if attType == attestation.SNP || attType == attestation.SNPvTPM {
|
||||
if len(teeNonce) > vtpm.SEVNonce {
|
||||
msg := color.New(color.FgRed).Sprintf("nonce must be a hex encoded string of length lesser or equal %d bytes ❌ ", vtpm.SEVNonce)
|
||||
cmd.Println(msg)
|
||||
return
|
||||
}
|
||||
|
||||
copy(fixedReportData[:], teeNonce)
|
||||
}
|
||||
|
||||
var fixedVtpmNonceByte [vtpm.Nonce]byte
|
||||
if attType != attestation.SNP || attestationType == AzureToken {
|
||||
if (len(nonce) > vtpm.Nonce) || (len(tokenNonce) > vtpm.Nonce) {
|
||||
msg := color.New(color.FgRed).Sprintf("vTPM nonce must be a hex encoded string of length lesser or equal %d bytes ❌ ", vtpm.Nonce)
|
||||
cmd.Println(msg)
|
||||
return
|
||||
}
|
||||
if attestationType == AzureToken {
|
||||
copy(fixedVtpmNonceByte[:], tokenNonce)
|
||||
} else {
|
||||
copy(fixedVtpmNonceByte[:], nonce)
|
||||
}
|
||||
}
|
||||
|
||||
filename := attestationFilePath
|
||||
|
||||
if attestationType == AzureToken {
|
||||
filename = azureAttestResultFilePath
|
||||
}
|
||||
|
||||
if getTextProtoAttestationReport {
|
||||
filename = attestationReportJson
|
||||
} else if getAzureTokenJWT {
|
||||
filename = azureAttestTokenFilePath
|
||||
}
|
||||
|
||||
attestationFile, err := os.Create(filename)
|
||||
reportData, err := hex.DecodeString(args[0])
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Error creating attestation file: %v ❌ ", err)
|
||||
return
|
||||
log.Fatalf("attestation validation and verification failed with error: %s", err)
|
||||
}
|
||||
if len(reportData) != agent.ReportDataSize {
|
||||
log.Fatalf("report data must be a hex encoded string of length %d bytes", agent.ReportDataSize)
|
||||
}
|
||||
|
||||
var returnJsonAzureToken bool
|
||||
|
||||
if attestationType == AzureToken {
|
||||
err := cli.agentSDK.AttestationToken(cmd.Context(), fixedVtpmNonceByte, int(attType), attestationFile)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Failed to get attestation token due to error: %v ❌", err)
|
||||
return
|
||||
}
|
||||
returnJsonAzureToken = !getAzureTokenJWT
|
||||
} else {
|
||||
err := cli.agentSDK.Attestation(cmd.Context(), fixedReportData, fixedVtpmNonceByte, int(attType), attestationFile)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Failed to get attestation due to error: %v ❌", err)
|
||||
return
|
||||
}
|
||||
result, err := cli.agentSDK.Attestation(cmd.Context(), [agent.ReportDataSize]byte(reportData))
|
||||
if err != nil {
|
||||
log.Fatalf("Error retrieving attestation: %v", err)
|
||||
}
|
||||
|
||||
if err := attestationFile.Close(); err != nil {
|
||||
cli.printError(cmd, "Error closing attestation file: %v ❌ ", err)
|
||||
return
|
||||
if err = os.WriteFile(attestationFilePath, result, 0o644); err != nil {
|
||||
log.Fatalf("Error saving attestation result: %v", err)
|
||||
}
|
||||
|
||||
if getTextProtoAttestationReport || returnJsonAzureToken {
|
||||
result, err := os.ReadFile(filename)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Error reading attestation file: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
|
||||
switch attestationType {
|
||||
case SNP:
|
||||
result, err = attestationToJSON(result)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Error converting SNP attestation to JSON: %v ❌", err)
|
||||
return
|
||||
}
|
||||
|
||||
case VTPM, SNPvTPM:
|
||||
marshalOptions := prototext.MarshalOptions{
|
||||
Multiline: true,
|
||||
EmitASCII: true,
|
||||
}
|
||||
var attvTPM tpmAttest.Attestation
|
||||
err = proto.Unmarshal(result, &attvTPM)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Failed to unmarshal the attestation report: %v ❌", err)
|
||||
return
|
||||
}
|
||||
result = []byte(marshalOptions.Format(&attvTPM))
|
||||
|
||||
case AzureToken:
|
||||
result, err = decodeJWTToJSON(result)
|
||||
if err != nil {
|
||||
cli.printError(cmd, "Error decoding Azure token: %v ❌", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err := os.WriteFile(filename, result, 0o644); err != nil {
|
||||
cli.printError(cmd, "Error writing attestation file: %v ❌ ", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
cmd.Println("Attestation retrieved and saved successfully!")
|
||||
log.Println("Attestation result retrieved and saved successfully!")
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
cmd.Flags().BoolVarP(&getAzureTokenJWT, "azurejwt", "t", false, "Get azure attestation token as jwt format")
|
||||
cmd.Flags().BoolVarP(&getTextProtoAttestationReport, "reporttextproto", "r", false, "Get attestation report in textproto format")
|
||||
cmd.Flags().BytesHexVar(&teeNonce, "tee", []byte{}, "Define the nonce for the SNP and TDX attestation report (must be used with attestation type snp, snp-vtpm, and tdx)")
|
||||
cmd.Flags().BytesHexVar(&nonce, "vtpm", []byte{}, "Define the nonce for the vTPM attestation report (must be used with attestation type vtpm and snp-vtpm)")
|
||||
cmd.Flags().BytesHexVar(&tokenNonce, "token", []byte{}, "Define the nonce for the Azure attestation token (must be used with attestation type azure-token)")
|
||||
func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "validate",
|
||||
Short: "Validate and verify attestation information. The report is provided as a file path.",
|
||||
Example: "validate <attestation_report_file_path>",
|
||||
Args: cobra.ExactArgs(1),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
log.Println("Checking attestation")
|
||||
|
||||
attestationFile = string(args[0])
|
||||
|
||||
if err := parseConfig(); err != nil {
|
||||
log.Fatalf("attestation validation and verification failed with error: %s", err)
|
||||
}
|
||||
if err := parseHashes(); err != nil {
|
||||
log.Fatalf("attestation validation and verification failed with error: %s", err)
|
||||
}
|
||||
if err := parseFiles(); err != nil {
|
||||
log.Fatalf("attestation validation and verification failed with error: %s", err)
|
||||
}
|
||||
// This format is the attestation report in AMD's specified ABI format, immediately
|
||||
// followed by the certificate table bytes.
|
||||
if len(attestation) < abi.ReportSize {
|
||||
log.Fatalf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes", len(attestation), abi.ReportSize)
|
||||
}
|
||||
if err := parseUints(); err != nil {
|
||||
log.Fatalf("attestation validation and verification failed with error: %s", err)
|
||||
}
|
||||
cfg.Policy.Vmpl = wrapperspb.UInt32(0)
|
||||
|
||||
if err := validateInput(); err != nil {
|
||||
log.Fatalf("attestation validation and verification failed with error: %s", err)
|
||||
}
|
||||
|
||||
if err := verifyAndValidateAttestation(attestation); err != nil {
|
||||
log.Fatalf("attestation validation and verification failed with error: %s", err)
|
||||
}
|
||||
log.Println("Attestation validation and verification is successful!")
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(
|
||||
&cfgString,
|
||||
"config",
|
||||
"",
|
||||
"Serialized json check.Config protobuf. This will overwrite individual flags. Unmarshalled as json. Example: "+exampleJSONConfig,
|
||||
)
|
||||
cmd.Flags().BytesHexVar(
|
||||
&cfg.Policy.HostData,
|
||||
"host_data",
|
||||
empty32[:],
|
||||
"The expected HOST_DATA field as a hex string. Must encode 32 bytes. Unchecked if unset.",
|
||||
)
|
||||
cmd.Flags().BytesHexVar(
|
||||
&cfg.Policy.ReportData,
|
||||
"report_data",
|
||||
empty64[:],
|
||||
"The expected REPORT_DATA field as a hex string. Must encode 64 bytes. Must be set.",
|
||||
)
|
||||
cmd.Flags().BytesHexVar(
|
||||
&cfg.Policy.FamilyId,
|
||||
"family_id",
|
||||
empty16[:],
|
||||
"The expected FAMILY_ID field as a hex string. Must encode 16 bytes. Unchecked if unset.",
|
||||
)
|
||||
cmd.Flags().BytesHexVar(
|
||||
&cfg.Policy.ImageId,
|
||||
"image_id",
|
||||
empty16[:],
|
||||
"The expected IMAGE_ID field as a hex string. Must encode 16 bytes. Unchecked if unset.",
|
||||
)
|
||||
cmd.Flags().BytesHexVar(
|
||||
&cfg.Policy.ReportId,
|
||||
"report_id",
|
||||
nil,
|
||||
"The expected REPORT_ID field as a hex string. Must encode 32 bytes. Unchecked if unset.",
|
||||
)
|
||||
cmd.Flags().BytesHexVar(
|
||||
&cfg.Policy.ReportIdMa,
|
||||
"report_id_ma",
|
||||
defaultReportIdMa,
|
||||
"The expected REPORT_ID_MA field as a hex string. Must encode 32 bytes. Unchecked if unset.",
|
||||
)
|
||||
cmd.Flags().BytesHexVar(
|
||||
&cfg.Policy.Measurement,
|
||||
"measurement",
|
||||
nil,
|
||||
"The expected MEASUREMENT field as a hex string. Must encode 48 bytes. Unchecked if unset.",
|
||||
)
|
||||
cmd.Flags().BytesHexVar(
|
||||
&cfg.Policy.ChipId,
|
||||
"chip_id",
|
||||
nil,
|
||||
"The expected MEASUREMENT field as a hex string. Must encode 48 bytes. Unchecked if unset.",
|
||||
)
|
||||
cmd.Flags().Uint64Var(
|
||||
&cfg.Policy.MinimumTcb,
|
||||
"minimum_tcb",
|
||||
defaultMinimumTcb,
|
||||
"The minimum acceptable value for CURRENT_TCB, COMMITTED_TCB, and REPORTED_TCB.",
|
||||
)
|
||||
cmd.Flags().Uint64Var(
|
||||
&cfg.Policy.MinimumLaunchTcb,
|
||||
"minimum_lauch_tcb",
|
||||
defaultMinimumLaunchTcb,
|
||||
"The minimum acceptable value for LAUNCH_TCB.",
|
||||
)
|
||||
cmd.Flags().Uint64Var(
|
||||
&cfg.Policy.Policy,
|
||||
"guest_policy",
|
||||
defaultGuestPolicy,
|
||||
"The most acceptable guest SnpPolicy.",
|
||||
)
|
||||
cmd.Flags().Uint32Var(
|
||||
&cfg.Policy.MinimumGuestSvn,
|
||||
"minimum_guest_svn",
|
||||
defaultMinimumGuestSvn,
|
||||
"The most acceptable GUEST_SVN.",
|
||||
)
|
||||
cmd.Flags().Uint32Var(
|
||||
&cfg.Policy.MinimumBuild,
|
||||
"minimum_build",
|
||||
defaultMinimumBuild,
|
||||
"The 8-bit minimum build number for AMD-SP firmware",
|
||||
)
|
||||
cmd.Flags().BoolVar(
|
||||
&cfg.RootOfTrust.CheckCrl,
|
||||
"check_crl",
|
||||
defaultCheckCrl,
|
||||
"Download and check the CRL for revoked certificates.",
|
||||
)
|
||||
cmd.Flags().DurationVar(
|
||||
&timeout,
|
||||
"timeout",
|
||||
defaultTimeout,
|
||||
"Duration to continue to retry failed HTTP requests.",
|
||||
)
|
||||
cmd.Flags().DurationVar(
|
||||
&maxRetryDelay,
|
||||
"max_retry_delay",
|
||||
defaultMaxRetryDelay,
|
||||
"Maximum Duration to wait between HTTP request retries.",
|
||||
)
|
||||
cmd.Flags().BoolVar(
|
||||
&cfg.Policy.RequireAuthorKey,
|
||||
"require_author_key",
|
||||
defaultRequireAuthor,
|
||||
"Require that AUTHOR_KEY_EN is 1.",
|
||||
)
|
||||
cmd.Flags().BoolVar(
|
||||
&cfg.Policy.RequireIdBlock,
|
||||
"require_id_block",
|
||||
defaultRequireIdBlock,
|
||||
"Require that the VM was launch with an ID_BLOCK signed by a trusted id key or author key",
|
||||
)
|
||||
cmd.Flags().StringVar(
|
||||
&platformInfo,
|
||||
"platform_info",
|
||||
"",
|
||||
"The maximum acceptable PLATFORM_INFO field bit-wise. May be empty or a 64-bit unsigned integer",
|
||||
)
|
||||
cmd.Flags().StringVar(
|
||||
&cfg.Policy.MinimumVersion,
|
||||
"minimum_version",
|
||||
defaultMinVersion,
|
||||
"Minimum AMD-SP firmware API version (major.minor). Each number must be 8-bit non-negative.",
|
||||
)
|
||||
cmd.Flags().StringArrayVar(
|
||||
&trustedAuthorKeys,
|
||||
"trusted_author_keys",
|
||||
[]string{},
|
||||
"Paths to x.509 certificates of trusted author keys",
|
||||
)
|
||||
cmd.Flags().StringArrayVar(
|
||||
&trustedAuthorHashes,
|
||||
"trusted_author_key_hashes",
|
||||
[]string{},
|
||||
"Hex-encoded SHA-384 hash values of trusted author keys in AMD public key format",
|
||||
)
|
||||
cmd.Flags().StringArrayVar(
|
||||
&trustedIdKeys,
|
||||
"trusted_id_keys",
|
||||
[]string{},
|
||||
"Paths to x.509 certificates of trusted author keys",
|
||||
)
|
||||
cmd.Flags().StringArrayVar(
|
||||
&trustedIdKeyHashes,
|
||||
"trusted_id_key_hashes",
|
||||
[]string{},
|
||||
"Hex-encoded SHA-384 hash values of trusted identity keys in AMD public key format",
|
||||
)
|
||||
cmd.Flags().StringVar(
|
||||
&cfg.RootOfTrust.Product,
|
||||
"product",
|
||||
"",
|
||||
"The AMD product name for the chip that generated the attestation report.",
|
||||
)
|
||||
cmd.Flags().StringVar(
|
||||
&stepping,
|
||||
"stepping",
|
||||
"",
|
||||
"The machine stepping for the chip that generated the attestation report. Default unchecked.",
|
||||
)
|
||||
cmd.Flags().StringArrayVar(
|
||||
&cfg.RootOfTrust.CabundlePaths,
|
||||
"CA_bundles_paths",
|
||||
[]string{},
|
||||
"Paths to CA bundles for the AMD product. Must be in PEM format, ASK, then ARK certificates. If unset, uses embedded root certificates.",
|
||||
)
|
||||
cmd.Flags().StringArrayVar(
|
||||
&cfg.RootOfTrust.Cabundles,
|
||||
"CA_bundles",
|
||||
[]string{},
|
||||
"PEM format CA bundles for the AMD product. Combined with contents of cabundle_paths.",
|
||||
)
|
||||
|
||||
if err := cmd.MarkFlagRequired("report_data"); err != nil {
|
||||
log.Fatalf("Failed to mark flag as required: %s", err)
|
||||
}
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
func attestationToJSON(report []byte) ([]byte, error) {
|
||||
if len(report) < abi.ReportSize {
|
||||
return nil, errors.Wrap(errReportSize, fmt.Errorf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes", len(report), abi.ReportSize))
|
||||
}
|
||||
attestationPB, err := abi.ReportCertsToProto(report[:abi.ReportSize])
|
||||
func verifyAndValidateAttestation(attestation []byte) error {
|
||||
sopts, err := verify.RootOfTrustToOptions(cfg.RootOfTrust)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return err
|
||||
}
|
||||
|
||||
return json.MarshalIndent(attestationPB, "", " ")
|
||||
sopts.Product = cfg.Policy.Product
|
||||
sopts.Getter = &trust.RetryHTTPSGetter{
|
||||
Timeout: timeout,
|
||||
MaxRetryDelay: maxRetryDelay,
|
||||
Getter: &trust.SimpleHTTPSGetter{},
|
||||
}
|
||||
attestationPB, err := abi.ReportCertsToProto(attestation)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = verify.SnpAttestation(attestationPB, sopts); err != nil {
|
||||
return err
|
||||
}
|
||||
opts, err := validate.PolicyToOptions(cfg.Policy)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = validate.SnpAttestation(attestationPB, opts); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "validate",
|
||||
Short: "Validate and verify attestation information (Deprecated)",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
cmd.Println("Validation via CLI using legacy policies is deprecated. Please use CoRIM tools.")
|
||||
},
|
||||
// parseConfig decodes config passed as json for check.Config struct.
|
||||
// example
|
||||
/* {
|
||||
"rootOfTrust":{
|
||||
"product":"test_product",
|
||||
"cabundlePaths":[
|
||||
"test_cabundlePaths"
|
||||
],
|
||||
"cabundles":[
|
||||
"test_Cabundles"
|
||||
],
|
||||
"checkCrl":true,
|
||||
"disallowNetwork":true
|
||||
},
|
||||
"policy":{
|
||||
"minimumGuestSvn":1,
|
||||
"policy":"1",
|
||||
"familyId":"AQIDBAUGBwgJCgsMDQ4PEA==",
|
||||
"imageId":"AQIDBAUGBwgJCgsMDQ4PEA==",
|
||||
"vmpl":0,
|
||||
"minimumTcb":"1",
|
||||
"minimumLaunchTcb":"1",
|
||||
"platformInfo":"1",
|
||||
"requireAuthorKey":true,
|
||||
"reportData":"J+60aXs8btm8VcGgaJYURGeNCu0FIyWMFXQ7ZUlJDC0FJGJizJsOzDIXgQ75UtPC+Zqe0A3dvnnf5VEeQ61RTg==",
|
||||
"measurement":"8s78ewoX7Xkfy1qsgVnkZwLDotD768Nqt6qTL5wtQOxHsLczipKM6bhDmWiHLdP4",
|
||||
"hostData":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
|
||||
"reportId":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
|
||||
"reportIdMa":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
|
||||
"chipId":"J+60aXs8btm8VcGgaJYURGeNCu0FIyWMFXQ7ZUlJDC0FJGJizJsOzDIXgQ75UtPC+Zqe0A3dvnnf5VEeQ61RTg==",
|
||||
"minimumBuild":1,
|
||||
"minimumVersion":"0.90",
|
||||
"permitProvisionalFirmware":true,
|
||||
"requireIdBlock":true,
|
||||
"trustedAuthorKeys":[
|
||||
"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
|
||||
],
|
||||
"trustedAuthorKeyHashes":[
|
||||
"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
|
||||
],
|
||||
"trustedIdKeys":[
|
||||
"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
|
||||
],
|
||||
"trustedIdKeyHashes":[
|
||||
"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
|
||||
],
|
||||
"product":{
|
||||
"name":"SEV_PRODUCT_MILAN",
|
||||
"stepping":1,
|
||||
"machineStepping":1
|
||||
}
|
||||
}
|
||||
}*/
|
||||
func parseConfig() error {
|
||||
if cfgString == "" {
|
||||
return nil
|
||||
}
|
||||
if err := protojson.Unmarshal([]byte(cfgString), &cfg); err != nil {
|
||||
return err
|
||||
}
|
||||
// Populate fields that should not be nil
|
||||
if cfg.RootOfTrust == nil {
|
||||
cfg.RootOfTrust = &check.RootOfTrust{}
|
||||
}
|
||||
if cfg.Policy == nil {
|
||||
cfg.Policy = &check.Policy{}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
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.`,
|
||||
func parseHashes() error {
|
||||
for _, hash := range trustedAuthorHashes {
|
||||
hashBytes, err := hex.DecodeString(hash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg.Policy.TrustedAuthorKeyHashes = append(cfg.Policy.TrustedAuthorKeyHashes, hashBytes)
|
||||
}
|
||||
for _, hash := range trustedIdKeyHashes {
|
||||
hashBytes, err := hex.DecodeString(hash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg.Policy.TrustedIdKeyHashes = append(cfg.Policy.TrustedIdKeyHashes, hashBytes)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
Args: cobra.MinimumNArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if len(args) == 0 {
|
||||
return fmt.Errorf("error: No input file provided")
|
||||
}
|
||||
func parseFiles() error {
|
||||
file, err := os.ReadFile(attestationFile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
attestation = file
|
||||
for _, path := range trustedAuthorKeys {
|
||||
file, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg.Policy.TrustedAuthorKeys = append(cfg.Policy.TrustedAuthorKeys, file)
|
||||
}
|
||||
for _, path := range trustedIdKeys {
|
||||
file, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg.Policy.TrustedIdKeys = append(cfg.Policy.TrustedIdKeys, file)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
inputFile := args[0]
|
||||
|
||||
measurement, err := cli.measurement.Run(inputFile)
|
||||
func parseUints() error {
|
||||
if stepping != "" {
|
||||
if base := getBase(stepping); base == 10 {
|
||||
num, err := strconv.ParseUint(stepping, getBase(stepping), 8)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
outputString := string(measurement)
|
||||
lines := strings.Split(strings.TrimSpace(outputString), "\n")
|
||||
|
||||
if len(lines) == 1 {
|
||||
outputString = strings.ToLower(outputString)
|
||||
} else {
|
||||
return fmt.Errorf("error: %s", outputString)
|
||||
cfg.Policy.Product.MachineStepping = wrapperspb.UInt32(uint32(num))
|
||||
} else {
|
||||
num, err := strconv.ParseUint(stepping[2:], base, 8)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cmd.Print(outputString)
|
||||
|
||||
return nil
|
||||
},
|
||||
cfg.Policy.Product.MachineStepping = wrapperspb.UInt32(uint32(num))
|
||||
}
|
||||
}
|
||||
|
||||
return igvmmeasureCmd
|
||||
if platformInfo != "" {
|
||||
if base := getBase(platformInfo); base == 10 {
|
||||
num, err := strconv.ParseUint(platformInfo, getBase(platformInfo), 8)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg.Policy.PlatformInfo = wrapperspb.UInt64(num)
|
||||
} else {
|
||||
num, err := strconv.ParseUint(platformInfo[2:], base, 8)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg.Policy.PlatformInfo = wrapperspb.UInt64(num)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeJWTToJSON(tokenBytes []byte) ([]byte, error) {
|
||||
token := string(tokenBytes) // convert to string
|
||||
parts := strings.Split(token, ".")
|
||||
if len(parts) < 2 {
|
||||
return nil, fmt.Errorf("invalid JWT: must have at least 2 parts")
|
||||
func getBase(val string) int {
|
||||
switch {
|
||||
case strings.HasPrefix(val, "0x"):
|
||||
return 16
|
||||
case strings.HasPrefix(val, "0o"):
|
||||
return 8
|
||||
case strings.HasPrefix(val, "0b"):
|
||||
return 2
|
||||
default:
|
||||
return 10
|
||||
}
|
||||
|
||||
decode := func(seg string) (map[string]any, error) {
|
||||
// Add padding if missing
|
||||
if m := len(seg) % 4; m != 0 {
|
||||
seg += strings.Repeat("=", 4-m)
|
||||
}
|
||||
|
||||
data, err := base64.URLEncoding.DecodeString(seg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal(data, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
header, err := decode(parts[0])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to decode header: %v", err)
|
||||
}
|
||||
|
||||
payload, err := decode(parts[1])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to decode payload: %v", err)
|
||||
}
|
||||
|
||||
combined := map[string]any{
|
||||
"header": header,
|
||||
"payload": payload,
|
||||
}
|
||||
|
||||
return json.MarshalIndent(combined, "", " ")
|
||||
}
|
||||
|
||||
func validateInput() error {
|
||||
if len(cfg.RootOfTrust.CabundlePaths) != 0 || len(cfg.RootOfTrust.Cabundles) != 0 && cfg.RootOfTrust.Product == "" {
|
||||
return fmt.Errorf("product name must be set if CA bundles are provided")
|
||||
}
|
||||
|
||||
if err := validateFieldLength("report_data", cfg.Policy.ReportData, size64); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateFieldLength("host_data", cfg.Policy.HostData, size32); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateFieldLength("family_id", cfg.Policy.FamilyId, size16); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateFieldLength("image_id", cfg.Policy.ImageId, size16); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateFieldLength("report_id", cfg.Policy.ReportId, size32); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateFieldLength("report_id_ma", cfg.Policy.ReportIdMa, size32); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateFieldLength("measurement", cfg.Policy.Measurement, size48); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateFieldLength("chip_id", cfg.Policy.ChipId, size64); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, hash := range cfg.Policy.TrustedAuthorKeyHashes {
|
||||
if err := validateFieldLength("trusted_author_key_hash", hash, size48); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for _, hash := range cfg.Policy.TrustedIdKeyHashes {
|
||||
if err := validateFieldLength("trusted_id_key_hash", hash, size48); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateFieldLength(fieldName string, field []byte, expectedLength int) error {
|
||||
if field != nil && len(field) != expectedLength {
|
||||
return fmt.Errorf("%s length should be at least %d bytes long", fieldName, expectedLength)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user