Skip to content

Commit 823fa08

Browse files
KbayeroKbayero
committed
feat[agent-manager]: own and validate the agent-installation connection key locally
1 parent 7bc2ead commit 823fa08

17 files changed

Lines changed: 713 additions & 167 deletions

File tree

‎agent-manager/agent/agent.pb.go‎

Lines changed: 158 additions & 67 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎agent-manager/agent/agent_grpc.pb.go‎

Lines changed: 87 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎agent-manager/agent/agent_imp.go‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,10 @@ type AgentService struct {
3232
CacheAgentKeyMutex sync.RWMutex
3333
CommandResultChannel map[string]chan *CommandResult
3434
CommandResultChannelM sync.Mutex
35-
36-
DBConnection *database.DB
35+
connKey string
36+
connKeyID uint
37+
connKeyMutex sync.RWMutex
38+
DBConnection *database.DB
3739
}
3840

3941
func (s *AgentService) ValidateAgentKey(key string, id uint) bool {
@@ -62,6 +64,11 @@ func InitAgentService() error {
6264
for _, agent := range agents {
6365
AgentServ.CacheAgentKey[agent.ID] = agent.AgentKey
6466
}
67+
68+
if e := AgentServ.loadConnectionKey(); e != nil {
69+
err = e
70+
return
71+
}
6572
})
6673
return err
6774
}

‎agent-manager/agent/connectionkey_imp.go‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package agent
2+
3+
import (
4+
"context"
5+
"crypto/rand"
6+
"crypto/subtle"
7+
"encoding/hex"
8+
"fmt"
9+
10+
"github.com/threatwinds/go-sdk/catcher"
11+
"github.com/utmstack/UTMStack/agent-manager/models"
12+
"google.golang.org/grpc/codes"
13+
"google.golang.org/grpc/status"
14+
"gorm.io/gorm"
15+
)
16+
17+
func (s *AgentService) loadConnectionKey() error {
18+
var rows []models.ConnectionKey
19+
if _, err := s.DBConnection.GetAll(&rows, ""); err != nil {
20+
return fmt.Errorf("failed to load connection key: %v", err)
21+
}
22+
23+
if len(rows) > 0 {
24+
s.connKeyMutex.Lock()
25+
s.connKeyID = rows[0].ID
26+
s.connKey = rows[0].Key
27+
s.connKeyMutex.Unlock()
28+
return nil
29+
}
30+
31+
key, err := generateConnectionKey()
32+
if err != nil {
33+
return err
34+
}
35+
row := models.ConnectionKey{Key: key}
36+
if err := s.DBConnection.Create(&row); err != nil {
37+
return fmt.Errorf("failed to create connection key: %v", err)
38+
}
39+
s.connKeyMutex.Lock()
40+
s.connKeyID = row.ID
41+
s.connKey = row.Key
42+
s.connKeyMutex.Unlock()
43+
catcher.Info("Generated agent connection key", map[string]any{"process": "agent-manager"})
44+
return nil
45+
}
46+
47+
// ValidateConnectionKey reports whether the presented key matches the current
48+
// connection key.
49+
func (s *AgentService) ValidateConnectionKey(key string) bool {
50+
s.connKeyMutex.RLock()
51+
defer s.connKeyMutex.RUnlock()
52+
if s.connKey == "" {
53+
return false
54+
}
55+
return subtle.ConstantTimeCompare([]byte(key), []byte(s.connKey)) == 1
56+
}
57+
58+
func (s *AgentService) GetConnectionKey(ctx context.Context, req *ConnectionKeyRequest) (*ConnectionKeyResponse, error) {
59+
s.connKeyMutex.RLock()
60+
key := s.connKey
61+
s.connKeyMutex.RUnlock()
62+
return &ConnectionKeyResponse{ConnectionKey: key}, nil
63+
}
64+
65+
func (s *AgentService) RotateConnectionKey(ctx context.Context, req *ConnectionKeyRequest) (*ConnectionKeyResponse, error) {
66+
key, err := generateConnectionKey()
67+
if err != nil {
68+
return nil, status.Error(codes.Internal, fmt.Sprintf("failed to generate connection key: %v", err))
69+
}
70+
71+
s.connKeyMutex.RLock()
72+
id := s.connKeyID
73+
s.connKeyMutex.RUnlock()
74+
75+
if err := s.DBConnection.Upsert(&models.ConnectionKey{Model: gorm.Model{ID: id}, Key: key}, "id = ?", map[string]interface{}{"key": key}, id); err != nil {
76+
return nil, status.Error(codes.Internal, fmt.Sprintf("failed to persist connection key: %v", err))
77+
}
78+
79+
s.connKeyMutex.Lock()
80+
s.connKey = key
81+
s.connKeyMutex.Unlock()
82+
83+
catcher.Info("Rotated agent connection key", map[string]any{"process": "agent-manager"})
84+
return &ConnectionKeyResponse{ConnectionKey: key}, nil
85+
}
86+
87+
func generateConnectionKey() (string, error) {
88+
b := make([]byte, 32)
89+
if _, err := rand.Read(b); err != nil {
90+
return "", fmt.Errorf("failed to read random bytes: %v", err)
91+
}
92+
return hex.EncodeToString(b), nil
93+
}

