RC-1

cmd/loadtest/main.go

@@ -0,0 +1,239 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"math/big"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"neuralhunt/internal/auth"
+	"neuralhunt/internal/core"
+)
+
+var b64 = base64.RawURLEncoding
+
+type apiClient struct {
+	base       string
+	hc         *http.Client
+	token, cid string
+	key        *ecdsa.PrivateKey
+}
+type taskDTO struct {
+	ID         string `json:"id"`
+	PublicSeed string `json:"public_seed"`
+	RangeBits  int    `json:"range_bits"`
+	NextSeq    int64  `json:"next_seq"`
+	SubmitSec  int    `json:"client_submit_interval_sec"`
+	Paused     bool   `json:"paused"`
+}
+
+func pad32(x *big.Int) []byte {
+	b := x.Bytes()
+	out := make([]byte, 32)
+	copy(out[32-len(b):], b)
+	return out
+}
+func signRaw(k *ecdsa.PrivateKey, msg string) (string, error) {
+	h := sha256.Sum256([]byte(msg))
+	r, s, err := ecdsa.Sign(rand.Reader, k, h[:])
+	if err != nil {
+		return "", err
+	}
+	raw := append(pad32(r), pad32(s)...)
+	return b64.EncodeToString(raw), nil
+}
+func (c *apiClient) do(method, path string, body any, out any) error {
+	var rd io.Reader
+	if body != nil {
+		b, _ := json.Marshal(body)
+		rd = bytes.NewReader(b)
+	}
+	req, _ := http.NewRequest(method, c.base+path, rd)
+	req.Header.Set("Content-Type", "application/json")
+	if c.token != "" {
+		req.Header.Set("Authorization", "Bearer "+c.token)
+	}
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
+	if resp.StatusCode/100 != 2 {
+		return fmt.Errorf("%s: %s", resp.Status, string(b))
+	}
+	if out != nil {
+		return json.Unmarshal(b, out)
+	}
+	return nil
+}
+func (c *apiClient) authn() error {
+	k, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		return err
+	}
+	c.key = k
+	j := auth.PublicJWK{Kty: "EC", Crv: "P-256", X: b64.EncodeToString(pad32(k.X)), Y: b64.EncodeToString(pad32(k.Y)), Ext: true}
+	var ch struct {
+		ClientID  string `json:"client_id"`
+		Challenge string `json:"challenge"`
+	}
+	if err = c.do("POST", "/api/auth/challenge", map[string]any{"public_jwk": j}, &ch); err != nil {
+		return err
+	}
+	c.cid = ch.ClientID
+	sig, _ := signRaw(k, "login|"+ch.Challenge+"|"+c.cid)
+	var lg struct {
+		Token string `json:"token"`
+	}
+	if err = c.do("POST", "/api/auth/login", map[string]any{"public_jwk": j, "challenge": ch.Challenge, "signature": sig}, &lg); err != nil {
+		return err
+	}
+	c.token = lg.Token
+	return nil
+}
+func (c *apiClient) current() (taskDTO, error) {
+	var t taskDTO
+	err := c.do("GET", "/api/tasks/current", nil, &t)
+	return t, err
+}
+func (c *apiClient) ws(ctx context.Context, maxNodes int) (*websocket.Conn, error) {
+	u, err := url.Parse(c.base)
+	if err != nil {
+		return nil, err
+	}
+	scheme := "ws"
+	if u.Scheme == "https" {
+		scheme = "wss"
+	}
+	q := url.Values{}
+	q.Set("token", c.token)
+	q.Set("max_nodes", fmt.Sprint(maxNodes))
+	wu := scheme + "://" + u.Host + "/api/ws?" + q.Encode()
+	conn, _, err := websocket.DefaultDialer.DialContext(ctx, wu, nil)
+	return conn, err
+}
+
+func main() {
+	base := flag.String("url", "http://127.0.0.1:8080", "server URL")
+	clients := flag.Int("clients", 1000, "virtual clients")
+	ramp := flag.Duration("ramp", 30*time.Second, "connection ramp")
+	duration := flag.Duration("duration", 2*time.Minute, "test duration after ramp")
+	nodes := flag.Int("max-nodes", 250, "snapshot budget per client")
+	flag.Parse()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	var connected, guesses, errs atomic.Uint64
+	var wg sync.WaitGroup
+	start := time.Now()
+	step := time.Duration(0)
+	if *clients > 0 {
+		step = *ramp / time.Duration(*clients)
+	}
+	for i := 0; i < *clients; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			if step > 0 {
+				select {
+				case <-ctx.Done():
+					return
+				case <-time.After(step * time.Duration(i)):
+				}
+			}
+			c := &apiClient{base: strings.TrimRight(*base, "/"), hc: &http.Client{Timeout: 10 * time.Second}}
+			if err := c.authn(); err != nil {
+				errs.Add(1)
+				return
+			}
+			t, err := c.current()
+			if err != nil {
+				errs.Add(1)
+				return
+			}
+			ws, err := c.ws(ctx, *nodes)
+			if err != nil {
+				errs.Add(1)
+				return
+			}
+			defer ws.Close()
+			connected.Add(1)
+			go func() {
+				for {
+					if _, _, e := ws.ReadMessage(); e != nil {
+						return
+					}
+				}
+			}()
+			interval := time.Duration(t.SubmitSec) * time.Second
+			if interval <= 0 {
+				interval = 11 * time.Second
+			}
+			ticker := time.NewTicker(interval)
+			defer ticker.Stop()
+			seq := t.NextSeq
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-ticker.C:
+					if t.Paused {
+						continue
+					}
+					guess := core.ExpectedGuess(t.ID, t.PublicSeed, c.cid, seq, t.RangeBits)
+					sig, _ := signRaw(c.key, fmt.Sprintf("guess|%s|%d|%s", t.ID, seq, guess))
+					var ok bool
+					err := c.do("POST", "/api/tasks/"+t.ID+"/guess", map[string]any{"seq": seq, "guess": guess, "signature": sig}, &ok)
+					if err != nil {
+						errs.Add(1)
+						nt, e := c.current()
+						if e == nil {
+							t = nt
+							seq = t.NextSeq
+						}
+						continue
+					}
+					guesses.Add(1)
+					seq++
+					if ok {
+						nt, e := c.current()
+						if e == nil {
+							t = nt
+							seq = t.NextSeq
+						}
+					}
+				}
+			}
+		}(i)
+	}
+	tick := time.NewTicker(5 * time.Second)
+	defer tick.Stop()
+	end := time.After(*ramp + *duration)
+	for {
+		select {
+		case <-end:
+			cancel()
+			wg.Wait()
+			fmt.Printf("done connected=%d accepted_guesses=%d errors=%d elapsed=%s\n", connected.Load(), guesses.Load(), errs.Load(), time.Since(start).Round(time.Second))
+			return
+		case <-tick.C:
+			log.Printf("connected=%d guesses=%d errors=%d", connected.Load(), guesses.Load(), errs.Load())
+		}
+	}
+}