mesh-catalog/modules/nats/cmd/nats-tools/snapshots.go

package main

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"

	"git.novox.be/novox/mesh-contracts/go/seats/meshbroker"
	stdio "git.novox.be/novox/mesh-sdk/go"
	"git.novox.be/novox/mesh-sdk/go/process"

	"nats-tools/internal/keeper"
)

// The seat mesh-broker's snapshot verbs and nats_snapshots (design 47 G §681, ADR 0675; A-44, A-107). The keeper's process answers the
// four verbs itself over the bus (process.Conn.AnswerSeat), with the caller the bus proved: snapshot-get answers the fields and
// sends the file in chunks, published by the process and acknowledged by the caller on a subscription of the process's own;
// snapshot-restore, snapshot-restore-cancel and snapshot-remove are acted on by the keeper at once, with no hand-over through the
// state requests. The tools bundle keeps nats_snapshots, which lists them.

// ChunkSize is at most what one chunk carries: 512 KiB.
const ChunkSize = 512 << 10

// AckWithin is how long the next chunk waits for its acknowledgement.
const AckWithin = 30 * time.Second

// callToken is the shape of the subject token a caller chooses for a send.
var callToken = regexp.MustCompile(`^[A-Za-z0-9_-]{16,64}$`)

// ChunkSubject and AckSubject are the subjects of the nth chunk of a send and of its acknowledgement.
func ChunkSubject(call string, n int64) string {
	return "mesh.seat." + meshbroker.Seat + ".chunk." + meshbroker.VerbSnapshotGet + "." + call + "." + strconv.FormatInt(n, 10)
}

// AckSubject is the subject the caller acknowledges the nth chunk on.
func AckSubject(call string, n int64) string {
	return "mesh.seat." + meshbroker.Seat + ".ack." + meshbroker.VerbSnapshotGet + "." + call + "." + strconv.FormatInt(n, 10)
}

// AckPattern is every acknowledgement of snapshot-get.
const AckPattern = "mesh.seat.mesh-broker.ack.snapshot-get.*.*"

// Chunk is the body of one chunk: its number, how many there are, and its bytes.
type Chunk struct {
	N    int64  `json:"n"`
	Of   int64  `json:"of"`
	Data []byte `json:"data"`
}

// snapshotWorld is what the keeper's snapshot verbs reach: the file, the process's own publish and subscription, and the bus's users.
type snapshotWorld struct {
	// Open opens the snapshot file the keeper wrote.
	Open func(path string) (io.ReadCloser, error)
	// Publish says one chunk on its subject, the process's own publish; Hear subscribes to the acknowledgements, once.
	Publish func(subject string, body any, id string) error
	Hear    func(ack func(subject string)) error
	// Granted says whether the caller may hear the chunks and acknowledge them (the bus's users, nats_user_can).
	Granted func(caller, call string) bool
	Logf    func(string, ...any)
	// Within is how long an acknowledgement is waited for.
	Within time.Duration

	once    sync.Once
	onceErr error
	mu      sync.Mutex
	waiting map[string]chan struct{}
}

// grantedBy is whether the bus's users let caller hear a send's chunks and acknowledge them.
func grantedBy(caller, call string) bool {
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	conf, err := ReadUsers(ctx)
	if err != nil {
		return false
	}
	a, err := UserCan(conf, caller, "subscribe", ChunkSubject(call, 1))
	if err != nil || a["allowed"] != true {
		return false
	}
	a, err = UserCan(conf, caller, "publish", AckSubject(call, 1))
	return err == nil && a["allowed"] == true
}

