RC-13

cmd/client/api.go

@@ -18,6 +18,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -64,12 +65,15 @@ func (e *apiError) Code() string {
 }
 
 type apiClient struct {
-	base  string
-	hc    *http.Client
-	token string
-	cid   string
-	id    identityFile
-	key   *ecdsa.PrivateKey
+	base string
+	hc   *http.Client
+
+	authMu sync.Mutex
+	mu     sync.RWMutex
+	token  string
+	cid    string
+	id     identityFile
+	key    *ecdsa.PrivateKey
 }
 
 func newAPI(base string, id identityFile, key *ecdsa.PrivateKey) *apiClient {
@@ -81,7 +85,19 @@ func newAPI(base string, id identityFile, key *ecdsa.PrivateKey) *apiClient {
 	}
 }
 
-func (c *apiClient) do(ctx context.Context, method, path string, body, out any) error {
+func (c *apiClient) tokenValue() string {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.token
+}
+
+func (c *apiClient) clientID() string {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.cid
+}
+
+func (c *apiClient) doRaw(ctx context.Context, method, path string, body, out any) error {
 	var rd io.Reader
 	if body != nil {
 		b, err := json.Marshal(body)
@@ -97,8 +113,8 @@ func (c *apiClient) do(ctx context.Context, method, path string, body, out any)
 	if body != nil {
 		req.Header.Set("Content-Type", "application/json")
 	}
-	if c.token != "" {
-		req.Header.Set("Authorization", "Bearer "+c.token)
+	if token := c.tokenValue(); token != "" {
+		req.Header.Set("Authorization", "Bearer "+token)
 	}
 	resp, err := c.hc.Do(req)
 	if err != nil {
@@ -117,6 +133,21 @@ func (c *apiClient) do(ctx context.Context, method, path string, body, out any)
 	return nil
 }
 
+// do retries authenticated API calls once after a fresh cryptographic login.
+// This matters for unattended hosted workers: a long network outage or an
+// expired JWT must not leave a healthy process permanently stuck on 401.
+func (c *apiClient) do(ctx context.Context, method, path string, body, out any) error {
+	err := c.doRaw(ctx, method, path, body, out)
+	var ae *apiError
+	if strings.HasPrefix(path, "/api/auth/") || !errors.As(err, &ae) || ae.Status != http.StatusUnauthorized {
+		return err
+	}
+	if loginErr := c.login(ctx); loginErr != nil {
+		return fmt.Errorf("session refresh after HTTP 401: %w", loginErr)
+	}
+	return c.doRaw(ctx, method, path, body, out)
+}
+
 func leadingZeroBitsClient(b []byte) int {
 	n := 0
 	for _, x := range b {
@@ -145,12 +176,15 @@ func solveProofClient(challenge, cid string, bits int) string {
 	}
 }
 func (c *apiClient) login(ctx context.Context) error {
+	c.authMu.Lock()
+	defer c.authMu.Unlock()
+
 	var ch struct {
 		ClientID        string `json:"client_id"`
 		Challenge       string `json:"challenge"`
 		ProofOfWorkBits int    `json:"proof_of_work_bits"`
 	}
-	if err := c.do(ctx, http.MethodPost, "/api/auth/challenge", map[string]any{"public_jwk": c.id.PublicJWK}, &ch); err != nil {
+	if err := c.doRaw(ctx, http.MethodPost, "/api/auth/challenge", map[string]any{"public_jwk": c.id.PublicJWK}, &ch); err != nil {
 		return fmt.Errorf("challenge: %w", err)
 	}
 	sig, err := signRaw(c.key, "login|"+ch.Challenge+"|"+ch.ClientID)
@@ -162,7 +196,7 @@ func (c *apiClient) login(ctx context.Context) error {
 		ClientID string `json:"client_id"`
 	}
 	pow := solveProofClient(ch.Challenge, ch.ClientID, ch.ProofOfWorkBits)
-	if err := c.do(ctx, http.MethodPost, "/api/auth/login", map[string]any{
+	if err := c.doRaw(ctx, http.MethodPost, "/api/auth/login", map[string]any{
 		"public_jwk":            c.id.PublicJWK,
 		"challenge":             ch.Challenge,
 		"signature":             sig,
@@ -170,7 +204,12 @@ func (c *apiClient) login(ctx context.Context) error {
 	}, &lg); err != nil {
 		return fmt.Errorf("login: %w", err)
 	}
+	if old := c.clientID(); old != "" && old != lg.ClientID {
+		return fmt.Errorf("server returned a different client identity after re-login: %s != %s", lg.ClientID, old)
+	}
+	c.mu.Lock()
 	c.token, c.cid = lg.Token, lg.ClientID
+	c.mu.Unlock()
 	return nil
 }
 
@@ -318,7 +357,7 @@ func (c *apiClient) ownedArtifacts(ctx context.Context, limit int) ([]ownedArtif
 }
 
 func (c *apiClient) guess(ctx context.Context, t taskDTO, seq int64, beaconPath string) (bool, error) {
-	guess := expectedGuess(t.ID, t.PublicSeed, c.cid, seq, t.RangeBits)
+	guess := expectedGuess(t.ID, t.PublicSeed, c.clientID(), seq, t.RangeBits)
 	msg := fmt.Sprintf("guess|%s|%d|%s", t.ID, seq, guess)
 	if t.BeaconHuntEnabled == 1 && t.GuessLotteryMaxAccepted > 0 {
 		beaconPath = strings.ToUpper(strings.TrimSpace(beaconPath))
@@ -333,7 +372,7 @@ func (c *apiClient) guess(ctx context.Context, t taskDTO, seq int64, beaconPath
 	return correct, err
 }
 
-func (c *apiClient) dialWS(ctx context.Context, maxNodes int) (*websocket.Conn, error) {
+func (c *apiClient) dialWSOnce(ctx context.Context, maxNodes int) (*websocket.Conn, error) {
 	u, err := url.Parse(c.base)
 	if err != nil {
 		return nil, err
@@ -346,16 +385,30 @@ func (c *apiClient) dialWS(ctx context.Context, maxNodes int) (*websocket.Conn,
 	q.Set("max_nodes", strconv.Itoa(maxNodes))
 	wu := scheme + "://" + u.Host + "/api/ws?" + q.Encode()
 	h := http.Header{}
-	h.Set("Authorization", "Bearer "+c.token)
+	if token := c.tokenValue(); token != "" {
+		h.Set("Authorization", "Bearer "+token)
+	}
 	conn, resp, err := websocket.DefaultDialer.DialContext(ctx, wu, h)
 	if err != nil && resp != nil {
 		b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
 		_ = resp.Body.Close()
-		return nil, fmt.Errorf("websocket %s: %s", resp.Status, strings.TrimSpace(string(b)))
+		return nil, &apiError{Status: resp.StatusCode, Body: string(b)}
 	}
 	return conn, err
 }
 
+func (c *apiClient) dialWS(ctx context.Context, maxNodes int) (*websocket.Conn, error) {
+	conn, err := c.dialWSOnce(ctx, maxNodes)
+	var ae *apiError
+	if !errors.As(err, &ae) || ae.Status != http.StatusUnauthorized {
+		return conn, err
+	}
+	if loginErr := c.login(ctx); loginErr != nil {
+		return nil, fmt.Errorf("websocket session refresh: %w", loginErr)
+	}
+	return c.dialWSOnce(ctx, maxNodes)
+}
+
 func (c *apiClient) downloadPreview(ctx context.Context, taskID, dest string) error {
 	path := "/api/public/artifacts/" + url.PathEscape(taskID) + "/preview"
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+path, nil)
@@ -391,8 +444,8 @@ func (c *apiClient) downloadOwnedArtifact(ctx context.Context, taskID, dest stri
 	if err != nil {
 		return err
 	}
-	if c.token != "" {
-		req.Header.Set("Authorization", "Bearer "+c.token)
+	if token := c.tokenValue(); token != "" {
+		req.Header.Set("Authorization", "Bearer "+token)
 	}
 	resp, err := c.hc.Do(req)
 	if err != nil {
@@ -439,7 +492,7 @@ func (c *apiClient) registerHostedWorker(ctx context.Context) error {
 	if registerURL == "" || workerID == "" || token == "" {
 		return errors.New("hosted worker registration requires NEURALHUNT_WORKER_REGISTER_URL, NEURALHUNT_WORKER_ID and NEURALHUNT_WORKER_REGISTER_TOKEN")
 	}
-	body, _ := json.Marshal(map[string]string{"worker_id": workerID, "client_id": c.cid})
+	body, _ := json.Marshal(map[string]string{"worker_id": workerID, "client_id": c.clientID()})
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, registerURL, bytes.NewReader(body))
 	if err != nil {
 		return err
@@ -492,6 +545,14 @@ func (c *apiClient) hostedWorkerLeaseLoop(ctx context.Context, cancel context.Ca
 						failures = 0
 						continue
 					}
+					// Explicit authentication/revocation responses are authoritative,
+					// not transient connectivity failures. Stop immediately; Docker's
+					// restart policy cannot bypass registration while the DB lease is
+					// revoked.
+					if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusConflict {
+						cancel()
+						return
+					}
 				}
 			}
 			failures++