arcad/emissary
16
0
Fork 0
Code Issues 14 Pull Requests Packages Projects Releases 3 Wiki Activity

feat: agent metadata with custom collectors

Browse Source
This commit is contained in:
wpetit
2023-03-02 13:05:24 +01:00
parent 3310c09320
commit 1ff29ae1fb
40 changed files with 998 additions and 256 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
internal/datastore/sqlite/agent_repository.go
View File

@ -3,10 +3,13 @@ package sqlite
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"time"
"forge.cadoles.com/Cadoles/emissary/internal/datastore"
"github.com/lestrrat-go/jwx/v2/jwk"
"github.com/pkg/errors"
"gitlab.com/wpetit/goweb/logger"
)
@ -116,7 +119,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, remote_id, status, created_at, updated_at FROM agents`
query := `SELECT id, thumbprint, status, created_at, updated_at FROM agents`
limit := 10
if options.Limit != nil {
@ -133,20 +136,18 @@ func (r *AgentRepository) Query(ctx context.Context, opts ...datastore.AgentQuer
args := []any{offset, limit}
if options.IDs != nil && len(options.IDs) > 0 {
filters += "id in ("
filter, newArgs, newParamIndex := inFilter("id", paramIndex, options.RemoteIDs)
filter, newArgs, newParamIndex := inFilter("id", paramIndex, options.IDs)
filters += filter
paramIndex = newParamIndex
args = append(args, newArgs...)
}
if options.RemoteIDs != nil && len(options.RemoteIDs) > 0 {
if options.Thumbprints != nil && len(options.Thumbprints) > 0 {
if filters != "" {
filters += " AND "
}
filter, newArgs, newParamIndex := inFilter("remote_id", paramIndex, options.RemoteIDs)
filter, newArgs, newParamIndex := inFilter("thumbprint", paramIndex, options.Thumbprints)
filters += filter
paramIndex = newParamIndex
args = append(args, newArgs...)
@ -180,10 +181,14 @@ func (r *AgentRepository) Query(ctx context.Context, opts ...datastore.AgentQuer
for rows.Next() {
agent := &datastore.Agent{}
if err := rows.Scan(&agent.ID, &agent.RemoteID, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil {
metadata := JSONMap{}
if err := rows.Scan(&agent.ID, &agent.Thumbprint, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil {
return errors.WithStack(err)
}
agent.Metadata = metadata
agents = append(agents, agent)
}
@ -202,12 +207,12 @@ func (r *AgentRepository) Query(ctx context.Context, opts ...datastore.AgentQuer
}
// Create implements datastore.AgentRepository
func (r *AgentRepository) Create(ctx context.Context, remoteID string, status datastore.AgentStatus) (*datastore.Agent, error) {
func (r *AgentRepository) Create(ctx context.Context, thumbprint string, keySet jwk.Set, metadata map[string]any) (*datastore.Agent, error) {
agent := &datastore.Agent{}
err := r.withTx(ctx, func(tx *sql.Tx) error {
query := `SELECT count(id) FROM agents WHERE remote_id = $1`
row := tx.QueryRowContext(ctx, query, remoteID)
query := `SELECT count(id) FROM agents WHERE thumbprint = $1`
row := tx.QueryRowContext(ctx, query, thumbprint)
var count int
@ -222,21 +227,37 @@ func (r *AgentRepository) Create(ctx context.Context, remoteID string, status da
now := time.Now().UTC()
query = `
INSERT INTO agents (remote_id, status, created_at, updated_at)
VALUES($1, $2, $3, $3)
RETURNING "id", "remote_id", "status", "created_at", "updated_at"
INSERT INTO agents (thumbprint, keyset, metadata, status, created_at, updated_at)
VALUES($1, $2, $3, $4, $5, $5)
RETURNING "id", "thumbprint", "keyset", "metadata", "status", "created_at", "updated_at"
`
row = tx.QueryRowContext(
ctx, query,
remoteID, status, now,
)
err := row.Scan(&agent.ID, &agent.RemoteID, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt)
rawKeySet, err := json.Marshal(keySet)
if err != nil {
return errors.WithStack(err)
}
row = tx.QueryRowContext(
ctx, query,
thumbprint, rawKeySet, JSONMap(metadata), datastore.AgentStatusPending, now,
)
metadata := JSONMap{}
err = row.Scan(&agent.ID, &agent.Thumbprint, &rawKeySet, &metadata, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt)
if err != nil {
return errors.WithStack(err)
}
agent.Metadata = metadata
keySet, err = jwk.Parse(rawKeySet)
if err != nil {
return errors.WithStack(err)
}
agent.KeySet = &datastore.SerializableKeySet{keySet}
return nil
})
if err != nil {
@ -266,14 +287,17 @@ func (r *AgentRepository) Get(ctx context.Context, id datastore.AgentID) (*datas
err := r.withTx(ctx, func(tx *sql.Tx) error {
query := `
SELECT "remote_id", "status", "created_at", "updated_at"
SELECT "id", "thumbprint", "keyset", "metadata", "status", "created_at", "updated_at"
FROM agents
WHERE id = $1
`
row := r.db.QueryRowContext(ctx, query, id)
if err := row.Scan(&agent.RemoteID, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil {
metadata := JSONMap{}
var rawKeySet []byte
if err := row.Scan(&agent.ID, &agent.Thumbprint, &rawKeySet, &metadata, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return datastore.ErrNotFound
}
@ -281,6 +305,15 @@ func (r *AgentRepository) Get(ctx context.Context, id datastore.AgentID) (*datas
return errors.WithStack(err)
}
agent.Metadata = metadata
keySet := jwk.NewSet()
if err := json.Unmarshal(rawKeySet, &keySet); err != nil {
return errors.WithStack(err)
}
agent.KeySet = &datastore.SerializableKeySet{keySet}
return nil
})
if err != nil {
@ -313,23 +346,60 @@ func (r *AgentRepository) Update(ctx context.Context, id datastore.AgentID, opts
if options.Status != nil {
query += fmt.Sprintf(`, status = $%d`, index)
args = append(args, *options.Status)
index++
}
if options.KeySet != nil {
rawKeySet, err := json.Marshal(*options.KeySet)
if err != nil {
return errors.WithStack(err)
}
query += fmt.Sprintf(`, keyset = $%d`, index)
args = append(args, rawKeySet)
index++
}
if options.Thumbprint != nil {
query += fmt.Sprintf(`, thumbprint = $%d`, index)
args = append(args, *options.Thumbprint)
index++
}
if options.Metadata != nil {
query += fmt.Sprintf(`, metadata = $%d`, index)
args = append(args, JSONMap(*options.Metadata))
index++
}
query += `
WHERE id = $1
RETURNING "id","remote_id","status","updated_at","created_at"
RETURNING "id", "thumbprint", "keyset", "metadata", "status", "created_at", "updated_at"
`
row := tx.QueryRowContext(ctx, query, args...)
if err := row.Scan(&agent.ID, &agent.RemoteID, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil {
metadata := JSONMap{}
var rawKeySet []byte
if err := row.Scan(&agent.ID, &agent.Thumbprint, &rawKeySet, &metadata, &agent.Status, &agent.CreatedAt, &agent.UpdatedAt); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return datastore.ErrNotFound
}
return errors.WithStack(err)
}
agent.Metadata = metadata
keySet := jwk.NewSet()
if err := json.Unmarshal(rawKeySet, &keySet); err != nil {
return errors.WithStack(err)
}
agent.KeySet = &datastore.SerializableKeySet{keySet}
return nil
})
if err != nil {