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 }