From 3515d392bfabc80080212aa5dcb60571bc408ed2 Mon Sep 17 00:00:00 2001 From: Henry Barreto Date: Thu, 27 Jun 2024 16:53:43 -0300 Subject: [PATCH] chore(agent): receive mode on initialization It was noticed that we could remove the `mode` field from the Agent's structure if we used it on the `Initialization` method, decreasing a bit the complexity of the structure and its requirements to be created. Beyond that, now we also create the SSH server in the initialization, taking advantage of this change. --- agent/main.go | 4 ++-- pkg/agent/agent.go | 28 +++++++++++++--------------- pkg/agent/agent_test.go | 20 +++----------------- pkg/agent/connector/docker.go | 4 ++-- 4 files changed, 20 insertions(+), 36 deletions(-) diff --git a/agent/main.go b/agent/main.go index 517c2b59a23..1c1f84e563e 100644 --- a/agent/main.go +++ b/agent/main.go @@ -85,7 +85,7 @@ func main() { "mode": mode, }).Info("Starting ShellHub") - ag, err := agent.NewAgentWithConfig(cfg, new(agent.HostMode)) + ag, err := agent.NewAgentWithConfig(cfg) if err != nil { log.WithError(err).WithFields(log.Fields{ "version": AgentVersion, @@ -93,7 +93,7 @@ func main() { }).Fatal("Failed to create agent") } - if err := ag.Initialize(); err != nil { + if err := ag.Initialize(new(agent.HostMode)); err != nil { log.WithError(err).WithFields(log.Fields{ "version": AgentVersion, "configuration": cfg, diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 563bbd1b18f..e1b0650dcf0 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -171,19 +171,18 @@ type Agent struct { tunnel *tunnel.Tunnel listening chan bool closed atomic.Bool - mode Mode } // NewAgent creates a new agent instance, requiring the ShellHub server's address to connect to, the namespace's tenant // where device own and the path to the private key on the file system. // // To create a new [Agent] instance with all configurations, you can use [NewAgentWithConfig]. -func NewAgent(address string, tenantID string, privateKey string, mode Mode) (*Agent, error) { +func NewAgent(address string, tenantID string, privateKey string) (*Agent, error) { return NewAgentWithConfig(&Config{ ServerAddress: address, TenantID: tenantID, PrivateKey: privateKey, - }, mode) + }) } var ( @@ -191,13 +190,13 @@ var ( ErrNewAgentWithConfigInvalidServerAddress = errors.New("address is invalid") ErrNewAgentWithConfigEmptyTenant = errors.New("tenant is empty") ErrNewAgentWithConfigEmptyPrivateKey = errors.New("private key is empty") - ErrNewAgentWithConfigNilMode = errors.New("agent's mode is nil") + ErrServerNil = errors.New("agent's mode is nil") ) // NewAgentWithConfig creates a new agent instance with all configurations. // // Check [Config] for more information. -func NewAgentWithConfig(config *Config, mode Mode) (*Agent, error) { +func NewAgentWithConfig(config *Config) (*Agent, error) { if config.ServerAddress == "" { return nil, ErrNewAgentWithConfigEmptyServerAddress } @@ -214,13 +213,8 @@ func NewAgentWithConfig(config *Config, mode Mode) (*Agent, error) { return nil, ErrNewAgentWithConfigEmptyPrivateKey } - if mode == nil { - return nil, ErrNewAgentWithConfigNilMode - } - return &Agent{ config: config, - mode: mode, }, nil } @@ -228,7 +222,7 @@ func NewAgentWithConfig(config *Config, mode Mode) (*Agent, error) { // key, reading public key, probing server information and authorizing device on ShellHub server. // // When any of the steps fails, the agent will return an error, and the agent will not be able to start. -func (a *Agent) Initialize() error { +func (a *Agent) Initialize(mode Mode) error { var err error a.cli, err = client.NewClient(a.config.ServerAddress) @@ -240,7 +234,7 @@ func (a *Agent) Initialize() error { return errors.Wrap(err, "failed to generate device identity") } - if err := a.loadDeviceInfo(); err != nil { + if err := a.loadDeviceInfo(mode); err != nil { return errors.Wrap(err, "failed to load device info") } @@ -260,6 +254,8 @@ func (a *Agent) Initialize() error { return errors.Wrap(err, "failed to authorize device") } + mode.Serve(a) + a.closed.Store(false) return nil @@ -310,8 +306,8 @@ func (a *Agent) generateDeviceIdentity() error { } // loadDeviceInfo load some device informations like OS name, version, arch and platform. -func (a *Agent) loadDeviceInfo() error { - info, err := a.mode.GetInfo() +func (a *Agent) loadDeviceInfo(mode Mode) error { + info, err := mode.GetInfo() if err != nil { return err } @@ -453,7 +449,9 @@ func closeHandler(a *Agent, serv *server.Server) func(c echo.Context) error { // Listen creates the SSH server and listening for connections. func (a *Agent) Listen(ctx context.Context) error { - a.mode.Serve(a) + if a.server == nil { + return ErrServerNil + } a.tunnel = tunnel.NewBuilder(). WithConnHandler(connHandler(a.server)). diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index 9038d9fdfd3..81b4a196426 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -18,14 +18,14 @@ func ExampleNewAgentWithConfig() { ServerAddress: "http://localhost:80", TenantID: "00000000-0000-4000-0000-000000000000", PrivateKey: "./shellhub.key", - }, new(HostMode)) + }) if err != nil { panic(err) } } func ExampleNewAgent() { - _, err := NewAgent("http://localhost:80", "00000000-0000-4000-0000-000000000000", "./shellhub.key", new(HostMode)) + _, err := NewAgent("http://localhost:80", "00000000-0000-4000-0000-000000000000", "./shellhub.key") if err != nil { panic(err) } @@ -211,19 +211,6 @@ func TestNewAgentWithConfig(t *testing.T) { err: ErrNewAgentWithConfigEmptyPrivateKey, }, }, - { - description: "fail when mode is nil", - config: &Config{ - ServerAddress: "http://localhost", - TenantID: "1c462afa-e4b6-41a5-ba54-7236a1770466", - PrivateKey: "/tmp/shellhub.key", - }, - mode: nil, - expected: expected{ - agent: nil, - err: ErrNewAgentWithConfigNilMode, - }, - }, { description: "success to create agent with config", config: config, @@ -231,7 +218,6 @@ func TestNewAgentWithConfig(t *testing.T) { expected: expected{ agent: &Agent{ config: config, - mode: new(HostMode), }, err: nil, }, @@ -240,7 +226,7 @@ func TestNewAgentWithConfig(t *testing.T) { for _, test := range tests { t.Run(test.description, func(t *testing.T) { - agent, err := NewAgentWithConfig(test.config, test.mode) + agent, err := NewAgentWithConfig(test.config) assert.Equal(t, test.expected.agent, agent) assert.ErrorIs(t, err, test.expected.err) diff --git a/pkg/agent/connector/docker.go b/pkg/agent/connector/docker.go index bc4e5dbcc9b..c590222d148 100644 --- a/pkg/agent/connector/docker.go +++ b/pkg/agent/connector/docker.go @@ -234,7 +234,7 @@ func initContainerAgent(ctx context.Context, cli *dockerclient.Client, container }).Fatal("Failed to create connector mode") } - ag, err := agent.NewAgentWithConfig(cfg, mode) + ag, err := agent.NewAgentWithConfig(cfg) if err != nil { log.WithError(err).WithFields(log.Fields{ "id": container.ID, @@ -243,7 +243,7 @@ func initContainerAgent(ctx context.Context, cli *dockerclient.Client, container }).Fatal("Failed to create agent") } - if err := ag.Initialize(); err != nil { + if err := ag.Initialize(mode); err != nil { log.WithError(err).WithFields(log.Fields{ "id": container.ID, "configuration": cfg,