Skip to content

Commit

Permalink
clean up connects
Browse files Browse the repository at this point in the history
  • Loading branch information
jsbroks committed Nov 15, 2024
1 parent 7672e25 commit 7ddd5e4
Show file tree
Hide file tree
Showing 7 changed files with 62 additions and 46 deletions.
9 changes: 9 additions & 0 deletions cmd/ctrlc/ctrlc.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,15 @@ var (
func init() {
cobra.OnInitialize(initConfig)
cmd.PersistentFlags().StringVar(&cfgFile, "config", "", "Config file (default is $HOME/.ctrlc.yaml)")
viper.BindEnv("config", "CTRLPLANE_CONFIG")

cmd.PersistentFlags().String("url", "", "API URL")
viper.BindPFlag("url", cmd.PersistentFlags().Lookup("url"))
viper.BindEnv("url", "CTRLPLANE_URL")

cmd.PersistentFlags().String("api-key", "", "API key for authentication")
viper.BindPFlag("api-key", cmd.PersistentFlags().Lookup("api-key"))
viper.BindEnv("api-key", "CTRLPLANE_API_KEY")
}

func main() {
Expand Down
34 changes: 11 additions & 23 deletions cmd/ctrlc/root/agent/run/run.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package run

import (
"log"
"os"
"strings"
"time"

Expand All @@ -12,45 +11,35 @@ import (
)

func NewAgentRunCmd() *cobra.Command {
var proxyAddr string
var agentName string
var workspace string
var metadata map[string]string
var insecure bool
var targetId string
var associatedResources []string

cmd := &cobra.Command{
Use: "run",
Short: "Run the agent",
Long: `Run the agent to establish connection with the proxy.`,
RunE: func(cmd *cobra.Command, args []string) error {
if proxyAddr == "" {
proxyAddr = viper.GetString("url")
proxyAddr = strings.TrimPrefix(proxyAddr, "https://")
proxyAddr = strings.TrimPrefix(proxyAddr, "http://")
}
proxyAddr := viper.GetString("url")
proxyAddr = strings.TrimPrefix(proxyAddr, "https://")
proxyAddr = strings.TrimPrefix(proxyAddr, "http://")

if insecure {
proxyAddr = "ws://" + proxyAddr
} else {
proxyAddr = "wss://" + proxyAddr
}

apiKey := os.Getenv("CTRLPLANE_API_KEY")
if apiKey == "" {
apiKey = viper.GetString("api-key")
}

opts := []func(*agent.Agent){
agent.WithMetadata(metadata),
agent.WithHeader("X-API-Key", apiKey),
agent.WithHeader("X-Workspace", workspace),
}

apiKey := viper.GetString("api-key")
agent := agent.NewAgent(
proxyAddr,
agentName,
opts...,
agent.WithMetadata(metadata),
agent.WithHeader("X-API-Key", apiKey),
agent.WithHeader("X-Workspace", workspace),
agent.WithAssociatedResources(associatedResources),
)

backoff := time.Second
Expand All @@ -72,12 +61,11 @@ func NewAgentRunCmd() *cobra.Command {
SilenceUsage: true,
}

cmd.Flags().StringVarP(&proxyAddr, "proxy", "p", "app.ctrlplane.dev", "Proxy address to connect through")
cmd.Flags().StringVarP(&agentName, "name", "n", "", "Name for this agent")
cmd.Flags().StringVarP(&workspace, "workspace", "w", "", "Workspace for this agent")
cmd.Flags().StringVarP(&targetId, "target", "t", "", "Target ID to link this agent too")
cmd.Flags().StringToStringVar(&metadata, "metadata", make(map[string]string), "Metadata key-value pairs (e.g. --metadata key=value)")
cmd.Flags().BoolVar(&insecure, "insecure", false, "Allow insecure connections")
cmd.Flags().BoolVar(&insecure, "insecure", false, "Allow insecure connections (a.k use ws://)")
cmd.Flags().StringArrayVarP(&associatedResources, "associated-resource", "r", []string{}, "Resource ID or Identifier to associate this agent with")

cmd.MarkFlagRequired("workspace")
cmd.MarkFlagRequired("name")
Expand Down
5 changes: 0 additions & 5 deletions cmd/ctrlc/root/api/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,14 +31,9 @@ func NewAPICmd() *cobra.Command {
return cmd.Help()
},
}
cmd.PersistentFlags().String("url", "", "API URL")
cmd.PersistentFlags().String("api-key", "", "API key for authentication")
cmd.PersistentFlags().String("template", "", "Template for output format. Accepts Go template format (e.g. --template='{{.status.phase}}')")
cmd.PersistentFlags().String("format", "json", "Output format. Accepts 'json' or 'yaml'")

viper.BindPFlag("url", cmd.PersistentFlags().Lookup("url"))
viper.BindPFlag("api-key", cmd.PersistentFlags().Lookup("api-key"))

cmd.AddCommand(get.NewGetCmd())
cmd.AddCommand(create.NewCreateCmd())
cmd.AddCommand(upsert.NewUpsertCmd())
Expand Down
1 change: 0 additions & 1 deletion cmd/ctrlc/root/api/delete/delete.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ func NewDeleteCmd() *cobra.Command {
},
}


cmd.AddCommand(resource.NewDeleteResourceCmd())

return cmd
Expand Down
49 changes: 32 additions & 17 deletions pkg/agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ import (
"log"
"net/http"
"os"
"runtime"
"strconv"
"time"

"github.com/creack/pty"
Expand Down Expand Up @@ -34,26 +36,28 @@ func GetControllerProxyURL() string {
}

type Agent struct {
headers http.Header
client *client.Client
serverURL string
agentName string
StopSignal chan struct{}
metadata map[string]string
manager *ptysession.Manager
headers http.Header
client *client.Client
serverURL string
agentName string
StopSignal chan struct{}
metadata map[string]string
associatedResources []string
manager *ptysession.Manager
}

func NewAgent(serverURL, agentName string, opts ...func(*Agent)) *Agent {
headers := make(http.Header)
headers.Set("User-Agent", "ctrlplane-cli")
headers.Set("X-Agent-Name", agentName)
agent := &Agent{
headers: headers,
serverURL: serverURL,
agentName: agentName,
StopSignal: make(chan struct{}),
metadata: make(map[string]string),
manager: ptysession.GetManager(),
headers: headers,
serverURL: serverURL,
agentName: agentName,
StopSignal: make(chan struct{}),
metadata: make(map[string]string),
manager: ptysession.GetManager(),
associatedResources: []string{},
}
for _, opt := range opts {
opt(agent)
Expand Down Expand Up @@ -135,12 +139,23 @@ func (a *Agent) handleConnect() {
log.Printf("Agent %s connected to server", a.agentName)

connectPayload := payloads.AgentConnectJson{
Type: payloads.AgentConnectJsonTypeAgentConnect,
Name: a.agentName,
Config: map[string]interface{}{},
Metadata: a.metadata,
Type: payloads.AgentConnectJsonTypeAgentConnect,
Name: a.agentName,
Config: map[string]interface{}{},
Metadata: a.metadata,
AssociatedResources: a.associatedResources,
}

var memStats runtime.MemStats
runtime.ReadMemStats(&memStats)
connectPayload.Metadata["go/memstats/totalalloc"] = strconv.FormatUint(memStats.TotalAlloc, 10)
connectPayload.Metadata["go/memstats/sys"] = strconv.FormatUint(memStats.Sys, 10)
connectPayload.Metadata["runtime/os"] = runtime.GOOS
connectPayload.Metadata["runtime/arch"] = runtime.GOARCH
connectPayload.Metadata["go/version"] = runtime.Version()
connectPayload.Metadata["go/compiler"] = runtime.Compiler
connectPayload.Metadata["go/numcpu"] = strconv.Itoa(runtime.NumCPU())

data, err := json.Marshal(connectPayload)
if err != nil {
log.Printf("Error marshaling connect payload: %v", err)
Expand Down
7 changes: 7 additions & 0 deletions pkg/agent/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,10 @@ func WithHeader(key string, value string) func(*Agent) {
a.headers.Set(key, value)
}
}


func WithAssociatedResources(resources []string) func(*Agent) {
return func(a *Agent) {
a.associatedResources = resources
}
}
3 changes: 3 additions & 0 deletions pkg/payloads/agent_connect_generated.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 7ddd5e4

Please sign in to comment.