// keeperSnapshots is the world as the keeper's process reaches it: its own JetStream publish of a chunk, as the seat's events are
// said, and its own core subscription for the acknowledgements, answered at once since the caller asks them as requests.
func keeperSnapshots(js jetstream.JetStream, nc *nats.Conn, a process.Account, logf func(string, ...any)) *snapshotWorld {
	return &snapshotWorld{
		Open: func(path string) (io.ReadCloser, error) { return os.Open(path) },
		Publish: func(subject string, body any, id string) error {
			data, err := json.Marshal(body)
			if err != nil {
				return err
			}
			msg := nats.NewMsg(subject)
			msg.Header.Set("x-event-id", a.Module+"."+id)
			msg.Header.Set("x-source", a.Module)
			msg.Header.Set("x-node", a.Node)
			msg.Header.Set("x-time", time.Now().UTC().Format(time.RFC3339Nano))
			msg.Header.Set("content-type", "application/json")
			msg.Data = data
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			_, err = js.PublishMsg(ctx, msg, jetstream.WithMsgID(a.Module+"."+id))
			return err
		},
		Hear: func(ack func(subject string)) error {
			_, err := nc.Subscribe(AckPattern, func(m *nats.Msg) {
				ack(m.Subject)
				if m.Reply != "" {
					_ = m.Respond([]byte("{}"))
				}
			})
			return err
		},
		Granted: grantedBy, Logf: logf, Within: AckWithin,
	}
}

// hearAcks subscribes once to the acknowledgements of every send.
func (w *snapshotWorld) hearAcks() error {
	w.once.Do(func() {
		w.waiting = map[string]chan struct{}{}
		w.onceErr = w.Hear(func(subject string) {
			p := strings.Split(subject, ".")
			if len(p) != 7 {
				return
			}
			w.mu.Lock()
			ch := w.waiting[p[5]+"."+p[6]]
			w.mu.Unlock()
			if ch != nil {
				select {
				case ch <- struct{}{}:
				default:
				}
			}
		})
	})
	return w.onceErr
}

// digestOf reads the file whole: its size and the `sha256:` of its bytes.
func (w *snapshotWorld) digestOf(path string) (int64, string, error) {
	f, err := w.Open(path)
	if err != nil {
		return 0, "", err
	}
	defer f.Close()
	h := sha256.New()
	n, err := io.Copy(h, f)
	if err != nil {
		return 0, "", err
	}
	return n, "sha256:" + hex.EncodeToString(h.Sum(nil)), nil
}

// get answers snapshot-get: the fields, then the chunks sent after the answer when the caller is granted the chunk subjects.
func (w *snapshotWorld) get(held func(id string) (*keeper.Request, error), caller, call, request string) (meshbroker.SnapshotGetAnswer, error) {
	refuse := func(why string) (meshbroker.SnapshotGetAnswer, error) {
		return meshbroker.SnapshotGetAnswer{Error: why}, nil
	}
	if !callToken.MatchString(call) {
		return refuse("call is one token of 16 to 64 letters, digits, - or _: 128 random bits the caller chooses")
	}
	rec, err := held(request)
	if err != nil {
		return meshbroker.SnapshotGetAnswer{}, fmt.Errorf("the snapshot of the request %s cannot be read: %w", request, err)
	}
	if rec == nil || (rec.Path == "" && rec.Snapshot == "") {
		return refuse("no snapshot of the request " + request + " is held")
	}
	if !rec.Recorded() {
		return refuse("snapshot taken before milestone G")
	}
	// A protected snapshot is never handed out: the keeper restores it in place (O-7).
	if rec.Protected {
		return refuse("snapshot protected")
	}
	size, digest, err := w.digestOf(rec.Path)
	if errors.Is(err, os.ErrNotExist) {
		return refuse("snapshot missing")
	}
	if err != nil {
		return meshbroker.SnapshotGetAnswer{}, fmt.Errorf("the snapshot %s cannot be read: %w", rec.Path, err)
	}
	if digest != rec.Digest || size != rec.Size {
		return refuse("snapshot damaged: request " + request)
	}
	at := rec.At
	out := meshbroker.SnapshotGetAnswer{Request: rec.ID, Module: rec.Module, Local: rec.Local, Stream: keeper.BusName(rec.Module, rec.Local),
		Size: size, Digest: digest, At: &at}
	if size == 0 || !w.Granted(caller, call) {
		return out, nil
	}
	if err := w.hearAcks(); err != nil {
		return meshbroker.SnapshotGetAnswer{}, fmt.Errorf("the acknowledgements cannot be heard, so no chunk is sent: %w", err)
	}
	out.Chunks = (size + ChunkSize - 1) / ChunkSize
	go w.send(rec.Path, call, out.Chunks)
	return out, nil
}

