Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 111 additions & 19 deletions cmd/start.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,21 @@ package cmd

import (
"context"
"crypto/ecdsa"
"crypto/x509"
"encoding/base64"
"encoding/pem"
"fmt"
"net/http"
"os"
"os/signal"
"runtime/debug"
"strings"
"sync"
"syscall"
"time"

"github.com/agentuity/go-common/gravity/provider"
_logger "github.com/agentuity/go-common/logger"
cnet "github.com/agentuity/go-common/network"
csys "github.com/agentuity/go-common/sys"
Expand Down Expand Up @@ -56,20 +64,48 @@ var rootCmd = &cobra.Command{
localPort, _ := flags.GetInt("port")
orgID, _ := flags.GetString("org-id")
projectID, _ := flags.GetString("project-id")
token, _ := flags.GetString("token")
maybePrivateKey, _ := flags.GetString("private-key")
endpointID, _ := flags.GetString("endpoint-id")
gravityUrl, _ := flags.GetString("url")
healthCheck, _ := flags.GetBool("health-check")
token, _ := flags.GetString("token")

if token != "" {
logger.Fatal("The --token flag is no longer supported. Please update your Agentuity CLI to the latest version.")
}
Comment on lines +67 to +75
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Make the --token deprecation message actionable.
The current fatal message implies “update CLI,” but the fix is to remove --token/AGENTUITY_SDK_KEY and use --private-key. Consider updating the text to reduce user confusion.

💡 Suggested wording tweak
-		if token != "" {
-			logger.Fatal("The --token flag is no longer supported. Please update your Agentuity CLI to the latest version.")
-		}
+		if token != "" {
+			logger.Fatal("The --token flag is no longer supported. Remove --token/AGENTUITY_SDK_KEY and use --private-key instead.")
+		}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
maybePrivateKey, _ := flags.GetString("private-key")
endpointID, _ := flags.GetString("endpoint-id")
gravityUrl, _ := flags.GetString("url")
healthCheck, _ := flags.GetBool("health-check")
token, _ := flags.GetString("token")
if token != "" {
logger.Fatal("The --token flag is no longer supported. Please update your Agentuity CLI to the latest version.")
}
maybePrivateKey, _ := flags.GetString("private-key")
endpointID, _ := flags.GetString("endpoint-id")
gravityUrl, _ := flags.GetString("url")
healthCheck, _ := flags.GetBool("health-check")
token, _ := flags.GetString("token")
if token != "" {
logger.Fatal("The --token flag is no longer supported. Remove --token/AGENTUITY_SDK_KEY and use --private-key instead.")
}
🤖 Prompt for AI Agents
In `@cmd/start.go` around lines 67 - 75, Change the unclear deprecation log for
the token flag: when reading token via flags.GetString("token") and detecting
token != "", update the logger.Fatal message (the code around maybePrivateKey
and flags) to explicitly instruct users to stop using --token/AGENTUITY_SDK_KEY
and to use --private-key (or the private key env/flag) instead — reference the
token variable and the logger.Fatal call and replace the generic "update your
Agentuity CLI" text with a clear, actionable message that says remove
--token/AGENTUITY_SDK_KEY and supply the private key via --private-key.


ipv4addr, err := utils.GetPrivateIPv4()
if err != nil {
logger.Fatal("failed to get private IPv4: %v", err)
}

privateKeyPEM, err := loadPrivateKeyPEM(maybePrivateKey)
if err != nil {
logger.Fatal("failed to load private key: %v", err)
}

block, _ := pem.Decode(privateKeyPEM)
if block == nil {
logger.Fatal("no PEM block found in --private-key")
}
key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
if err != nil {
// Try SEC 1/EC format as fallback
key, err = x509.ParseECPrivateKey(block.Bytes)
if err != nil {
logger.Fatal("failed to parse private key: %v (pem length: %d)", err, len(privateKeyPEM))
}
}
privateKey, ok := key.(*ecdsa.PrivateKey)
if !ok {
logger.Fatal("private key is not an ECDSA key")
}

agent := stack.AgentMetadata{
OrgID: orgID,
ProjectID: projectID,
InstanceID: endpointID,
PrivateKey: privateKey,
}
Comment on lines +82 to 109
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# First, let's find and read the file cmd/start.go
find . -name "cmd" -type d | head -5
find . -name "start.go" -type f | head -5

Repository: agentuity/gravity

Length of output: 81


🏁 Script executed:

#!/bin/bash
# Read the specific lines in question (82-109)
if [ -f "cmd/start.go" ]; then
  echo "=== cmd/start.go lines 75-115 ===" 
  sed -n '75,115p' "cmd/start.go" | cat -n
else
  echo "File not found, searching..."
  find . -name "start.go" -type f -exec wc -l {} \; -exec echo "File: {}" \;
fi

Repository: agentuity/gravity

Length of output: 1535


🏁 Script executed:

#!/bin/bash
# Also check if there are any tests or examples showing how privateKeyPEM is used
rg -A 5 -B 5 "pem.Decode" --type go

Repository: agentuity/gravity

Length of output: 1754


🏁 Script executed:

#!/bin/bash
# Find the loadPrivateKeyPEM function
rg -A 15 "func loadPrivateKeyPEM" cmd/start.go

Repository: agentuity/gravity

Length of output: 524


🏁 Script executed:

#!/bin/bash
# Check for tests or examples that show input format
find . -name "*test*.go" -type f -exec rg -l "privateKeyPEM\|loadPrivateKeyPEM" {} \;

Repository: agentuity/gravity

Length of output: 43


🏁 Script executed:

#!/bin/bash
# Check Go's crypto/x509 documentation via looking at how EC PARAMETERS blocks work
# EC PARAMETERS are sometimes generated in openssl EC key output
# Let's see if there are any examples in the codebase
rg -i "EC PARAMETERS|ec parameters" .

Repository: agentuity/gravity

Length of output: 43


🏁 Script executed:

#!/bin/bash
# Get the full loadPrivateKeyPEM function
rg -A 30 "func loadPrivateKeyPEM" cmd/start.go

Repository: agentuity/gravity

Length of output: 1116


🏁 Script executed:

#!/bin/bash
# Look at the parsePEMBundle implementation in stack.go to see the pattern
rg -A 25 "func parsePEMBundle" internal/stack/stack.go

Repository: agentuity/gravity

Length of output: 665


🏁 Script executed:

#!/bin/bash
# Check Go's crypto documentation online to confirm pem.Decode behavior
# Let's also look for any comments about key format expectations
rg -B 5 -A 5 "PKCS8\|ParseEC" cmd/start.go

Repository: agentuity/gravity

Length of output: 43


Loop through PEM blocks to find a private key block instead of parsing the first block unconditionally.

PEM files may contain EC PARAMETERS blocks, certificates, or other data before the private key. The current code calls pem.Decode() once and attempts to parse whatever block it returns, which fails if that block isn't a key. Search for a block matching "PRIVATE KEY" to ensure valid keys aren't rejected.

This pattern is already established in parsePEMBundle() in internal/stack/stack.go, which iterates through blocks and checks block.Type.

Suggested fix
-		block, _ := pem.Decode(privateKeyPEM)
-		if block == nil {
-			logger.Fatal("no PEM block found in --private-key")
-		}
+		var block *pem.Block
+		rest := privateKeyPEM
+		for {
+			var b *pem.Block
+			b, rest = pem.Decode(rest)
+			if b == nil {
+				break
+			}
+			if strings.Contains(b.Type, "PRIVATE KEY") {
+				block = b
+				break
+			}
+		}
+		if block == nil {
+			logger.Fatal("no private key PEM block found in --private-key")
+		}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
privateKeyPEM, err := loadPrivateKeyPEM(maybePrivateKey)
if err != nil {
logger.Fatal("failed to load private key: %v", err)
}
block, _ := pem.Decode(privateKeyPEM)
if block == nil {
logger.Fatal("no PEM block found in --private-key")
}
key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
if err != nil {
// Try SEC 1/EC format as fallback
key, err = x509.ParseECPrivateKey(block.Bytes)
if err != nil {
logger.Fatal("failed to parse private key: %v (pem length: %d)", err, len(privateKeyPEM))
}
}
privateKey, ok := key.(*ecdsa.PrivateKey)
if !ok {
logger.Fatal("private key is not an ECDSA key")
}
agent := stack.AgentMetadata{
OrgID: orgID,
ProjectID: projectID,
InstanceID: endpointID,
PrivateKey: privateKey,
}
privateKeyPEM, err := loadPrivateKeyPEM(maybePrivateKey)
if err != nil {
logger.Fatal("failed to load private key: %v", err)
}
var block *pem.Block
rest := privateKeyPEM
for {
var b *pem.Block
b, rest = pem.Decode(rest)
if b == nil {
break
}
if strings.Contains(b.Type, "PRIVATE KEY") {
block = b
break
}
}
if block == nil {
logger.Fatal("no private key PEM block found in --private-key")
}
key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
if err != nil {
// Try SEC 1/EC format as fallback
key, err = x509.ParseECPrivateKey(block.Bytes)
if err != nil {
logger.Fatal("failed to parse private key: %v (pem length: %d)", err, len(privateKeyPEM))
}
}
privateKey, ok := key.(*ecdsa.PrivateKey)
if !ok {
logger.Fatal("private key is not an ECDSA key")
}
agent := stack.AgentMetadata{
OrgID: orgID,
ProjectID: projectID,
InstanceID: endpointID,
PrivateKey: privateKey,
}
🤖 Prompt for AI Agents
In `@cmd/start.go` around lines 82 - 109, The current logic decodes only the first
PEM block from privateKeyPEM and fails if that block isn't a key; update the
parsing near pem.Decode(privateKeyPEM) to iterate with pem.Decode in a loop over
the remaining bytes and select a block whose block.Type contains "PRIVATE KEY"
(or ends with "PRIVATE KEY") before attempting x509.ParsePKCS8PrivateKey /
x509.ParseECPrivateKey; for each candidate block try PKCS#8 then EC parsing (as
the existing fallback does) and only break/assign privateKey when parsing
succeeds, otherwise continue to the next PEM block; if no private-key block
parses successfully, call logger.Fatal with an informative message (include
len(privateKeyPEM) as currently done).


ipv6Address := cnet.NewIPv6Address(cnet.GetRegion(""), cnet.NetworkHadron, agent.OrgID, agent.InstanceID, ipv4addr)
Expand All @@ -83,32 +119,42 @@ var rootCmd = &cobra.Command{
IPv6Addr: ipv6Address.String(),
LocalPort: localPort,
ProxyPort: proxyPort,
Token: token,
URL: gravityUrl,
Version: version,
}

provResp, err := stack.ProvisionGravity(ctx, logger, agent, urls)
if err != nil {
logger.Fatal("failed to provision gravity: %v", err)
}
tlsConfig, err := stack.GenerateCertificate(ctx, logger, provResp)
if err != nil {
logger.Fatal("failed to generate certificate: %v", err)
}
server, err := stack.StartServer(ctx, logger, tlsConfig, urls)
if err != nil {
logger.Fatal("failed to start server: %v", err)
}

netStack, linkEP, err := stack.CreateNetworkStack(logger, urls)
if err != nil {
logger.Fatal("failed to create network stack: %v", err)
}
defer netStack.Close()
defer linkEP.Close()

provider, client, err := stack.CreateNetworkProvider(ctx, logger, linkEP, provResp, urls, agent)
var server *http.Server
var serverMu sync.Mutex
provider, client, err := stack.CreateNetworkProvider(ctx, logger, linkEP, urls, agent, func(c *provider.Configuration) error {
serverMu.Lock()
oldServer := server
server = nil
serverMu.Unlock()
if oldServer != nil {
if err := oldServer.Shutdown(ctx); err != nil {
logger.Error("failed to shutdown server: %v", err)
}
}
tlsConfig, err := stack.GenerateCertificate(ctx, logger, c.MachineCertBundle)
if err != nil {
return fmt.Errorf("failed to generate certificate: %w", err)
}
newServer, err := stack.StartServer(ctx, logger, tlsConfig, urls)
if err != nil {
return fmt.Errorf("failed to start server: %w", err)
}
serverMu.Lock()
server = newServer
serverMu.Unlock()
Comment on lines +133 to +155
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

❓ Verification inconclusive

Script executed:

git ls-files cmd/ | head -20

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

cmd/start.go

Script executed:

cat -n cmd/start.go | sed -n '120,170p'

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

   120				LocalPort: localPort,
   121				ProxyPort: proxyPort,
   122				URL:       gravityUrl,
   123				Version:   version,
   124			}
   125	
   126			netStack, linkEP, err := stack.CreateNetworkStack(logger, urls)
   127			if err != nil {
   128				logger.Fatal("failed to create network stack: %v", err)
   129			}
   130			defer netStack.Close()
   131			defer linkEP.Close()
   132	
   133			var server *http.Server
   134			var serverMu sync.Mutex
   135			provider, client, err := stack.CreateNetworkProvider(ctx, logger, linkEP, urls, agent, func(c *provider.Configuration) error {
   136				serverMu.Lock()
   137				oldServer := server
   138				server = nil
   139				serverMu.Unlock()
   140				if oldServer != nil {
   141					if err := oldServer.Shutdown(ctx); err != nil {
   142						logger.Error("failed to shutdown server: %v", err)
   143					}
   144				}
   145				tlsConfig, err := stack.GenerateCertificate(ctx, logger, c.MachineCertBundle)
   146				if err != nil {
   147					return fmt.Errorf("failed to generate certificate: %w", err)
   148				}
   149				newServer, err := stack.StartServer(ctx, logger, tlsConfig, urls)
   150				if err != nil {
   151					return fmt.Errorf("failed to start server: %w", err)
   152				}
   153				serverMu.Lock()
   154				server = newServer
   155				serverMu.Unlock()
   156				return nil
   157			})
   158			if err != nil {
   159				logger.Fatal("failed to create network provider: %v", err)
   160			}
   161	
   162			// Wait for provider connection
   163			select {
   164			case <-ctx.Done():
   165				logger.Fatal("context done: %v", ctx.Err())
   166			case <-time.After(time.Second * 10):
   167				logger.Error("timed out waiting for provider connection")
   168				os.Exit(1)
   169			case <-provider.Connected:
   170				logger.Debug("✅ Connected to Gravity! Proxy is ready.")

Script executed:

cat -n cmd/start.go | sed -n '1,140p' | grep -E "(func |ctx\s*:=|ctx\s*=)"

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 1


Script executed:

cat -n cmd/start.go | head -80

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

     1	package cmd
     2	
     3	import (
     4		"context"
     5		"crypto/ecdsa"
     6		"crypto/x509"
     7		"encoding/base64"
     8		"encoding/pem"
     9		"fmt"
    10		"net/http"
    11		"os"
    12		"os/signal"
    13		"runtime/debug"
    14		"strings"
    15		"sync"
    16		"syscall"
    17		"time"
    18	
    19		"github.com/agentuity/go-common/gravity/provider"
    20		_logger "github.com/agentuity/go-common/logger"
    21		cnet "github.com/agentuity/go-common/network"
    22		csys "github.com/agentuity/go-common/sys"
    23		"github.com/agentuity/gravity/internal/heartbeat"
    24		"github.com/agentuity/gravity/internal/stack"
    25		"github.com/agentuity/gravity/internal/utils"
    26		"github.com/spf13/cobra"
    27	)
    28	
    29	var version string = "dev"
    30	
    31	var rootCmd = &cobra.Command{
    32		Use:   "gravity",
    33		Short: "Run the gravity client",
    34		Args:  cobra.NoArgs,
    35		Run: func(cmd *cobra.Command, args []string) {
    36			flags := cmd.Flags()
    37	
    38			var logLevel _logger.LogLevel
    39			logLevelStr, _ := flags.GetString("log-level")
    40			if logLevelStr == "" {
    41				logLevel = _logger.GetLevelFromEnv()
    42			} else {
    43				switch logLevelStr {
    44				case "info", "INFO":
    45					logLevel = _logger.LevelInfo
    46				case "debug", "DEBUG":
    47					logLevel = _logger.LevelDebug
    48				case "warn", "WARN":
    49					logLevel = _logger.LevelWarn
    50				case "trace", "TRACE":
    51					logLevel = _logger.LevelTrace
    52				case "error", "ERROR":
    53					logLevel = _logger.LevelError
    54				default:
    55					logLevel = _logger.LevelTrace
    56				}
    57			}
    58	
    59			logger := _logger.NewConsoleLogger(logLevel)
    60	
    61			ctx, cancel := signal.NotifyContext(cmd.Context(), os.Interrupt, syscall.SIGTERM)
    62			defer cancel()
    63	
    64			localPort, _ := flags.GetInt("port")
    65			orgID, _ := flags.GetString("org-id")
    66			projectID, _ := flags.GetString("project-id")
    67			maybePrivateKey, _ := flags.GetString("private-key")
    68			endpointID, _ := flags.GetString("endpoint-id")
    69			gravityUrl, _ := flags.GetString("url")
    70			healthCheck, _ := flags.GetBool("health-check")
    71			token, _ := flags.GetString("token")
    72	
    73			if token != "" {
    74				logger.Fatal("The --token flag is no longer supported. Please update your Agentuity CLI to the latest version.")
    75			}
    76	
    77			ipv4addr, err := utils.GetPrivateIPv4()
    78			if err != nil {
    79				logger.Fatal("failed to get private IPv4: %v", err)
    80			}

Script executed:

fd -t f "\.go$" internal/stack/ | head -20

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

internal/stack/certs.go
internal/stack/gravityClient.go
internal/stack/networkProvider.go
internal/stack/stack.go

Script executed:

rg "func CreateNetworkProvider" --type go -A 20

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

internal/stack/stack.go:func CreateNetworkProvider(
internal/stack/stack.go-	ctx context.Context,
internal/stack/stack.go-	logger _logger.Logger,
internal/stack/stack.go-	linkEP *channel.Endpoint,
internal/stack/stack.go-	urls UrlsMetadata,
internal/stack/stack.go-	agent AgentMetadata,
internal/stack/stack.go-	onConnect func(*provider.Configuration) error,
internal/stack/stack.go-) (*GravityClient, *gravity.GravityClient, error) {
internal/stack/stack.go-
internal/stack/stack.go-	// Egress pump to send outbound packets to Gravity
internal/stack/stack.go-	var network networkProvider
internal/stack/stack.go-
internal/stack/stack.go-	go func() {
internal/stack/stack.go-		for {
internal/stack/stack.go-			select {
internal/stack/stack.go-			case <-ctx.Done():
internal/stack/stack.go-				return
internal/stack/stack.go-			default:
internal/stack/stack.go-				pkt := linkEP.ReadContext(ctx)
internal/stack/stack.go-				if pkt == nil {
internal/stack/stack.go-					continue

Script executed:

rg "onConnect" internal/stack/stack.go -B 3 -A 3

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

	linkEP *channel.Endpoint,
	urls UrlsMetadata,
	agent AgentMetadata,
	onConnect func(*provider.Configuration) error,
) (*GravityClient, *gravity.GravityClient, error) {

	// Egress pump to send outbound packets to Gravity
--
	}()

	prov := GravityClient{
		onConnect: onConnect,
		logger:    logger,
		ep:        linkEP,
		Connected: make(chan struct{}, 1),

Script executed:

cat -n internal/stack/stack.go | wc -l

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

455

Script executed:

cat -n internal/stack/stack.go | sed -n '50,150p'

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

    50		Token     string
    51		URL       string
    52		Version   string
    53	}
    54	
    55	func ProvisionGravity(ctx context.Context, logger _logger.Logger, agent AgentMetadata, urls UrlsMetadata) (*proto.ProvisionResponse, error) {
    56	
    57		hostname, err := os.Hostname()
    58		if err != nil {
    59			return nil, fmt.Errorf("error getting hostname: %w", err)
    60		}
    61	
    62		requestObj := gravity.ProvisionRequest{
    63			Context:    ctx,
    64			GravityURL: urls.URL,
    65			InstanceID: agent.InstanceID,
    66			Region:     "unknown",
    67			Provider:   "other",
    68			PrivateIP:  urls.IPv4Addr,
    69			Token:      urls.Token,
    70			Hostname:   hostname,
    71			Ephemeral:  true,
    72			Capabilities: &proto.ClientCapabilities{
    73				DynamicHostname: true,
    74				// DynamicProjectRouting: "",
    75			},
    76		}
    77	
    78		logger.Debug("Provisioning gravity connection...")
    79		resp, err := gravity.Provision(requestObj)
    80		if err != nil {
    81			return nil, fmt.Errorf("failed to provision: %w", err)
    82		}
    83	
    84		return resp, nil
    85	}
    86	
    87	func GenerateCertificate(_ context.Context, logger _logger.Logger, bundle string) (*tls.Config, error) {
    88		certPEM, caCertPEM, keyPEM, err := parsePEMBundle(bundle)
    89		if err != nil {
    90			return nil, fmt.Errorf("failed to parse certificate bundle: %w", err)
    91		}
    92	
    93		// Set up TLS config
    94		cert, err := tls.X509KeyPair(certPEM, keyPEM)
    95		if err != nil {
    96			return nil, fmt.Errorf("failed to load certificate: %w", err)
    97		}
    98	
    99		// Parse the certificate to log details
   100		if len(cert.Certificate) > 0 {
   101			x509Cert, err := x509.ParseCertificate(cert.Certificate[0])
   102			if err == nil {
   103				cert.Leaf = x509Cert
   104				logger.Debug("Loaded certificate: CN=%s, SANs=%v, NotBefore=%v, NotAfter=%v",
   105					x509Cert.Subject.CommonName, x509Cert.DNSNames, x509Cert.NotBefore, x509Cert.NotAfter)
   106			} else {
   107				logger.Warn("Failed to parse certificate for logging: %v", err)
   108			}
   109		}
   110	
   111		caCertPool := x509.NewCertPool()
   112		if !caCertPool.AppendCertsFromPEM(caCertPEM) {
   113			return nil, fmt.Errorf("failed to parse CA certificate")
   114		}
   115	
   116		// Parse and log CA certificate details
   117		caCert, err := x509.ParseCertificate(caCertPEM)
   118		if err == nil {
   119			logger.Debug("Loaded CA certificate: CN=%s, Issuer=%s", caCert.Subject.CommonName, caCert.Issuer.CommonName)
   120		} else {
   121			// Try parsing as PEM
   122			block, _ := pem.Decode(caCertPEM)
   123			if block != nil {
   124				caCert, err = x509.ParseCertificate(block.Bytes)
   125				if err == nil {
   126					logger.Debug("Loaded CA certificate (PEM): CN=%s, Issuer=%s", caCert.Subject.CommonName, caCert.Issuer.CommonName)
   127				}
   128			}
   129		}
   130	
   131		tlsConfig := &tls.Config{
   132			Certificates:     []tls.Certificate{cert},
   133			RootCAs:          caCertPool,
   134			ClientCAs:        caCertPool, // Also set ClientCAs for mutual TLS
   135			ClientAuth:       tls.VerifyClientCertIfGiven,
   136			MinVersion:       tls.VersionTLS13,
   137			CurvePreferences: []tls.CurveID{tls.X25519, tls.X25519MLKEM768, tls.CurveP256},
   138			NextProtos:       []string{"h2", "http/1.1"},
   139			GetCertificate: func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
   140				remoteAddr := "<no-remote>"
   141				if hello != nil && hello.Conn != nil {
   142					remoteAddr = hello.Conn.RemoteAddr().String()
   143				}
   144				logger.Trace("TLS GetCertificate: ServerName=%s, SupportedProtos=%v, RemoteAddr=%v",
   145					hello.ServerName, hello.SupportedProtos, remoteAddr)
   146				return &cert, nil
   147			},
   148			VerifyPeerCertificate: func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {
   149				logger.Trace("TLS VerifyPeerCertificate: chains=%d, rawCerts=%d", len(verifiedChains), len(rawCerts))
   150				if len(rawCerts) > 0 {

Script executed:

rg "\.onConnect\(" internal/stack/ -B 5 -A 5

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

internal/stack/gravityClient.go-
internal/stack/gravityClient.go-func (p *GravityClient) Configure(config provider.Configuration) error {
internal/stack/gravityClient.go-	p.logger.Debug("configuring provider")
internal/stack/gravityClient.go-	p.config = config
internal/stack/gravityClient.go-	p.Connected <- struct{}{}
internal/stack/gravityClient.go:	return p.onConnect(&config)
internal/stack/gravityClient.go-}
internal/stack/gravityClient.go-
internal/stack/gravityClient.go-func (p *GravityClient) Deprovision(ctx context.Context, resourceID string, reason provider.DeprovisionReason) error {
internal/stack/gravityClient.go-	return nil
internal/stack/gravityClient.go-}

Script executed:

cat -n internal/stack/gravityClient.go

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

     1	package stack
     2	
     3	import (
     4		"context"
     5	
     6		"github.com/agentuity/go-common/gravity/proto"
     7		"github.com/agentuity/go-common/gravity/provider"
     8		"github.com/agentuity/go-common/logger"
     9		"gvisor.dev/gvisor/pkg/buffer"
    10		"gvisor.dev/gvisor/pkg/tcpip/link/channel"
    11		"gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
    12		"gvisor.dev/gvisor/pkg/tcpip/stack"
    13	)
    14	
    15	type GravityClient struct {
    16		ep        *channel.Endpoint
    17		logger    logger.Logger
    18		config    provider.Configuration
    19		Connected chan struct{}
    20		onConnect func(*provider.Configuration) error
    21	}
    22	
    23	func (p *GravityClient) Configure(config provider.Configuration) error {
    24		p.logger.Debug("configuring provider")
    25		p.config = config
    26		p.Connected <- struct{}{}
    27		return p.onConnect(&config)
    28	}
    29	
    30	func (p *GravityClient) Deprovision(ctx context.Context, resourceID string, reason provider.DeprovisionReason) error {
    31		return nil
    32	}
    33	
    34	func (p *GravityClient) Resources() []*proto.ExistingDeployment {
    35		return nil
    36	}
    37	
    38	func (p *GravityClient) SetMetricsCollector(collector provider.ProjectRuntimeStatsCollector) {
    39	}
    40	
    41	func (p *GravityClient) ProcessInPacket(payload []byte) {
    42		if p.ep == nil {
    43			p.logger.Error("ProcessInPacket called but endpoint is nil")
    44			return
    45		}
    46		pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{})
    47		view := buffer.NewView(len(payload))
    48		view.Write(payload)
    49		pkt.Data().AppendView(view)
    50		p.ep.InjectInbound(ipv6.ProtocolNumber, pkt)
    51		// Note: Don't call DecRef here - the netstack will handle the packet lifecycle
    52	}

Script executed:

rg "\.Configure\(" --type go | head -20

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0


Script executed:

rg "provider\.Configuration" --type go -B 3 -A 3 | head -50

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

internal/stack/gravityClient.go-type GravityClient struct {
internal/stack/gravityClient.go-	ep        *channel.Endpoint
internal/stack/gravityClient.go-	logger    logger.Logger
internal/stack/gravityClient.go:	config    provider.Configuration
internal/stack/gravityClient.go-	Connected chan struct{}
internal/stack/gravityClient.go:	onConnect func(*provider.Configuration) error
internal/stack/gravityClient.go-}
internal/stack/gravityClient.go-
internal/stack/gravityClient.go:func (p *GravityClient) Configure(config provider.Configuration) error {
internal/stack/gravityClient.go-	p.logger.Debug("configuring provider")
internal/stack/gravityClient.go-	p.config = config
internal/stack/gravityClient.go-	p.Connected <- struct{}{}
--
internal/stack/stack.go-	linkEP *channel.Endpoint,
internal/stack/stack.go-	urls UrlsMetadata,
internal/stack/stack.go-	agent AgentMetadata,
internal/stack/stack.go:	onConnect func(*provider.Configuration) error,
internal/stack/stack.go-) (*GravityClient, *gravity.GravityClient, error) {
internal/stack/stack.go-
internal/stack/stack.go-	// Egress pump to send outbound packets to Gravity
--
cmd/start.go-
cmd/start.go-		var server *http.Server
cmd/start.go-		var serverMu sync.Mutex
cmd/start.go:		provider, client, err := stack.CreateNetworkProvider(ctx, logger, linkEP, urls, agent, func(c *provider.Configuration) error {
cmd/start.go-			serverMu.Lock()
cmd/start.go-			oldServer := server
cmd/start.go-			server = nil

Script executed:

rg "func.*Shutdown" --type go -B 2 -A 5

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

internal/heartbeat/heartbeat.go-}
internal/heartbeat/heartbeat.go-
internal/heartbeat/heartbeat.go:func (s *Server) Shutdown() {
internal/heartbeat/heartbeat.go-	s.shutdownOnce.Do(func() {
internal/heartbeat/heartbeat.go-		if s.cancel != nil {
internal/heartbeat/heartbeat.go-			s.cancel()
internal/heartbeat/heartbeat.go-		}
internal/heartbeat/heartbeat.go-		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)

Script executed:

cat -n cmd/start.go | sed -n '135,160p'

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 0

stdout:

   135			provider, client, err := stack.CreateNetworkProvider(ctx, logger, linkEP, urls, agent, func(c *provider.Configuration) error {
   136				serverMu.Lock()
   137				oldServer := server
   138				server = nil
   139				serverMu.Unlock()
   140				if oldServer != nil {
   141					if err := oldServer.Shutdown(ctx); err != nil {
   142						logger.Error("failed to shutdown server: %v", err)
   143					}
   144				}
   145				tlsConfig, err := stack.GenerateCertificate(ctx, logger, c.MachineCertBundle)
   146				if err != nil {
   147					return fmt.Errorf("failed to generate certificate: %w", err)
   148				}
   149				newServer, err := stack.StartServer(ctx, logger, tlsConfig, urls)
   150				if err != nil {
   151					return fmt.Errorf("failed to start server: %w", err)
   152				}
   153				serverMu.Lock()
   154				server = newServer
   155				serverMu.Unlock()
   156				return nil
   157			})
   158			if err != nil {
   159				logger.Fatal("failed to create network provider: %v", err)
   160			}

Script executed:

web_search
http.Server.Shutdown behavior with context timeout

Repository: agentuity/gravity


Repository: agentuity/gravity
Exit code: 127

stderr:

/bin/bash: line 1: web_search: command not found
/bin/bash: line 2: http.Server.Shutdown: command not found

http.Server.Shutdown context timeout behavior golang


In Go’s net/http, (*http.Server).Shutdown(ctx) is graceful and the ctx timeout only bounds how long Shutdown will wait, not how long active requests are allowed to run.

What Shutdown does

Shutdown performs this sequence: (1) close listeners (stop accepting new connections), (2) close idle connections, (3) wait for remaining connections to become idle so they can be closed. If ctx expires first, Shutdown returns ctx.Err() (typically context.DeadlineExceeded). It also does not close/wait for hijacked connections (e.g., WebSockets). [1]

What happens when the context times out

From the actual implementation, when the context is done, Shutdown just returns ctx.Err(); it does not forcibly close active connections for you. Active handlers/connections can continue running until they finish or until something else closes them. [2]

Practical implications

  • Your in-flight request r.Context() is not automatically canceled just because Shutdown was called; Shutdown is explicitly “without interrupting any active connections.” [1][3]
  • If you need a hard cutoff after the graceful period, a common pattern is: call Shutdown(timeoutCtx), and if it returns context.DeadlineExceeded, then call srv.Close() to force-close remaining connections. (Close closes connections in StateNew, StateActive, and StateIdle.) [1]

Sources:
[1] Go net/http docs for Server.Shutdown / Server.Close (pkg.go.dev) (pkg.go.dev)
[2] Go net/http Server.Shutdown implementation showing case <-ctx.Done(): return ctx.Err() (go.googlesource.com) (go.googlesource.com)
[3] Discussion confirming Shutdown doesn’t interrupt active connections / doesn’t cancel request contexts (Stack Overflow) (stackoverflow.com)

Citations:


Use a bounded shutdown context to prevent blocking reconfiguration.

During provider reconfiguration, oldServer.Shutdown(ctx) uses the long-lived root context which can block indefinitely if a client connection is stuck. This blocks the entire onConnect callback, preventing the new server from starting and effectively halting provider reconfiguration.

Use a timeout context for shutdown to ensure the callback completes within a reasonable grace period:

⏱️ Use a bounded shutdown context
 			if oldServer != nil {
+				shutdownCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+				defer cancel()
-				if err := oldServer.Shutdown(ctx); err != nil {
+				if err := oldServer.Shutdown(shutdownCtx); err != nil {
 					logger.Error("failed to shutdown server: %v", err)
 				}
 			}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
var server *http.Server
var serverMu sync.Mutex
provider, client, err := stack.CreateNetworkProvider(ctx, logger, linkEP, urls, agent, func(c *provider.Configuration) error {
serverMu.Lock()
oldServer := server
server = nil
serverMu.Unlock()
if oldServer != nil {
if err := oldServer.Shutdown(ctx); err != nil {
logger.Error("failed to shutdown server: %v", err)
}
}
tlsConfig, err := stack.GenerateCertificate(ctx, logger, c.MachineCertBundle)
if err != nil {
return fmt.Errorf("failed to generate certificate: %w", err)
}
newServer, err := stack.StartServer(ctx, logger, tlsConfig, urls)
if err != nil {
return fmt.Errorf("failed to start server: %w", err)
}
serverMu.Lock()
server = newServer
serverMu.Unlock()
var server *http.Server
var serverMu sync.Mutex
provider, client, err := stack.CreateNetworkProvider(ctx, logger, linkEP, urls, agent, func(c *provider.Configuration) error {
serverMu.Lock()
oldServer := server
server = nil
serverMu.Unlock()
if oldServer != nil {
shutdownCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
if err := oldServer.Shutdown(shutdownCtx); err != nil {
logger.Error("failed to shutdown server: %v", err)
}
}
tlsConfig, err := stack.GenerateCertificate(ctx, logger, c.MachineCertBundle)
if err != nil {
return fmt.Errorf("failed to generate certificate: %w", err)
}
newServer, err := stack.StartServer(ctx, logger, tlsConfig, urls)
if err != nil {
return fmt.Errorf("failed to start server: %w", err)
}
serverMu.Lock()
server = newServer
serverMu.Unlock()
🤖 Prompt for AI Agents
In `@cmd/start.go` around lines 133 - 155, The shutdown of oldServer in the
provider reconfigure callback currently calls oldServer.Shutdown(ctx) with the
long-lived root context and can block indefinitely; change that call to use a
bounded timeout context (e.g., via context.WithTimeout(parentCtx, timeout)) so
the shutdown completes or aborts within a fixed grace period, ensure you call
cancel() and handle/log Shutdown errors including context deadline exceeded, and
keep the rest of the callback (GenerateCertificate, StartServer, and assignment
to server under serverMu) unchanged; update the closure where oldServer is set
and Shutdown is invoked (the variables server, serverMu, oldServer in the
CreateNetworkProvider callback).

return nil
})
if err != nil {
logger.Fatal("failed to create network provider: %v", err)
}
Expand Down Expand Up @@ -161,8 +207,14 @@ var rootCmd = &cobra.Command{

shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
defer shutdownCancel()
if err := server.Shutdown(shutdownCtx); err != nil {
logger.Error("error shutting down proxy server: %v", err)
serverMu.Lock()
serverToShutdown := server
server = nil
serverMu.Unlock()
if serverToShutdown != nil {
if err := serverToShutdown.Shutdown(shutdownCtx); err != nil {
logger.Error("error shutting down proxy server: %v", err)
}
}
},
}
Expand Down Expand Up @@ -202,12 +254,52 @@ func init() {
rootCmd.Flags().IntP("port", "p", 3500, "Local port for the proxy")
rootCmd.Flags().StringP("org-id", "o", os.Getenv("AGENTUITY_CLOUD_ORG_ID"), "Organization ID")
rootCmd.Flags().StringP("project-id", "i", os.Getenv("AGENTUITY_CLOUD_PROJECT_ID"), "Project ID")
rootCmd.Flags().StringP("token", "t", os.Getenv("AGENTUITY_SDK_KEY"), "Project SDK Key")
rootCmd.Flags().StringP("token", "t", "", "Deprecated")
rootCmd.Flags().StringP("endpoint-id", "e", "", "The endpoint id")
rootCmd.Flags().StringP("url", "u", "grpc://devmode.agentuity.com", "The gravity url")
rootCmd.Flags().String("log-level", "", "The log level to use")
rootCmd.Flags().Bool("health-check", false, "Enable health check server for heartbeat monitoring")
rootCmd.Flags().String("private-key", "", "EC private key for authentication (PEM string, base64-encoded PEM, or file path)")

// Mark required flags that must be passed in
rootCmd.MarkFlagRequired("endpoint-id")
rootCmd.MarkFlagRequired("private-key")
}

// loadPrivateKeyPEM loads a private key from various formats:
// - PEM-encoded string (starts with "-----BEGIN")
// - Base64-encoded PEM
// - File path containing PEM data
func loadPrivateKeyPEM(input string) ([]byte, error) {
if input == "" {
return nil, fmt.Errorf("private key is required")
}

// Check if it's already PEM format
if strings.HasPrefix(strings.TrimSpace(input), "-----BEGIN") {
return []byte(input), nil
}

// Check if it's a file path
if data, err := os.ReadFile(input); err == nil {
if strings.HasPrefix(strings.TrimSpace(string(data)), "-----BEGIN") {
return data, nil
}
// File contains base64
decoded, err := base64.StdEncoding.DecodeString(strings.TrimSpace(string(data)))
if err == nil && strings.HasPrefix(string(decoded), "-----BEGIN") {
return decoded, nil
}
return nil, fmt.Errorf("file does not contain valid PEM or base64-encoded PEM")
}

// Try base64 decode
decoded, err := base64.StdEncoding.DecodeString(strings.TrimSpace(input))
if err != nil {
return nil, fmt.Errorf("input is not valid PEM, base64, or file path: %w", err)
}
if !strings.HasPrefix(string(decoded), "-----BEGIN") {
return nil, fmt.Errorf("base64-decoded content is not valid PEM")
}
return decoded, nil
}
8 changes: 2 additions & 6 deletions go.mod
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
module github.com/agentuity/gravity

go 1.25.4
go 1.25.5

require (
github.com/agentuity/go-common v1.0.127
github.com/agentuity/go-common v1.0.142
github.com/spf13/cobra v1.10.1
gvisor.dev/gvisor v0.0.0-20240423190808-9d7a357edefe
)
Expand Down Expand Up @@ -51,7 +51,3 @@ require (
google.golang.org/grpc v1.76.0 // indirect
google.golang.org/protobuf v1.36.10 // indirect
)

replace github.com/agentuity/cli => ..

replace pkg.agentuity.com/cloud/go => github.com/agentuity/go-common v1.0.111
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
github.com/agentuity/go-common v1.0.127 h1:KW1QJHtVizlyTqqq4LEmCnpJl9q0ZzlwxjFnDgjFMzM=
github.com/agentuity/go-common v1.0.127/go.mod h1:C6bHmnDizFTbkB0jL7m3YVJTcaPKfc5bRRIYEva4ZTQ=
github.com/agentuity/go-common v1.0.142 h1:hcrf9y0pYSne47B1ReM+VZdhuAJxDmJn1mM1VXiBemw=
github.com/agentuity/go-common v1.0.142/go.mod h1:rvM6VhfltrwrfbjKjNToBQ7kNx3lR6955+n/G1C8soM=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/cockroachdb/errors v1.12.0 h1:d7oCs6vuIMUQRVbi6jWWWEJZahLCfJpnJSVobd1/sUo=
Expand Down
70 changes: 70 additions & 0 deletions internal/stack/certs.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
package stack

import (
"net/url"
"strings"

"github.com/agentuity/go-common/logger"
)

const gravityPublicCACertProd = `-----BEGIN CERTIFICATE-----
MIIBvDCCAWKgAwIBAgIIQpGyfo9xbKUwCgYIKoZIzj0EAwIwQjELMAkGA1UEBhMC
VVMxFzAVBgNVBAoTDkFnZW50dWl0eSwgSW5jMRowGAYDVQQDExFBZ2VudHVpdHkg
Um9vdCBDQTAeFw0yNTA4MjgyMDAyNDJaFw0zNTA4MjgyMTAyNDJaMEIxCzAJBgNV
BAYTAlVTMRcwFQYDVQQKEw5BZ2VudHVpdHksIEluYzEaMBgGA1UEAxMRQWdlbnR1
aXR5IFJvb3QgQ0EwWTATBgcqhkjOPQIBBggqhkjOPQMBBwNCAASuQLIe47OC2EtE
3cwkXJ3siuBeoi5FuJ5wmny1BqA60FrHqHnCiYHIZSyv79WvGs6NAxlJsCQbBtcE
FNvrFfUso0IwQDAOBgNVHQ8BAf8EBAMCAQYwDwYDVR0TAQH/BAUwAwEB/zAdBgNV
HQ4EFgQURmQweSNnxsTZ/ckG+A08IW7S+fwwCgYIKoZIzj0EAwIDSAAwRQIhAJzN
k5ZrKQPMCAEh1zBuUbWTcuRCnqdk583gcCBkUo58AiBny+nX/KLE46z1B1NK8qg9
/K75YibTFYFFQAXMF10aNg==
-----END CERTIFICATE-----`

const gravityPublicCACertDev = `-----BEGIN CERTIFICATE-----
MIIBtTCCAVugAwIBAgIBATAKBggqhkjOPQQDAjBCMQswCQYDVQQGEwJVUzEXMBUG
A1UEChMOQWdlbnR1aXR5LCBJbmMxGjAYBgNVBAMTEUFnZW50dWl0eSBSb290IENB
MB4XDTI1MDgyNzIwNTIxM1oXDTM1MDgyNzIxNTIxM1owQjELMAkGA1UEBhMCVVMx
FzAVBgNVBAoTDkFnZW50dWl0eSwgSW5jMRowGAYDVQQDExFBZ2VudHVpdHkgUm9v
dCBDQTBZMBMGByqGSM49AgEGCCqGSM49AwEHA0IABMleqtFE+8V/nC2eUEOEkfX+
SZDT9+QSIr2B2HoCllDVnyzOKZFGxURvr90dKa/SUiwfS6LjZp6xNUv2gP7GOIWj
QjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNVHRMBAf8EBTADAQH/MB0GA1UdDgQWBBSg
O9ViNhPYNISETcjqCsgs+TN80TAKBggqhkjOPQQDAgNIADBFAiAPcCANK2aU/Q/t
J4R6Su68r+iGJKpONFOiD3DX3p8oMgIhAI32feTJZwUVkNtjTYBRxoYUdZlyobp1
lTpZE2H0bHCM
-----END CERTIFICATE-----
`

// GravityCACertificate returns the Agentuity gravity CA certificate
func GravityCACertificate(logger logger.Logger, gravityURL string) string {
// Select the appropriate hash based on the API URL
if isDevelopmentGravityURL(gravityURL) {
logger.Debug("Using development CA")
return gravityPublicCACertDev
}

logger.Debug("Using production CA")
return gravityPublicCACertProd
}

// isDevelopmentGravityURL determines if we're in a development environment based on the API URL
func isDevelopmentGravityURL(apiURL string) bool {
parsedURL, err := url.Parse(apiURL)
if err != nil {
return false
}

// Extract hostname, stripping any port
hostname := parsedURL.Hostname()

// Exact matches for development environments
if hostname == "localhost" || hostname == "127.0.0.1" || hostname == "host.docker" {
return true
}

// Safe suffix check for agentuity.io domain
if hostname == "agentuity.io" || strings.HasSuffix(hostname, ".agentuity.io") {
return true
}

return false
}
3 changes: 2 additions & 1 deletion internal/stack/gravityClient.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,14 @@ type GravityClient struct {
logger logger.Logger
config provider.Configuration
Connected chan struct{}
onConnect func(*provider.Configuration) error
}

func (p *GravityClient) Configure(config provider.Configuration) error {
p.logger.Debug("configuring provider")
p.config = config
p.Connected <- struct{}{}
return nil
return p.onConnect(&config)
}

func (p *GravityClient) Deprovision(ctx context.Context, resourceID string, reason provider.DeprovisionReason) error {
Expand Down
Loading