refactor

Signed-off-by: WashingtonKK <[email protected]>
WashingtonKK · Nov 3, 2024 · ed5284e · ed5284e
1 parent 2593797
commit ed5284e
Show file tree

Hide file tree

Showing 8 changed files with 40 additions and 23 deletions.
diff --git a/cmd/cli/main.go b/cmd/cli/main.go
@@ -15,7 +15,7 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/ultravioletrs/cocos/cli"
-	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	"github.com/ultravioletrs/cocos/internal/server"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
 	"github.com/ultravioletrs/cocos/pkg/sdk"
 	cmd "github.com/virtee/sev-snp-measure-go/sevsnpmeasure/cmd"
@@ -94,7 +94,7 @@ func main() {
 		return
 	}
 
-	agentGRPCConfig := grpc.ManagerConfig{}
+	agentGRPCConfig := server.AgentConfig{}
 	if err := env.ParseWithOptions(&agentGRPCConfig, env.Options{Prefix: envPrefixAgentGRPC}); err != nil {
 		message := color.New(color.FgRed).Sprintf("failed to load %s gRPC client configuration : %s", svcName, err)
 		rootCmd.Println(message)

diff --git a/cmd/manager/main.go b/cmd/manager/main.go
@@ -25,7 +25,6 @@ import (
 	"github.com/ultravioletrs/cocos/manager/events"
 	"github.com/ultravioletrs/cocos/manager/qemu"
 	"github.com/ultravioletrs/cocos/manager/tracing"
-	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	managergrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/manager"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/errgroup"
@@ -91,7 +90,7 @@ func main() {
 	args := qemuCfg.ConstructQemuArgs()
 	logger.Info(strings.Join(args, " "))
 
-	managerGRPCConfig := grpc.ManagerConfig{}
+	managerGRPCConfig := managergrpc.ManagerConfig{}
 	if err := env.ParseWithOptions(&managerGRPCConfig, env.Options{Prefix: envPrefixGRPC}); err != nil {
 		logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
 		exitCode = 1

diff --git a/pkg/clients/grpc/agent/agent.go b/pkg/clients/grpc/agent/agent.go
@@ -7,15 +7,20 @@ import (
 
 	"github.com/absmach/magistrala/pkg/errors"
 	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/internal/server"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	grpchealth "google.golang.org/grpc/health/grpc_health_v1"
 )
 
 var ErrAgentServiceUnavailable = errors.New("agent service is unavailable")
 
 // NewAgentClient creates new agent gRPC client instance.
-func NewAgentClient(ctx context.Context, cfg grpc.ManagerConfig) (grpc.Client, agent.AgentServiceClient, error) {
-	client, err := grpc.NewClient(cfg)
+func NewAgentClient(ctx context.Context, cfg server.AgentConfig) (grpc.Client, agent.AgentServiceClient, error) {
+	conf := grpc.ClientConfig{
+		BaseConfig:  cfg.BaseConfig,
+		AttestedTLS: cfg.AttestedTLS,
+	}
+	client, err := grpc.NewClient(conf)
 	if err != nil {
 		return nil, nil, err
 	}

diff --git a/pkg/clients/grpc/agent/agent_test.go b/pkg/clients/grpc/agent/agent_test.go
@@ -15,6 +15,7 @@ import (
 	"github.com/ultravioletrs/cocos/agent"
 	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
 	"github.com/ultravioletrs/cocos/agent/mocks"
+	"github.com/ultravioletrs/cocos/internal/server"
 	pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health"
@@ -102,7 +103,7 @@ func TestAgentClientIntegration(t *testing.T) {
 				testServer.health.SetServingStatus("agent", grpchealth.HealthCheckResponse_SERVING)
 			}
 
-			cfg := pkggrpc.ManagerConfig{
+			cfg := server.AgentConfig{
 				BaseConfig: pkggrpc.BaseConfig{
 					URL:     testServer.listenAddr,
 					Timeout: 1,

diff --git a/pkg/clients/grpc/connect.go b/pkg/clients/grpc/connect.go
@@ -74,10 +74,11 @@ type BaseConfig struct {
 	ClientCAFile string        `env:"CLIENT_CA_CERTS"    envDefault:""`
 }
 
-type ManagerConfig struct {
+type ClientConfig struct {
 	BaseConfig
 	BackendInfo string `env:"BACKEND_INFO"    envDefault:""`
 	ClientTLS   bool   `env:"CLIENT_TLS"    envDefault:"false"`
+	AttestedTLS bool   `env:"ATTESTED_TLS"       envDefault:"false"`
 }
 
 type AttestationConfiguration struct {
@@ -98,13 +99,13 @@ type Client interface {
 
 type client struct {
 	*grpc.ClientConn
-	cfg    ManagerConfig
+	cfg    ClientConfig
 	secure security
 }
 
 var _ Client = (*client)(nil)
 
-func NewClient(cfg ManagerConfig) (Client, error) {
+func NewClient(cfg ClientConfig) (Client, error) {
 	conn, secure, err := connect(cfg)
 	if err != nil {
 		return nil, err
@@ -143,7 +144,7 @@ func (c *client) Connection() *grpc.ClientConn {
 }
 
 // connect creates new gRPC client and connect to gRPC server.
-func connect(cfg ManagerConfig) (*grpc.ClientConn, security, error) {
+func connect(cfg ClientConfig) (*grpc.ClientConn, security, error) {
 	opts := []grpc.DialOption{
 		grpc.WithStatsHandler(otelgrpc.NewClientHandler()),
 	}

diff --git a/pkg/clients/grpc/connect_test.go b/pkg/clients/grpc/connect_test.go
@@ -38,13 +38,13 @@ func TestNewClient(t *testing.T) {
 
 	tests := []struct {
 		name    string
-		cfg     ManagerConfig
+		cfg     ClientConfig
 		wantErr bool
 		err     error
 	}{
 		{
 			name: "Success without TLS",
-			cfg: ManagerConfig{
+			cfg: ClientConfig{
 				BaseConfig: BaseConfig{
 					URL: "localhost:7001",
 				},
@@ -54,7 +54,7 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Success with TLS",
-			cfg: ManagerConfig{
+			cfg: ClientConfig{
 				BaseConfig: BaseConfig{
 					URL:          "localhost:7001",
 					ServerCAFile: caCertFile,
@@ -65,7 +65,7 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Success with mTLS",
-			cfg: ManagerConfig{
+			cfg: ClientConfig{
 				BaseConfig: BaseConfig{
 					URL:          "localhost:7001",
 					ServerCAFile: caCertFile,
@@ -78,7 +78,7 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Fail with invalid ServerCAFile",
-			cfg: ManagerConfig{
+			cfg: ClientConfig{
 				BaseConfig: BaseConfig{
 					URL:          "localhost:7001",
 					ServerCAFile: "nonexistent.pem",
@@ -89,7 +89,7 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Fail with invalid ClientCert",
-			cfg: ManagerConfig{
+			cfg: ClientConfig{
 				BaseConfig: BaseConfig{
 					URL:          "localhost:7001",
 					ServerCAFile: caCertFile,
@@ -102,7 +102,7 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Fail with invalid ClientKey",
-			cfg: ManagerConfig{
+			cfg: ClientConfig{
 				BaseConfig: BaseConfig{
 					URL:          "localhost:7001",
 					ServerCAFile: caCertFile,

diff --git a/pkg/clients/grpc/manager/manager.go b/pkg/clients/grpc/manager/manager.go
@@ -7,9 +7,20 @@ import (
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 )
 
+type ManagerConfig struct {
+	grpc.BaseConfig
+	BackendInfo string `env:"BACKEND_INFO"    envDefault:""`
+	ClientTLS   bool   `env:"CLIENT_TLS"    envDefault:"false"`
+}
+
 // NewManagerClient creates new manager gRPC client instance.
-func NewManagerClient(cfg grpc.ManagerConfig) (grpc.Client, manager.ManagerServiceClient, error) {
-	client, err := grpc.NewClient(cfg)
+func NewManagerClient(cfg ManagerConfig) (grpc.Client, manager.ManagerServiceClient, error) {
+	conf := grpc.ClientConfig{
+		BaseConfig:  cfg.BaseConfig,
+		ClientTLS:   cfg.ClientTLS,
+		BackendInfo: cfg.BackendInfo,
+	}
+	client, err := grpc.NewClient(conf)
 	if err != nil {
 		return nil, nil, err
 	}

diff --git a/pkg/clients/grpc/manager/manager_test.go b/pkg/clients/grpc/manager/manager_test.go
@@ -13,12 +13,12 @@ import (
 func TestNewManagerClient(t *testing.T) {
 	tests := []struct {
 		name string
-		cfg  grpc.ManagerConfig
+		cfg  ManagerConfig
 		err  error
 	}{
 		{
 			name: "Valid config",
-			cfg: grpc.ManagerConfig{
+			cfg: ManagerConfig{
 				BaseConfig: grpc.BaseConfig{
 					URL: "localhost:7001",
 				},
@@ -27,7 +27,7 @@ func TestNewManagerClient(t *testing.T) {
 		},
 		{
 			name: "invalid config, missing BackendInfo with aTLS",
-			cfg:  grpc.ManagerConfig{ClientTLS: true},
+			cfg:  ManagerConfig{ClientTLS: true},
 			err:  grpc.ErrBackendInfoMissing,
 		},
 	}