// send sends the file in chunks, each acknowledged before the next; a missed acknowledgement ends the send, said in the journal.
func (w *snapshotWorld) send(path, call string, total int64) {
	f, err := w.Open(path)
	if err != nil {
		w.Logf("the snapshot %s cannot be opened to send (%v): the send %s ends", path, err, call)
		return
	}
	defer f.Close()
	buf := make([]byte, ChunkSize)
	for n := int64(1); n <= total; n++ {
		read, err := io.ReadFull(f, buf)
		if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) && !errors.Is(err, io.EOF) {
			w.Logf("the snapshot %s cannot be read (%v): the send %s ends at chunk %d", path, err, call, n)
			return
		}
		key := call + "." + strconv.FormatInt(n, 10)
		ack := make(chan struct{}, 1)
		w.mu.Lock()
		w.waiting[key] = ack
		w.mu.Unlock()
		err = w.Publish(ChunkSubject(call, n), Chunk{N: n, Of: total, Data: buf[:read]}, key)
		if err == nil {
			select {
			case <-ack:
			case <-time.After(w.Within):
				err = fmt.Errorf("no acknowledgement within %s", w.Within)
			}
		}
		w.mu.Lock()
		delete(w.waiting, key)
		w.mu.Unlock()
		if err != nil {
			w.Logf("the send %s of the snapshot %s ends at chunk %d of %d: %v", call, path, n, total, err)
			return
		}
	}
}

// snapshotArgs are the arguments of the four verbs.
type snapshotArgs struct {
	Request string `json:"request"`
	Call    string `json:"call"`
	Act     string `json:"act"`
	Why     string `json:"why"`
	Verify  bool   `json:"verify"`
	Expired bool   `json:"expired"`
}

// argsOf reads a call's arguments and refuses one that lacks a name in need.
func argsOf(body json.RawMessage, need ...string) (snapshotArgs, error) {
	var a snapshotArgs
	if len(body) > 0 {
		if err := json.Unmarshal(body, &a); err != nil {
			return a, fmt.Errorf("the arguments are not JSON: %w", err)
		}
	}
	a.Why = strings.TrimSpace(a.Why)
	have := map[string]string{"request": a.Request, "call": a.Call, "act": a.Act, "why": a.Why}
	for _, n := range need {
		if have[n] == "" {
			return a, fmt.Errorf("%s is required", n)
		}
	}
	return a, nil
}

