RC-6

cmd/loadtest/main.go

@@ -83,6 +83,34 @@ func (c *apiClient) do(method, path string, body any, out any) error {
 	}
 	return nil
 }
+
+func leadingZeroBitsLoad(b []byte) int {
+	n := 0
+	for _, x := range b {
+		if x == 0 {
+			n += 8
+			continue
+		}
+		for m := byte(0x80); m != 0 && x&m == 0; m >>= 1 {
+			n++
+		}
+		break
+	}
+	return n
+}
+
+func solveProofLoad(challenge, cid string, bits int) string {
+	if bits <= 0 {
+		return ""
+	}
+	for i := uint64(0); ; i++ {
+		counter := fmt.Sprint(i)
+		h := sha256.Sum256([]byte("nh-pow-v1|" + challenge + "|" + cid + "|" + counter))
+		if leadingZeroBitsLoad(h[:]) >= bits {
+			return counter
+		}
+	}
+}
 func (c *apiClient) authn() error {
 	k, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 	if err != nil {
@@ -91,8 +119,9 @@ func (c *apiClient) authn() error {
 	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"`
+		ClientID        string `json:"client_id"`
+		Challenge       string `json:"challenge"`
+		ProofOfWorkBits int    `json:"proof_of_work_bits"`
 	}
 	if err = c.do("POST", "/api/auth/challenge", map[string]any{"public_jwk": j}, &ch); err != nil {
 		return err
@@ -102,7 +131,8 @@ func (c *apiClient) authn() error {
 	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 {
+	pow := solveProofLoad(ch.Challenge, c.cid, ch.ProofOfWorkBits)
+	if err = c.do("POST", "/api/auth/login", map[string]any{"public_jwk": j, "challenge": ch.Challenge, "signature": sig, "proof_of_work_counter": pow}, &lg); err != nil {
 		return err
 	}
 	c.token = lg.Token
@@ -123,10 +153,11 @@ func (c *apiClient) ws(ctx context.Context, maxNodes int) (*websocket.Conn, erro
 		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)
+	h := http.Header{}
+	h.Set("Authorization", "Bearer "+c.token)
+	conn, _, err := websocket.DefaultDialer.DialContext(ctx, wu, h)
 	return conn, err
 }