ITWorx-Pulse-Public/internal/unraid/client.go

package unraid

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const ProtocolVersion = "v1"

var allowedQueries = map[string]string{
	"server": "query server { server { name status } }",
	// The fields below are the documented monitoring-only selection sets. Keeping the
	// complete documents here, rather than accepting caller supplied GraphQL, means a
	// compromised caller cannot turn this client into a mutation or schema-enumeration
	// proxy. See docs.unraid.net/API/how-to-use-the-api/.
	"array":      "query array { array { state capacity { disks { free used total } } parities { id name type size status } disks { id name type size status fsSize fsUsed fsFree fsType temp isSpinning } caches { id name type size status fsSize fsUsed fsFree fsType temp isSpinning } } }",
	"containers": "query containers { docker { containers { id names state status autoStart } } }",
	"shares":     "query shares { shares { id name used size include exclude cache allocator } }",
}

type Client struct {
	endpoint    *url.URL
	token       string
	http        *http.Client
	timeout     time.Duration
	maxResponse int64
}

func New(endpoint, token string, httpClient *http.Client) (*Client, error) {
	return newClient(endpoint, token, httpClient)
}

// NewWithCAPEM builds the production client with one explicitly trusted private
// CA/leaf certificate. It never disables hostname or chain verification.
func NewWithCAPEM(endpoint, token string, caPEM []byte) (*Client, error) {
	if len(caPEM) == 0 || len(caPEM) > 1<<20 {
		return nil, errors.New("unraid CA certificate must be between 1 byte and 1 MiB")
	}
	roots, err := x509.SystemCertPool()
	if err != nil || roots == nil {
		roots = x509.NewCertPool()
	}
	if !roots.AppendCertsFromPEM(caPEM) {
		return nil, errors.New("unraid CA certificate is invalid")
	}
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: roots}
	return newClient(endpoint, token, &http.Client{Transport: transport})
}

func newClient(endpoint, token string, httpClient *http.Client) (*Client, error) {
	u, err := url.Parse(strings.TrimSpace(endpoint))
	if err != nil || u.Scheme != "https" || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" {
		return nil, errors.New("unraid endpoint must be an HTTPS URL without credentials or query")
	}
	if strings.TrimSpace(token) == "" {
		return nil, errors.New("unraid token is required")
	}
	if httpClient == nil {
		httpClient = &http.Client{}
	}
	return &Client{endpoint: u, token: token, http: httpClient, timeout: 10 * time.Second, maxResponse: 2 << 20}, nil
}

type Response struct {
	Data   json.RawMessage `json:"data"`
	Errors []struct {
		Message string `json:"message"`
	} `json:"errors"`
}

func (c *Client) Query(ctx context.Context, operation string) (json.RawMessage, error) {
	query, ok := allowedQueries[operation]
	if !ok {
		return nil, fmt.Errorf("unsupported read-only Unraid operation %q", operation)
	}
	body, _ := json.Marshal(map[string]string{"query": query, "operationName": operation})
	requestContext, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()
	req, err := http.NewRequestWithContext(requestContext, http.MethodPost, c.endpoint.String(), bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	// Unraid API keys are sent through x-api-key. Do not use a browser/session bearer
	// token here: the agent/API holds a narrowly scoped API key and never forwards it.
	req.Header.Set("x-api-key", c.token)
	req.Header.Set("Content-Type", "application/json")
	response, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("unraid request: %w", err)
	}
	defer response.Body.Close()
	if response.ContentLength > c.maxResponse {
		return nil, errors.New("unraid response exceeds bounds")
	}
	limited, err := io.ReadAll(io.LimitReader(response.Body, c.maxResponse+1))
	if err != nil || int64(len(limited)) > c.maxResponse {
		return nil, errors.New("unraid response exceeds bounds")
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return nil, fmt.Errorf("unraid HTTP status %d", response.StatusCode)
	}
	var result Response
	if err := json.Unmarshal(limited, &result); err != nil {
		return nil, errors.New("invalid Unraid response")
	}
	if len(result.Errors) > 0 {
		return nil, errors.New("unraid GraphQL query failed")
	}
	return result.Data, nil
}