// snapshotAnswers are the four verbs of the seat mesh-broker that concern the keeper's snapshots, as the keeper's process answers
// them (process.Conn.AnswerSeat): each with the caller the bus proved, "unproven" for a call that named none.
func snapshotAnswers(k *keeper.Keeper, held func(id string) (*keeper.Request, error), w *snapshotWorld) map[string]process.SeatHandler {
	acting := func(verb, caller, of string, a snapshotArgs) keeper.Request {
		if caller == "" {
			caller = stdio.Unproven
		}
		return k.AnswerSnapshot(keeper.Request{ID: requestID(time.Now().UTC()), Verb: verb, Of: of, Verify: a.Verify, AllExpired: a.Expired, Why: a.Why, Caller: caller})
	}
	return map[string]process.SeatHandler{
		meshbroker.VerbSnapshotGet: func(caller string, body json.RawMessage) (any, error) {
			a, err := argsOf(body, "request", "call")
			if err != nil {
				return nil, err
			}
			if caller == "" {
				caller = stdio.Unproven
			}
			return w.get(held, caller, a.Call, a.Request)
		},
		meshbroker.VerbSnapshotRestore: func(caller string, body json.RawMessage) (any, error) {
			a, err := argsOf(body, "request", "why")
			if err != nil {
				return nil, err
			}
			r := acting(keeper.VerbSnapshotRestore, caller, a.Request, a)
			switch {
			case r.Error != "":
				return meshbroker.SnapshotRestoreAnswer{Error: r.Error}, nil
			case r.Outcome == "busy":
				return meshbroker.SnapshotRestoreAnswer{Busy: r.ActID}, nil
			case r.ActID != "":
				return meshbroker.SnapshotRestoreAnswer{Accepted: r.ActID}, nil
			}
			return meshbroker.SnapshotRestoreAnswer{Accepted: r.ID}, nil
		},
		meshbroker.VerbSnapshotRestoreCancel: func(caller string, body json.RawMessage) (any, error) {
			a, err := argsOf(body, "act", "why")
			if err != nil {
				return nil, err
			}
			r := acting(keeper.VerbSnapshotRestoreCancel, caller, a.Act, a)
			if r.Error != "" {
				return meshbroker.SnapshotRestoreCancelAnswer{Error: r.Error}, nil
			}
			return meshbroker.SnapshotRestoreCancelAnswer{Accepted: r.ActID}, nil
		},
		meshbroker.VerbSnapshotRemove: func(caller string, body json.RawMessage) (any, error) {
			a, err := argsOf(body, "why")
			if err != nil {
				return nil, err
			}
			r := acting(keeper.VerbSnapshotRemove, caller, a.Request, a)
			if r.Error != "" {
				return meshbroker.SnapshotRemoveAnswer{Error: r.Error}, nil
			}
			return meshbroker.SnapshotRemoveAnswer{Accepted: r.ID}, nil
		},
	}
}

// heldRequest reads one record of the keeper's own state requests.
func heldRequest(store keeper.Store) func(id string) (*keeper.Request, error) {
	return func(id string) (*keeper.Request, error) {
		raw, _, found, err := store.Get(keeper.StateRequests, id)
		if err != nil || !found {
			return nil, err
		}
		var r keeper.Request
		return &r, json.Unmarshal(raw, &r)
	}
}

// snapshotListTool is nats_snapshots, which the tools bundle keeps beside the verbs the keeper's process answers.
func snapshotListTool(rs requestStore, keepDays int) []stdio.Tool {
	return []stdio.Tool{
		{Name: "nats_snapshots",
			Description: "The keeper's snapshots of kept data taken before a drop or recreate, from its state requests: each one's request, module, local name, stream, size, digest, time, whether it is protected (restored only in place, never sent), whether it is expired (past snapshot-keep-days, removed at the keeper's next request) and whether it is recorded (a snapshot of before milestone G holds no digest and is sent or removed by hand only).",
			Run: func(map[string]any) (any, error) {
				keys, err := rs.Keys()
				if err != nil {
					return nil, err
				}
				sort.Strings(keys)
				out := []map[string]any{}
				for _, key := range keys {
					if strings.HasPrefix(key, keeper.ActedPrefix) {
						continue
					}
					r, err := rs.Get(key)
					if err != nil || r == nil || (r.Path == "" && r.Snapshot == "") {
						continue
					}
					out = append(out, map[string]any{"request": r.ID, "module": r.Module, "local": r.Local, "stream": keeper.BusName(r.Module, r.Local),
						"size": r.Size, "digest": r.Digest, "at": r.At, "protected": r.Protected, "expired": r.Expired(rs.Now(), keepDays), "recorded": r.Recorded()})
				}
				return map[string]any{"snapshots": out}, nil
			}},
	}
}