domrichardson
92 lines
2.3 KiB
Go
92 lines
2.3 KiB
Go
package services
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/mrhid6/keymanager/server/internal/grpc/pb"
|
|
)
|
|
|
|
type commandDispatcher struct {
|
|
mu sync.RWMutex
|
|
channels map[string]chan *pb.ServerCommand
|
|
}
|
|
|
|
// Dispatcher is the singleton command dispatcher used by both the gRPC server
|
|
// and the REST API to push commands to connected agents.
|
|
var Dispatcher = &commandDispatcher{
|
|
channels: make(map[string]chan *pb.ServerCommand),
|
|
}
|
|
|
|
// Connect registers an agent's command channel. Returns the channel to drain.
|
|
func (d *commandDispatcher) Connect(serverID string) chan *pb.ServerCommand {
|
|
ch := make(chan *pb.ServerCommand, 16)
|
|
d.mu.Lock()
|
|
d.channels[serverID] = ch
|
|
d.mu.Unlock()
|
|
return ch
|
|
}
|
|
|
|
// Disconnect removes the agent's channel on stream close.
|
|
func (d *commandDispatcher) Disconnect(serverID string) {
|
|
d.mu.Lock()
|
|
delete(d.channels, serverID)
|
|
d.mu.Unlock()
|
|
}
|
|
|
|
// IsConnected reports whether an agent is currently holding a CommandStream.
|
|
func (d *commandDispatcher) IsConnected(serverID string) bool {
|
|
d.mu.RLock()
|
|
_, ok := d.channels[serverID]
|
|
d.mu.RUnlock()
|
|
return ok
|
|
}
|
|
|
|
func (d *commandDispatcher) dispatch(serverID string, cmd *pb.ServerCommand) error {
|
|
d.mu.RLock()
|
|
ch, ok := d.channels[serverID]
|
|
d.mu.RUnlock()
|
|
if !ok {
|
|
return fmt.Errorf("agent for server %s is not connected", serverID)
|
|
}
|
|
select {
|
|
case ch <- cmd:
|
|
return nil
|
|
default:
|
|
return fmt.Errorf("command queue full for server %s", serverID)
|
|
}
|
|
}
|
|
|
|
// KeyGenParams carries all options for a generate-key command.
|
|
type KeyGenParams struct {
|
|
Label string
|
|
KeyType string
|
|
KeySize int
|
|
Passphrase string
|
|
Comment string
|
|
}
|
|
|
|
// DispatchGenerateKey sends a generate-key command to the named server's agent.
|
|
// Returns the command ID that can be used to correlate the agent's result.
|
|
func DispatchGenerateKey(serverID string, p KeyGenParams) (string, error) {
|
|
if !Dispatcher.IsConnected(serverID) {
|
|
return "", fmt.Errorf("agent is not connected to the command stream")
|
|
}
|
|
cmdID := uuid.New().String()
|
|
cmd := &pb.ServerCommand{
|
|
CommandId: cmdID,
|
|
GenerateKey: &pb.GenerateKeyCmd{
|
|
Label: p.Label,
|
|
KeyType: p.KeyType,
|
|
KeySize: p.KeySize,
|
|
Passphrase: p.Passphrase,
|
|
Comment: p.Comment,
|
|
},
|
|
}
|
|
if err := Dispatcher.dispatch(serverID, cmd); err != nil {
|
|
return "", err
|
|
}
|
|
return cmdID, nil
|
|
}
|