Warp/server/internal/storage/pocketbase/client.go

package pocketbase

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"strings"
	"sync"
)

// client speaks PocketBase's REST API as a superuser. Everything the driver
// does goes through it: the SQL endpoint for reads, the batch endpoint for
// writes, and the record and file endpoints for media.
type client struct {
	base     *url.URL
	email    string
	password string
	http     *http.Client

	mu    sync.Mutex
	token string
}

// newClient parses the server address. Credentials are checked on first use,
// not here: a PocketBase that is still starting is not a configuration
// error.
func newClient(rawURL, email, password string) (*client, error) {
	if rawURL == "" {
		return nil, errors.New("pocketbase: connection_string must be the server's URL, e.g. http://pocketbase:8070")
	}
	base, err := url.Parse(strings.TrimSuffix(rawURL, "/"))
	if err != nil || (base.Scheme != "http" && base.Scheme != "https") || base.Host == "" {
		return nil, fmt.Errorf("pocketbase: connection_string %q is not an http(s) URL", rawURL)
	}
	if email == "" || password == "" {
		return nil, errors.New("pocketbase: a superuser email and password are required")
	}
	return &client{base: base, email: email, password: password, http: &http.Client{}}, nil
}

// apiError is a non-2xx answer from PocketBase.
type apiError struct {
	Status  int
	Message string
	Data    json.RawMessage
}

func (e *apiError) Error() string {
	msg := fmt.Sprintf("pocketbase: HTTP %d: %s", e.Status, e.Message)
	if len(e.Data) > 0 && string(e.Data) != "{}" && string(e.Data) != "null" {
		msg += " " + string(e.Data)
	}
	return msg
}

// isNotFound reports a 404 from PocketBase.
func isNotFound(err error) bool {
	var apiErr *apiError
	return errors.As(err, &apiErr) && apiErr.Status == http.StatusNotFound
}

// url resolves an API path against the server address.
func (c *client) url(path string) string { return c.base.String() + path }

// authToken returns the current superuser token, signing in if there is
// none yet.
func (c *client) authToken(ctx context.Context) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.token != "" {
		return c.token, nil
	}
	body, err := json.Marshal(map[string]string{"identity": c.email, "password": c.password})
	if err != nil {
		return "", err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.url("/api/collections/_superusers/auth-with-password"), bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	var out struct {
		Token string `json:"token"`
	}
	if err := c.send(req, &out); err != nil {
		return "", fmt.Errorf("pocketbase: sign in as superuser %q: %w", c.email, err)
	}
	if out.Token == "" {
		return "", errors.New("pocketbase: sign in returned no token")
	}
	c.token = out.Token
	return c.token, nil
}

// forgetToken drops a token PocketBase refused, so the next call signs in
// again. A superuser token outlives most deployments' restarts, but not all.
func (c *client) forgetToken(token string) {
	c.mu.Lock()
	if c.token == token {
		c.token = ""
	}
	c.mu.Unlock()
}

// do sends one authenticated request, signing in again once if the token has
// expired. body is JSON-encoded unless it is an io.Reader, which is sent as
// is with contentType.
func (c *client) do(ctx context.Context, method, path string, body any, contentType string, out any) error {
	var payload []byte
	if body != nil {
		if r, ok := body.(io.Reader); ok {
			// A streamed body cannot be replayed, so it is sent once and a
			// stale token is the caller's error to retry.
			return c.doOnce(ctx, method, path, r, contentType, out)
		}
		var err error
		if payload, err = json.Marshal(body); err != nil {
			return err
		}
		contentType = "application/json"
	}
	for attempt := 0; ; attempt++ {
		var r io.Reader
		if payload != nil {
			r = bytes.NewReader(payload)
		}
		err := c.doOnce(ctx, method, path, r, contentType, out)
		var apiErr *apiError
		if attempt == 0 && errors.As(err, &apiErr) && apiErr.Status == http.StatusUnauthorized {
			continue
		}
		return err
	}
}

func (c *client) doOnce(ctx context.Context, method, path string, body io.Reader, contentType string, out any) error {
	token, err := c.authToken(ctx)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, method, c.url(path), body)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", token)
	if contentType != "" {
		req.Header.Set("Content-Type", contentType)
	}
	err = c.send(req, out)
	var apiErr *apiError
	if errors.As(err, &apiErr) && apiErr.Status == http.StatusUnauthorized {
		c.forgetToken(token)
	}
	return err
}

