add agent-bench

cmd/agent-bench/go.mod

@@ -0,0 +1,12 @@
+module agent-bench
+
+go 1.22.3
+
+require (
+	github.com/Microsoft/go-winio v0.6.2
+	github.com/davidmz/go-pageant v1.0.2
+	github.com/kelseyhightower/envconfig v1.4.0
+	golang.org/x/crypto v0.26.0
+)
+
+require golang.org/x/sys v0.23.0 // indirect

cmd/agent-bench/go.sum

@@ -0,0 +1,18 @@
+github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
+github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
+github.com/davidmz/go-pageant v1.0.2 h1:bPblRCh5jGU+Uptpz6LgMZGD5hJoOt7otgT454WvHn0=
+github.com/davidmz/go-pageant v1.0.2/go.mod h1:P2EDDnMqIwG5Rrp05dTRITj9z2zpGcD9efWSkTNKLIE=
+github.com/kelseyhightower/envconfig v1.4.0 h1:Im6hONhd3pLkfDFsbRgu68RDNkGF1r3dvMUtDTo2cv8=
+github.com/kelseyhightower/envconfig v1.4.0/go.mod h1:cccZRl6mQpaq41TPp5QxidR+Sa3axMbJDNb//FQX6Gg=
+golang.org/x/crypto v0.0.0-20210513164829-c07d793c2f9a/go.mod h1:P+XmwS30IXTQdn5tA2iutPOUgjI07+tq3H3K9MVA1s8=
+golang.org/x/crypto v0.26.0 h1:RrRspgV4mU+YwB4FYnuBoKsUapNIL5cohGAmSH3azsw=
+golang.org/x/crypto v0.26.0/go.mod h1:GY7jblb9wI+FOo5y8/S2oY4zWP07AkOJ4+jxCqdqn54=
+golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
+golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.23.0 h1:YfKFowiIMvtgl1UERQoTPPToxltDeZfbj4H7dVUCwmM=
+golang.org/x/sys v0.23.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/term v0.23.0 h1:F6D4vR+EHoL9/sWAWgAR1H2DcHr4PareCbAaCo1RpuU=
+golang.org/x/term v0.23.0/go.mod h1:DgV24QBUrK6jhZXl+20l6UWznPlwAHm1Q1mGHtydmSk=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

cmd/agent-bench/main.go

@@ -0,0 +1,157 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"sort"
+	"sync"
+	"time"
+
+	"github.com/kelseyhightower/envconfig"
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
+)
+
+type Specification struct {
+	PERSISTENT  bool `default:"false"`
+	CONCURRENCY int  `default:"10"`
+	RUN_COUNT   int  `default:"1000"`
+}
+
+type sshAgent interface {
+	agent.Agent
+	Close() error
+}
+type Agent struct {
+	agent.ExtendedAgent
+	net.Conn
+}
+
+type exAgent struct {
+	agent.Agent
+}
+
+func (e *exAgent) SignWithFlags(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (*ssh.Signature, error) {
+	return nil, nil
+}
+
+func (e *exAgent) Extension(string, []byte) ([]byte, error) {
+	return nil, nil
+}
+
+func main() {
+	s := Specification{}
+	err := envconfig.Process("", &s)
+	if err != nil {
+		log.Fatal(err)
+	}
+	taskCh := make(chan struct{})
+	doneCh := make(chan []time.Duration, s.CONCURRENCY)
+
+	var wg sync.WaitGroup
+	for i := 0; i < s.CONCURRENCY; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			worker(taskCh, doneCh, s.PERSISTENT)
+		}()
+	}
+
+	start := time.Now()
+	go func() {
+		for i := 0; i < s.RUN_COUNT; i++ {
+			taskCh <- struct{}{}
+		}
+		close(taskCh)
+	}()
+
+	go func() {
+		wg.Wait()
+		close(doneCh)
+	}()
+	var allExecutionTimes []time.Duration
+	for times := range doneCh {
+		allExecutionTimes = append(allExecutionTimes, times...)
+	}
+
+	totalTime := time.Duration(0)
+	var minTime, maxTime time.Duration
+	minTime = allExecutionTimes[0]
+	maxTime = allExecutionTimes[0]
+
+	for _, t := range allExecutionTimes {
+		totalTime += t
+		if t < minTime {
+			minTime = t
+		}
+		if t > maxTime {
+			maxTime = t
+		}
+	}
+
+	averageTime := totalTime / time.Duration(len(allExecutionTimes))
+
+	sort.Slice(allExecutionTimes, func(i, j int) bool {
+		return allExecutionTimes[i] < allExecutionTimes[j]
+	})
+	p99Time := allExecutionTimes[int(float64(len(allExecutionTimes))*0.99)-1]
+
+	fmt.Printf("Real Time: %v\n", time.Since(start))
+	fmt.Printf("Total Executions: %d\n", s.RUN_COUNT)
+	fmt.Printf("Concurrency: %d\n", s.CONCURRENCY)
+	fmt.Printf("Persistent Mode: %v\n", s.PERSISTENT)
+	fmt.Printf("Total Time: %v\n", totalTime)
+	fmt.Printf("Average Execution Time: %v\n", averageTime)
+	fmt.Printf("Min Execution Time: %v\n", minTime)
+	fmt.Printf("Max Execution Time: %v\n", maxTime)
+	fmt.Printf("99th Percentile Execution Time: %v\n", p99Time)
+}
+
+func worker(taskCh <-chan struct{}, doneCh chan<- []time.Duration, persistent bool) {
+	var executionTimes []time.Duration
+
+	var agentClient sshAgent
+	var key *agent.Key
+
+	a, err := newAgent()
+	if err != nil {
+		log.Fatal(err)
+	}
+	keys, err := a.List()
+	if err != nil {
+		log.Fatalf("Failed to list keys: %v", err)
+	}
+	if len(keys) == 0 {
+		log.Fatalf("No keys found in SSH agent")
+	}
+	key = keys[0]
+	a.Close()
+	agentClient = nil
+
+	log.Printf("key:%s", key.String())
+
+	for range taskCh {
+		start := time.Now()
+		if agentClient == nil {
+			agentClient, err = newAgent()
+			if err != nil {
+				log.Fatal(err)
+			}
+		}
+		data := []byte("Benchmark data")
+		_, err := agentClient.Sign(key, data)
+		if !persistent {
+			agentClient.Close()
+			agentClient = nil
+		}
+		if err != nil {
+			log.Printf("Failed to sign data: %v", err)
+			continue
+		}
+		duration := time.Since(start)
+		executionTimes = append(executionTimes, duration)
+	}
+
+	doneCh <- executionTimes
+}