‎agent-manager/agent/interceptor.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import (
88
"strings"
99

1010
"github.com/utmstack/UTMStack/agent-manager/config"
11-
"github.com/utmstack/UTMStack/agent-manager/utils"
1211
"google.golang.org/grpc"
1312
"google.golang.org/grpc/codes"
1413
"google.golang.org/grpc/metadata"
@@ -90,7 +89,7 @@ func authHeaders(md metadata.MD, fullMethod string) error {
9089
return status.Error(codes.PermissionDenied, "invalid type")
9190
}
9291
case "connection-key":
93-
if !utils.IsConnectionKeyValid(fmt.Sprintf(config.PanelConnectionKeyUrl, config.UTMHost), authConnectionKey[0]) {
92+
if !AgentServ.ValidateConnectionKey(authConnectionKey[0]) {
9493
return status.Error(codes.PermissionDenied, "invalid connection key")
9594
}
9695
case "internal-key":

‎agent-manager/config/global_const.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,14 +39,15 @@ func InternalKeyRoutes() []string {
3939
"/agent.CollectorService/ListCollector",
4040

4141
"/agent.PanelService/ProcessCommand",
42+
"/agent.PanelService/GetConnectionKey",
43+
"/agent.PanelService/RotateConnectionKey",
4244
"/agent.PanelCollectorService/RegisterCollectorConfig",
4345

4446
"/grpc.health.v1.Health/Check",
4547
}
4648
}
4749

4850
var (
49-
PanelConnectionKeyUrl = "%s/api/authenticateFederationServiceManager"
5051
CheckEvery = 5 * time.Minute
5152
CertPath = "/cert/utm.crt"
5253
CertKeyPath = "/cert/utm.key"

‎agent-manager/database/migration.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import "github.com/utmstack/UTMStack/agent-manager/models"
44

55
func MigrateDatabase() error {
66
db := GetDB()
7-
err := db.Migrate(&models.Agent{}, &models.AgentCommand{}, &models.LastSeen{}, &models.Collector{})
7+
err := db.Migrate(&models.Agent{}, &models.AgentCommand{}, &models.LastSeen{}, &models.Collector{}, &models.ConnectionKey{})
88
if err != nil {
99
return err
1010
}

‎agent-manager/models/connection_key.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
package models
2+
3+
import "gorm.io/gorm"
4+
5+
type ConnectionKey struct {
6+
gorm.Model
7+
Key string `gorm:"not null"`
8+
}

‎agent-manager/protos/agent.proto‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,14 @@ service AgentService {
1818

1919
service PanelService {
2020
rpc ProcessCommand (stream UtmCommand) returns (stream CommandResult){}
21+
rpc GetConnectionKey(ConnectionKeyRequest) returns (ConnectionKeyResponse) {}
22+
rpc RotateConnectionKey(ConnectionKeyRequest) returns (ConnectionKeyResponse) {}
23+
}
24+
25+
message ConnectionKeyRequest {}
26+
27+
message ConnectionKeyResponse {
28+
string connection_key = 1;
2129
}
2230

2331
enum AgentCommandStatus {

‎agent-manager/utils/auth.go‎

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,23 +2,8 @@ package utils
22

33
import (
44
"crypto/subtle"
5-
"crypto/tls"
6-
"net/http"
7-
"strings"
85
)
96

10-
func IsConnectionKeyValid(panelUrl string, token string) bool {
11-
requestBody := strings.NewReader(token)
12-
tlsConfig := &tls.Config{InsecureSkipVerify: true}
13-
transport := &http.Transport{TLSClientConfig: tlsConfig}
14-
client := &http.Client{Transport: transport}
15-
resp, err := client.Post(panelUrl, "application/json", requestBody)
16-
if err != nil || resp.StatusCode != http.StatusOK {
17-
return false
18-
}
19-
return true
20-
}
21-
227
func IsKeyPairValid(key string, id uint, cache map[uint]string) (string, bool) {
238
agentKey, ok := cache[id]
249
if !ok {

0 commit comments

Comments
 (0)