// send runs a request and decodes a JSON answer into out, if out is set.
func (c *client) send(req *http.Request, out any) error {
	resp, err := c.http.Do(req)
	if err != nil {
		return fmt.Errorf("pocketbase: %s %s: %w", req.Method, req.URL.Path, err)
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode >= 300 {
		return readAPIError(resp)
	}
	if out == nil || resp.StatusCode == http.StatusNoContent {
		_, _ = io.Copy(io.Discard, resp.Body)
		return nil
	}
	if err := json.NewDecoder(resp.Body).Decode(out); err != nil {
		return fmt.Errorf("pocketbase: decode the answer to %s %s: %w", req.Method, req.URL.Path, err)
	}
	return nil
}

func readAPIError(resp *http.Response) error {
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	apiErr := &apiError{Status: resp.StatusCode}
	var body struct {
		Message string          `json:"message"`
		Data    json.RawMessage `json:"data"`
	}
	if json.Unmarshal(raw, &body) == nil && body.Message != "" {
		apiErr.Message = body.Message
		apiErr.Data = body.Data
	} else {
		apiErr.Message = strings.TrimSpace(string(raw))
	}
	return apiErr
}

// health checks the server answers at all; it needs no token.
func (c *client) health(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.url("/api/health"), nil)
	if err != nil {
		return err
	}
	return c.send(req, nil)
}

// sqlResult is what PocketBase's SQL endpoint answers with. Every value comes
// back as text, or null.
type sqlResult struct {
	Columns []struct {
		Name string `json:"name"`
	} `json:"columns"`
	Rows [][]*string `json:"rows"`
}

// query runs one SELECT through the SQL endpoint.
func (c *client) query(ctx context.Context, query string) (*sqlResult, error) {
	var out sqlResult
	if err := c.do(ctx, http.MethodPost, "/api/sql", map[string]string{"query": query}, "", &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// batchRequest is one record operation inside a batch.
type batchRequest struct {
	Method string         `json:"method"`
	URL    string         `json:"url"`
	Body   map[string]any `json:"body,omitempty"`
}

// batch runs record operations in one PocketBase transaction: all of them
// happen, or none do.
func (c *client) batch(ctx context.Context, reqs []batchRequest) error {
	if len(reqs) == 0 {
		return nil
	}
	return c.do(ctx, http.MethodPost, "/api/batch", map[string]any{"requests": reqs}, "", nil)
}

// patchSettings changes part of PocketBase's application settings.
func (c *client) patchSettings(ctx context.Context, settings map[string]any) error {
	return c.do(ctx, http.MethodPatch, "/api/settings", settings, "", nil)
}

// importCollections creates or updates collections, leaving any not named
// alone.
func (c *client) importCollections(ctx context.Context, collections []map[string]any) error {
	return c.do(ctx, http.MethodPut, "/api/collections/import",
		map[string]any{"collections": collections, "deleteMissing": false}, "", nil)
}

// getRecord reads one record as JSON fields.
func (c *client) getRecord(ctx context.Context, collection, id string) (map[string]any, error) {
	var out map[string]any
	if err := c.do(ctx, http.MethodGet,
		"/api/collections/"+collection+"/records/"+url.PathEscape(id), nil, "", &out); err != nil {
		return nil, err
	}
	return out, nil
}

// deleteRecord removes one record.
func (c *client) deleteRecord(ctx context.Context, collection, id string) error {
	return c.do(ctx, http.MethodDelete,
		"/api/collections/"+collection+"/records/"+url.PathEscape(id), nil, "", nil)
}

// createRecordWithFile creates a record with one file field, streaming the
// file as multipart so that it is never held in memory whole.
func (c *client) createRecordWithFile(ctx context.Context, collection string, fields map[string]string,
	fileField, fileName string, file io.Reader) error {
	pr, pw := io.Pipe()
	mw := multipart.NewWriter(pw)
	go func() {
		err := func() error {
			for k, v := range fields {
				if err := mw.WriteField(k, v); err != nil {
					return err
				}
			}
			part, err := mw.CreateFormFile(fileField, fileName)
			if err != nil {
				return err
			}
			if _, err := io.Copy(part, file); err != nil {
				return err
			}
			return mw.Close()
		}()
		_ = pw.CloseWithError(err)
	}()
	err := c.do(ctx, http.MethodPost, "/api/collections/"+collection+"/records", pr, mw.FormDataContentType(), nil)
	_ = pr.Close()
	return err
}

// fileToken mints the short-lived token protected files are read with.
func (c *client) fileToken(ctx context.Context) (string, error) {
	var out struct {
		Token string `json:"token"`
	}
	if err := c.do(ctx, http.MethodPost, "/api/files/token", nil, "", &out); err != nil {
		return "", err
	}
	return out.Token, nil
}

// openFile starts a download of a stored file from offset on.
func (c *client) openFile(ctx context.Context, collection, id, name, token string, offset int64) (*http.Response, error) {
	path := "/api/files/" + collection + "/" + url.PathEscape(id) + "/" + url.PathEscape(name) +
		"?token=" + url.QueryEscape(token)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.url(path), nil)
	if err != nil {
		return nil, err
	}
	if offset > 0 {
		req.Header.Set("Range", fmt.Sprintf("bytes=%d-", offset))
	}
	resp, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("pocketbase: download %s: %w", name, err)
	}
	if resp.StatusCode >= 300 {
		defer func() { _ = resp.Body.Close() }()
		return nil, readAPIError(resp)
	}
	return resp, nil
}