cmd/agent-bench/unix.go

@@ -0,0 +1,42 @@
+//go:build unix
+
+package main
+
+import (
+	"log"
+	"net"
+	"os"
+
+	"golang.org/x/crypto/ssh/agent"
+)
+
+func listKeys() {
+	socketPath := os.Getenv("SSH_AUTH_SOCK")
+	conn, err := net.Dial("unix", socketPath)
+	if err != nil {
+		log.Fatal(err)
+	}
+	agentClient := agent.NewClient(conn)
+	list, err := agentClient.List()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	for _, key := range list {
+		log.Println(key.String())
+	}
+}
+
+func NewUnixDomain() (sshAgent, error) {
+	socketPath := os.Getenv("SSH_AUTH_SOCK")
+	conn, err := net.Dial("unix", socketPath)
+	if err != nil {
+		log.Fatal(err)
+	}
+	a := agent.NewClient(conn)
+	return &Agent{ExtendedAgent: &exAgent{a}, Conn: conn}, nil
+}
+
+func newAgent() (sshAgent, error) {
+	return NewUnixDomain()
+}

cmd/agent-bench/win.go

@@ -0,0 +1,44 @@
+//go:build windows
+
+package main
+
+import (
+	"errors"
+
+	"github.com/Microsoft/go-winio"
+	"github.com/davidmz/go-pageant"
+	"golang.org/x/crypto/ssh/agent"
+)
+
+const (
+	sshAgentPipe = `\\.\pipe\openssh-ssh-agent`
+)
+
+func NewPageant() (sshAgent, error) {
+	ok := pageant.Available()
+	if !ok {
+		return nil, errors.New("pageant is not available")
+	}
+	p := pageant.New()
+	return &Agent{ExtendedAgent: &exAgent{p}}, nil
+}
+
+func (a *Agent) Close() error {
+	if a.Conn != nil {
+		a.Conn.Close()
+	}
+	return nil
+}
+
+func NewNamedPipe() (sshAgent, error) {
+	conn, err := winio.DialPipe(sshAgentPipe, nil)
+	if err != nil {
+		return nil, err
+	}
+	a := agent.NewClient(conn)
+	return &Agent{ExtendedAgent: &exAgent{a}, Conn: conn}, nil
+}
+
+func newAgent() (sshAgent, error) {
+	return NewNamedPipe()
+}