worker/internal/distworker/client.go

package distworker

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"rocketgit.ru/rsmon/worker/internal/wire"
)

// HTTPStatusError reports an HTTP response returned while dialing the websocket.
type HTTPStatusError struct {
	StatusCode int
	Status     string
	Body       string
}

func (e *HTTPStatusError) Error() string {
	if e.Body == "" {
		return fmt.Sprintf("websocket dial failed: %s", e.Status)
	}
	return fmt.Sprintf("websocket dial failed: %s: %s", e.Status, e.Body)
}

// Client is the HTTP client for communicating with the control plane
type Client struct {
	endpoint   string
	authToken  string
	httpClient *http.Client
}

// NewClient creates a new API client
func NewClient(endpoint, authToken string) *Client {
	return &Client{
		endpoint:  endpoint,
		authToken: authToken,
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
		},
	}
}

// postJSON is a helper for sending JSON POST requests
func (c *Client) postJSON(path string, payload interface{}) (*http.Response, error) {
	return c.postJSONContext(context.Background(), path, payload)
}

func (c *Client) postJSONContext(ctx context.Context, path string, payload interface{}) (*http.Response, error) {
	body, err := json.Marshal(payload)
	if err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", c.endpoint+path, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}

	httpReq.Header.Set("Content-Type", "application/json")
	if c.authToken != "" {
		httpReq.Header.Set("Authorization", "Bearer "+c.authToken)
	}

	return c.httpClient.Do(httpReq)
}

// checkStatusCode checks if the response status is OK, returns error otherwise
func checkStatusCode(resp *http.Response, errorMsg string) error {
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("%s failed: %s: %s", errorMsg, resp.Status, string(body))
	}
	return nil
}

// Heartbeat sends a heartbeat to the control plane
func (c *Client) Heartbeat(req wire.HeartbeatRequest) error {
	resp, err := c.postJSON("/api/internal/workers/heartbeat", req)
	if err != nil {
		return err
	}
	defer resp.Body.Close() //nolint:errcheck

	return checkStatusCode(resp, "heartbeat")
}

// RotateToken asks the main app's internal API to mint a new
// bearer token for this worker. The current bearer is used for
// authentication; the response carries the freshly issued token.
//
// Returns the new token string. The main app invalidates the old token
// immediately.
func (c *Client) RotateToken(ctx context.Context) (string, error) {
	resp, err := c.postJSONContext(ctx, "/api/internal/workers/rotate-token", struct{}{})
	if err != nil {
		return "", err
	}
	defer resp.Body.Close() //nolint:errcheck
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("rotate-token failed: %s: %s", resp.Status, string(body))
	}
	var out struct {
		AuthToken string `json:"auth_token"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return "", fmt.Errorf("decode rotate-token response: %w", err)
	}
	if out.AuthToken == "" {
		return "", fmt.Errorf("rotate-token response empty")
	}
	return out.AuthToken, nil
}

func (c *Client) Bootstrap(ctx context.Context, workerID, token string) (wire.BootstrapResponse, error) {
	var out wire.BootstrapResponse
	resp, err := c.postJSONContext(ctx, "/api/internal/workers/bootstrap", wire.BootstrapRequest{WorkerID: workerID, BootstrapToken: token})
	if err != nil {
		return out, err
	}
	defer resp.Body.Close() //nolint:errcheck
	if resp.StatusCode != http.StatusOK {
		return out, checkStatusCode(resp, "bootstrap")
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return out, err
	}
	if out.AuthToken == "" || out.WorkerID != workerID || out.ConfigVerificationKey == "" {
		return out, fmt.Errorf("invalid bootstrap response")
	}
	return out, nil
}

// GetJobs fetches available check jobs from the control plane
func (c *Client) GetJobs() (*wire.JobsResponse, error) {
	httpReq, err := http.NewRequest("GET", c.endpoint+"/api/internal/workers/jobs", http.NoBody)
	if err != nil {
		return nil, err
	}

	httpReq.Header.Set("Authorization", "Bearer "+c.authToken)

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close() //nolint:errcheck

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("get jobs failed: %s: %s", resp.Status, string(body))
	}

	var result wire.JobsResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, err
	}

	return &result, nil
}

// ReportResults sends check results to the control plane
func (c *Client) ReportResults(req wire.ResultsRequest) error {
	resp, err := c.postJSON("/api/internal/workers/results", req)
	if err != nil {
		return err
	}
	defer resp.Body.Close() //nolint:errcheck

	return checkStatusCode(resp, "report results")
}

// WorkerSocket connects to the websocket task channel.
func (c *Client) WorkerSocket() (*websocket.Conn, error) {
	return c.WorkerSocketContext(context.Background())
}

// WorkerSocketContext connects to the websocket task channel, cancelling the
// dial when the caller's context ends.
func (c *Client) WorkerSocketContext(ctx context.Context) (*websocket.Conn, error) {
	endpoint := strings.TrimRight(c.endpoint, "/")
	wsURL := endpoint
	if !strings.HasSuffix(wsURL, "/worker") && !strings.HasSuffix(wsURL, "/api/worker") {
		wsURL += "/worker"
	}
	if strings.HasPrefix(wsURL, "https://") {
		wsURL = "wss://" + strings.TrimPrefix(wsURL, "https://")
	} else if strings.HasPrefix(wsURL, "http://") {
		wsURL = "ws://" + strings.TrimPrefix(wsURL, "http://")
	}

	u, err := url.Parse(wsURL)
	if err != nil {
		return nil, err
	}
	q := u.Query()
	q.Set("token", c.authToken)
	u.RawQuery = q.Encode()

	dialer := *websocket.DefaultDialer
	var (
		connMu   sync.Mutex
		dialConn net.Conn
	)
	stopClose := context.AfterFunc(ctx, func() {
		connMu.Lock()
		if dialConn != nil {
			_ = dialConn.Close()
		}
		connMu.Unlock()
	})
	defer stopClose()
	dialer.NetDialContext = func(dialCtx context.Context, network, address string) (net.Conn, error) {
		conn, err := (&net.Dialer{}).DialContext(dialCtx, network, address)
		if err != nil {
			return nil, err
		}
		connMu.Lock()
		dialConn = conn
		if ctx.Err() != nil {
			_ = conn.Close()
		}
		connMu.Unlock()
		return conn, nil
	}
	conn, resp, err := dialer.DialContext(ctx, u.String(), nil)
	if err != nil && resp != nil {
		defer resp.Body.Close() //nolint:errcheck
		body, _ := io.ReadAll(resp.Body)
		return nil, &HTTPStatusError{StatusCode: resp.StatusCode, Status: resp.Status, Body: string(body)}
	}
	return conn, err
}