mesh-tools/node-tools/internal/bus/chunked.go

package bus

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

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

	"git.novox.be/novox/mesh-sdk/go/buscalls"
)

// FileField is the member of a chunked verb's answer that carries the file received: its whole bytes, base64 (design 47 G
// A-249).
const FileField = "file"

// AskAsChunked is AskAs of a seat's verb whose definition says `chunks` (design 47 G §681, A-237, A-249): it chooses the call
// token (the body's `call`), hears the chunk subjects of that call before it asks, receives the chunks the verb's answer
// counts with buscalls.SeatChunks, each acknowledged before the next, and answers the verb's own answer with the whole
// file added as FileField. The asking bundle receives no chunk itself.
func (c *Conn) AskAsChunked(module, key string, body any) (Answered, error) {
	// The verb may be asked of the machine that ran the ask (`@<node>`, A-254): the chunk subjects name the seat, the verb and
	// the call alone.
	name, _, _ := strings.Cut(key, "@")
	seat, verb, ok := strings.Cut(strings.TrimPrefix(name, "seat:"), ".")
	if !strings.HasPrefix(key, "seat:") || !ok {
		return Answered{}, fmt.Errorf("%q is no seat:<seat>.<verb>, which a chunked answer needs", key)
	}
	args := map[string]any{}
	if raw, err := json.Marshal(body); err != nil {
		return Answered{}, err
	} else if string(raw) != "null" {
		if err := json.Unmarshal(raw, &args); err != nil {
			return Answered{}, fmt.Errorf("the arguments of %s are not an object: %w", key, err)
		}
	}
	var bits [16]byte
	if _, err := rand.Read(bits[:]); err != nil {
		return Answered{}, err
	}
	send := buscalls.SeatChunks{Seat: seat, Verb: verb, Call: hex.EncodeToString(bits[:])}
	args["call"] = send.Call
	chunks := make(chan *nats.Msg, 4)
	sub, err := c.nc.ChanSubscribe(send.Chunks(), chunks)
	if err != nil {
		return Answered{}, fmt.Errorf("the chunks of %s cannot be heard: %w", key, err)
	}
	defer func() { _ = sub.Unsubscribe() }()
	answered, err := c.AskAs(module, key, args)
	if err != nil {
		return answered, err
	}
	fields := map[string]json.RawMessage{}
	if err := json.Unmarshal(answered.Result, &fields); err != nil {
		return answered, fmt.Errorf("%s answered what is no object: %w", key, err)
	}
	var count int64
	if raw, has := fields["chunks"]; has {
		if err := json.Unmarshal(raw, &count); err != nil {
			return answered, fmt.Errorf("%s answered chunks that is no count: %w", key, err)
		}
	}
	var file bytes.Buffer
	if count > 0 {
		next := func() (string, []byte, error) {
			select {
			case m := <-chunks:
				return m.Subject, m.Data, nil
			case <-time.After(buscalls.SeatAckWithin):
				return "", nil, buscalls.ErrSeatChunkMissed
			}
		}
		ack := func(subject string) error {
			_, err := c.nc.Request(subject, []byte("{}"), buscalls.SeatAckWithin)
			return err
		}
		if _, err := send.Receive(count, next, ack, &file); err != nil {
			return answered, fmt.Errorf("the chunks of %s: %w", key, err)
		}
	}
	if _, taken := fields[FileField]; taken {
		return answered, errors.New(key + " answers a member named " + FileField + " itself")
	}
	enc, _ := json.Marshal(base64.StdEncoding.EncodeToString(file.Bytes()))
	fields[FileField] = enc
	out, err := json.Marshal(fields)
	if err != nil {
		return answered, err
	}
	answered.Result = out
	return answered, nil
}

// AnswerChunked is a seat verb's answer as its definition says `chunks` (design 47 G A-257, §681): a bundle answers the whole file
// as the member FileField (base64), and the runner that serves the verb sends it in chunks with buscalls.SeatChunks on the call
// token the body names, each acknowledged by the caller before the next, and answers the verb's own members with `chunks`, the
// count, in place of the file. The bundle sends no chunk itself. An answer with no file is answered as it is. The send runs after
// the answer is returned: the caller hears the chunks before it asks (AskAsChunked) and acknowledges each once it has the answer.
func (c *Conn) AnswerChunked(seat, verb string, body json.RawMessage, answer any) (any, error) {
	raw, err := json.Marshal(answer)
	if err != nil {
		return nil, err
	}
	fields := map[string]json.RawMessage{}
	if err := json.Unmarshal(raw, &fields); err != nil {
		return answer, nil
	}
	enc, has := fields[FileField]
	if !has {
		return answer, nil
	}
	var text string
	if err := json.Unmarshal(enc, &text); err != nil {
		return nil, fmt.Errorf("%s.%s answers a %s that is no string: %w", seat, verb, FileField, err)
	}
	file, err := base64.StdEncoding.DecodeString(text)
	if err != nil {
		return nil, fmt.Errorf("%s.%s answers a %s in no base64: %w", seat, verb, FileField, err)
	}
	var asked struct {
		Call string `json:"call"`
	}
	if err := json.Unmarshal(body, &asked); err != nil || asked.Call == "" {
		return nil, fmt.Errorf("%s.%s is asked in chunks with no call token", seat, verb)
	}
	delete(fields, FileField)
	count := int64((len(file) + buscalls.SeatChunkBytes - 1) / buscalls.SeatChunkBytes)
	fields["chunks"], _ = json.Marshal(count)
	send := buscalls.SeatChunks{Seat: seat, Verb: verb, Call: asked.Call}
	go func() { _ = send.Send(c.nc, file, buscalls.SeatAckWithin, c.Logf) }()
	return json.RawMessage(mustMarshal(fields)), nil
}

func mustMarshal(v any) []byte {
	b, _ := json.Marshal(v)
	return b
}