From d2bcdd2999c9c80c04d1411475b4d83000fa7232 Mon Sep 17 00:00:00 2001 From: William Petit Date: Sat, 1 Apr 2023 13:28:18 +0200 Subject: [PATCH] feat(storage,agent): add label attribute --- .gitignore | 2 +- cmd/agent/main.go | 1 + internal/client/update_agent.go | 11 +++++++++++ internal/command/api/agent/update.go | 10 ++++++++++ internal/command/api/agent/util.go | 1 + internal/datastore/agent.go | 1 + internal/datastore/agent_repository.go | 7 +++++++ internal/datastore/sqlite/agent_repository.go | 18 ++++++++++++------ internal/server/agent_api.go | 5 +++++ migrations/sqlite/0000001_agent_label.down.sql | 1 + migrations/sqlite/0000001_agent_label.up.sql | 1 + 11 files changed, 51 insertions(+), 7 deletions(-) create mode 100644 migrations/sqlite/0000001_agent_label.down.sql create mode 100644 migrations/sqlite/0000001_agent_label.up.sql diff --git a/.gitignore b/.gitignore index 46dc536..fc7ef18 100644 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,7 @@ dist/ /tools /tmp /state.json -/emissary.sqlite +/emissary.sqlite* /.gitea-release /agent-key.json /apps diff --git a/cmd/agent/main.go b/cmd/agent/main.go index 102c049..c741325 100644 --- a/cmd/agent/main.go +++ b/cmd/agent/main.go @@ -7,6 +7,7 @@ import ( "forge.cadoles.com/Cadoles/emissary/internal/command/agent" "forge.cadoles.com/Cadoles/emissary/internal/command/api" + _ "forge.cadoles.com/Cadoles/emissary/internal/imports/format" _ "forge.cadoles.com/Cadoles/emissary/internal/imports/spec" ) diff --git a/internal/client/update_agent.go b/internal/client/update_agent.go index 118efce..0ddeec5 100644 --- a/internal/client/update_agent.go +++ b/internal/client/update_agent.go @@ -10,6 +10,7 @@ import ( type UpdateAgentOptions struct { Status *int + Label *string Options []OptionFunc } @@ -21,6 +22,12 @@ func WithAgentStatus(status int) UpdateAgentOptionFunc { } } +func WithAgentLabel(label string) UpdateAgentOptionFunc { + return func(opts *UpdateAgentOptions) { + opts.Label = &label + } +} + func WithUpdateAgentsOptions(funcs ...OptionFunc) UpdateAgentOptionFunc { return func(opts *UpdateAgentOptions) { opts.Options = funcs @@ -39,6 +46,10 @@ func (c *Client) UpdateAgent(ctx context.Context, agentID datastore.AgentID, fun payload["status"] = *opts.Status } + if opts.Label != nil { + payload["label"] = *opts.Label + } + response := withResponse[struct { Agent *datastore.Agent `json:"agent"` }]() diff --git a/internal/command/api/agent/update.go b/internal/command/api/agent/update.go index 9b87db4..121b97d 100644 --- a/internal/command/api/agent/update.go +++ b/internal/command/api/agent/update.go @@ -22,6 +22,11 @@ func UpdateCommand() *cli.Command { Usage: "Set `STATUS` to selected agent", Value: -1, }, + &cli.StringFlag{ + Name: "label", + Usage: "Set `LABEL` to selected agent", + Value: "", + }, ), Action: func(ctx *cli.Context) error { baseFlags := clientFlag.GetBaseFlags(ctx) @@ -43,6 +48,11 @@ func UpdateCommand() *cli.Command { options = append(options, client.WithAgentStatus(status)) } + label := ctx.String("label") + if label != "" { + options = append(options, client.WithAgentLabel(label)) + } + client := client.New(baseFlags.ServerURL, client.WithToken(token)) agent, err := client.UpdateAgent(ctx.Context, agentID, options...) diff --git a/internal/command/api/agent/util.go b/internal/command/api/agent/util.go index faf89f0..78cfeec 100644 --- a/internal/command/api/agent/util.go +++ b/internal/command/api/agent/util.go @@ -7,6 +7,7 @@ func agentHints(outputMode format.OutputMode) format.Hints { OutputMode: outputMode, Props: []format.Prop{ format.NewProp("ID", "ID"), + format.NewProp("Label", "Label"), format.NewProp("Thumbprint", "Thumbprint"), format.NewProp("Status", "Status"), format.NewProp("CreatedAt", "CreatedAt"), diff --git a/internal/datastore/agent.go b/internal/datastore/agent.go index 3e0efc9..7709683 100644 --- a/internal/datastore/agent.go +++ b/internal/datastore/agent.go @@ -21,6 +21,7 @@ const ( type Agent struct { ID AgentID `json:"id"` + Label string `json:"label"` Thumbprint string `json:"thumbprint"` KeySet *SerializableKeySet `json:"keyset,omitempty"` Metadata map[string]any `json:"metadata,omitempty"` diff --git a/internal/datastore/agent_repository.go b/internal/datastore/agent_repository.go index 25a10aa..a643d67 100644 --- a/internal/datastore/agent_repository.go +++ b/internal/datastore/agent_repository.go @@ -68,6 +68,7 @@ func WithAgentQueryThumbprints(thumbprints ...string) AgentQueryOptionFunc { type AgentUpdateOptionFunc func(*AgentUpdateOptions) type AgentUpdateOptions struct { + Label *string Status *AgentStatus Metadata *map[string]any KeySet *jwk.Set @@ -97,3 +98,9 @@ func WithAgentUpdateThumbprint(thumbprint string) AgentUpdateOptionFunc { opts.Thumbprint = &thumbprint } } + +func WithAgentUpdateLabel(label string) AgentUpdateOptionFunc { + return func(opts *AgentUpdateOptions) { + opts.Label = &label + } +} diff --git a/internal/datastore/sqlite/agent_repository.go b/internal/datastore/sqlite/agent_repository.go index eda1c55..9141885 100644 --- a/internal/datastore/sqlite/agent_repository.go +++ b/internal/datastore/sqlite/agent_repository.go @@ -127,7 +127,7 @@ func (r *AgentRepository) Query(ctx context.Context, opts ...datastore.AgentQuer count := 0 err := r.withTx(ctx, func(tx *sql.Tx) error { - query := `SELECT id, thumbprint, status, created_at, updated_at FROM agents` + query := `SELECT id, label, thumbprint, status, created_at, updated_at FROM agents` limit := 10 if options.Limit != nil { @@ -195,7 +195,7 @@ func (r *AgentRepository) Query(ctx context.Context, opts ...datastore.AgentQuer metadata := JSONMap{} - if err := rows.Scan(&agent.ID, &agent.Thumbprint, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil { + if err := rows.Scan(&agent.ID, &agent.Label, &agent.Thumbprint, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil { return errors.WithStack(err) } @@ -315,7 +315,7 @@ func (r *AgentRepository) Get(ctx context.Context, id datastore.AgentID) (*datas err := r.withTx(ctx, func(tx *sql.Tx) error { query := ` - SELECT "id", "thumbprint", "keyset", "metadata", "status", "created_at", "updated_at" + SELECT "id", "label", "thumbprint", "keyset", "metadata", "status", "created_at", "updated_at" FROM agents WHERE id = $1 ` @@ -325,7 +325,7 @@ func (r *AgentRepository) Get(ctx context.Context, id datastore.AgentID) (*datas metadata := JSONMap{} var rawKeySet []byte - if err := row.Scan(&agent.ID, &agent.Thumbprint, &rawKeySet, &metadata, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil { + if err := row.Scan(&agent.ID, &agent.Label, &agent.Thumbprint, &rawKeySet, &metadata, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil { if errors.Is(err, sql.ErrNoRows) { return datastore.ErrNotFound } @@ -395,6 +395,12 @@ func (r *AgentRepository) Update(ctx context.Context, id datastore.AgentID, opts index++ } + if options.Label != nil { + query += fmt.Sprintf(`, label = $%d`, index) + args = append(args, *options.Label) + index++ + } + if options.Metadata != nil { query += fmt.Sprintf(`, metadata = $%d`, index) args = append(args, JSONMap(*options.Metadata)) @@ -403,7 +409,7 @@ func (r *AgentRepository) Update(ctx context.Context, id datastore.AgentID, opts query += ` WHERE id = $1 - RETURNING "id", "thumbprint", "keyset", "metadata", "status", "created_at", "updated_at" + RETURNING "id", "label", "thumbprint", "keyset", "metadata", "status", "created_at", "updated_at" ` row := tx.QueryRowContext(ctx, query, args...) @@ -411,7 +417,7 @@ func (r *AgentRepository) Update(ctx context.Context, id datastore.AgentID, opts metadata := JSONMap{} var rawKeySet []byte - if err := row.Scan(&agent.ID, &agent.Thumbprint, &rawKeySet, &metadata, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil { + if err := row.Scan(&agent.ID, &agent.Label, &agent.Thumbprint, &rawKeySet, &metadata, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil { if errors.Is(err, sql.ErrNoRows) { return datastore.ErrNotFound } diff --git a/internal/server/agent_api.go b/internal/server/agent_api.go index 01d2c5a..92f568c 100644 --- a/internal/server/agent_api.go +++ b/internal/server/agent_api.go @@ -145,6 +145,7 @@ func (s *Server) registerAgent(w http.ResponseWriter, r *http.Request) { type updateAgentRequest struct { Status *datastore.AgentStatus `json:"status" validate:"omitempty,oneof=0 1 2 3"` + Label *string `json:"label" validate:"omitempty"` } func (s *Server) updateAgent(w http.ResponseWriter, r *http.Request) { @@ -166,6 +167,10 @@ func (s *Server) updateAgent(w http.ResponseWriter, r *http.Request) { options = append(options, datastore.WithAgentUpdateStatus(*updateAgentReq.Status)) } + if updateAgentReq.Label != nil { + options = append(options, datastore.WithAgentUpdateLabel(*updateAgentReq.Label)) + } + agent, err := s.agentRepo.Update( ctx, datastore.AgentID(agentID), diff --git a/migrations/sqlite/0000001_agent_label.down.sql b/migrations/sqlite/0000001_agent_label.down.sql new file mode 100644 index 0000000..5ebb4c5 --- /dev/null +++ b/migrations/sqlite/0000001_agent_label.down.sql @@ -0,0 +1 @@ +ALTER TABLE agents DROP COLUMN label; \ No newline at end of file diff --git a/migrations/sqlite/0000001_agent_label.up.sql b/migrations/sqlite/0000001_agent_label.up.sql new file mode 100644 index 0000000..921914f --- /dev/null +++ b/migrations/sqlite/0000001_agent_label.up.sql @@ -0,0 +1 @@ +ALTER TABLE agents ADD COLUMN label TEXT DEFAULT ""; \ No newline at end of file