mesh-tools/node-tools/internal/runtime/sessions.go

package runtime

import (
	"bufio"
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

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

	"github.com/novox/mesh-tools/node-tools/internal/launch"
)

// A bundle's put and get sessions (design 47 G §518, A-21, A-25, §620): a value over buscalls.ChunkBytes crosses the
// bundle's channel as calls the runner holds the session of, nothing of it on the bus. A put writes each chunk into the
// module's object store as it arrives (through a pipe into the store's own chunking, so no more than one chunk is held) and
// seals the object on the last chunk; a get reads the object out as the bundle asks its chunks, in order.

// SessionBound is how long a session has, counted once for the whole session from its first chunk (§518 "The bound"): the
// call's bound, launch.CallTimeout.
var SessionBound = launch.CallTimeout

// objectSessions is what a session needs of the bus: *bus.Conn.
type objectSessions interface {
	ObjectsPutStream(ctx context.Context, module string, r buscalls.ObjectsNameRequest, description string, in io.Reader) (buscalls.ObjectInfo, error)
	ObjectsOpen(ctx context.Context, module string, r buscalls.ObjectsNameRequest) (buscalls.ObjectInfo, io.ReadCloser, bool, error)
	SentOpen(ctx context.Context, r buscalls.ObjectsSentGetRequest) (buscalls.ObjectInfo, io.ReadCloser, error)
}

// sessionTable holds the sessions of every module, each under its module's name and the bundle's call id, so a bundle
// reaches its own sessions alone.
type sessionTable struct {
	mu   sync.Mutex
	puts map[string]*putSession
	gets map[string]*getSession
	// store replaces the bus where a test has none.
	store objectSessions
}

type putSession struct {
	mu     sync.Mutex
	next   uint32
	pw     *io.PipeWriter
	done   chan struct{} // closed once res is set
	res    putResult
	cancel context.CancelFunc
	timer  *time.Timer
	store  string
}

type putResult struct {
	info buscalls.ObjectInfo
	err  error
}

type getSession struct {
	mu     sync.Mutex
	next   uint32
	rd     *bufio.Reader
	closer io.Closer
	timer  *time.Timer
}

var errSessionBound = errors.New("the session was not completed within the call's bound")

func sessionKey(module, call string) string { return module + " " + call }

// putChunk takes one chunk of a put session: it answers {seq}, or {error}, and the put's own answer for the last chunk.
func (c *consumers) putChunk(module string, params json.RawMessage) (any, error) {
	var r buscalls.ObjectsPutChunkRequest
	if err := decode(buscalls.ObjectsPutChunk, params, &r); err != nil {
		return nil, err
	}
	if r.Call == "" {
		return nil, errors.New("a chunk names the put's call")
	}
	t := &c.sessions
	key := sessionKey(module, r.Call)
	t.mu.Lock()
	s := t.puts[key]
	if r.Seq == 1 && s == nil {
		var err error
		if s, err = c.beginPut(module, key, r); err != nil {
			t.mu.Unlock()
			return buscalls.ObjectsChunkTaken{Error: err.Error()}, nil
		}
	}
	t.mu.Unlock()
	if s == nil {
		return buscalls.ObjectsChunkTaken{Error: fmt.Sprintf("no put %s of %s is open at chunk %d: it was dropped at its bound, or never begun", r.Call, r.Name, r.Seq)}, nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if r.Seq != s.next {
		c.dropPut(key, s, errors.New("a chunk out of order"))
		return buscalls.ObjectsChunkTaken{Error: fmt.Sprintf("putting %s: chunk %d came where chunk %d was next", r.Name, r.Seq, s.next)}, nil
	}
	if len(r.Data) > buscalls.ChunkBytes {
		c.dropPut(key, s, errors.New("a chunk over the limit"))
		return buscalls.ObjectsChunkTaken{Error: fmt.Sprintf("putting %s: chunk %d is %d bytes, over the %d a chunk carries", r.Name, r.Seq, len(r.Data), buscalls.ChunkBytes)}, nil
	}
	if _, err := s.pw.Write(r.Data); err != nil {
		// The store refused the chunk (its cap) or the bound passed: its own reason is the answer.
		<-s.done
		res := s.res
		c.forgetPut(key, s)
		reason := res.err
		if reason == nil || errors.Is(reason, io.ErrClosedPipe) {
			reason = err
		}
		return buscalls.ObjectsChunkTaken{Error: fmt.Sprintf("putting %s into %s, at chunk %d: %v", r.Name, s.store, r.Seq, reason)}, nil
	}
	if !r.Last {
		s.next++
		return buscalls.ObjectsChunkTaken{Seq: r.Seq}, nil
	}
	_ = s.pw.Close()
	<-s.done
	res := s.res
	c.forgetPut(key, s)
	if res.err != nil {
		return buscalls.ObjectsChunkTaken{Error: fmt.Sprintf("putting %s into %s: %v", r.Name, s.store, res.err)}, nil
	}
	return res.info, nil
}

// beginPut opens the put of the first chunk, with t.mu held.
func (c *consumers) beginPut(module, key string, r buscalls.ObjectsPutChunkRequest) (*putSession, error) {
	if r.Name == "" {
		return nil, errors.New("an object is put under a name")
	}
	st := c.objects()
	if st == nil {
		return nil, errors.New("the runner has no bus to put on")
	}
	pr, pw := io.Pipe()
	ctx, cancel := context.WithCancel(context.Background())
	s := &putSession{next: 1, pw: pw, done: make(chan struct{}), cancel: cancel, store: r.Objects}
	s.timer = time.AfterFunc(SessionBound, func() {
		c.sessions.mu.Lock()
		current := c.sessions.puts[key] == s
		c.sessions.mu.Unlock()
		if current {
			_ = pw.CloseWithError(errSessionBound)
			cancel()
		}
	})
	if c.sessions.puts == nil {
		c.sessions.puts = map[string]*putSession{}
	}
	c.sessions.puts[key] = s
	go func() {
		info, err := st.ObjectsPutStream(ctx, module, buscalls.ObjectsNameRequest{Objects: r.Objects, Name: r.Name}, r.Description, pr)
		_ = pr.CloseWithError(err)
		s.res = putResult{info, err}
		close(s.done)
	}()
	return s, nil
}

// dropPut ends a put that will not be sealed: the store purges what it wrote.
func (c *consumers) dropPut(key string, s *putSession, why error) {
	_ = s.pw.CloseWithError(why)
	<-s.done
	c.forgetPut(key, s)
}

func (c *consumers) forgetPut(key string, s *putSession) {
	s.timer.Stop()
	s.cancel()
	c.sessions.mu.Lock()
	delete(c.sessions.puts, key)
	c.sessions.mu.Unlock()
}

// objects is the bus a session reads and writes objects through.
func (c *consumers) objects() objectSessions {
	if c.sessions.store != nil {
		return c.sessions.store
	}
	if c.conn == nil {
		return nil
	}
	return c.conn
}

// openGet registers the get session of an opened object and answers the chunked answer that begins it.
func (c *consumers) openGet(module string, info buscalls.ObjectInfo, rd io.ReadCloser, digest string, sent bool) (any, error) {
	call := newCall()
	g := &getSession{next: 1, rd: bufio.NewReaderSize(rd, buscalls.ChunkBytes), closer: rd}
	key := sessionKey(module, call)
	g.timer = time.AfterFunc(SessionBound, func() { c.dropGet(key, g) })
	c.sessions.mu.Lock()
	if c.sessions.gets == nil {
		c.sessions.gets = map[string]*getSession{}
	}
	c.sessions.gets[key] = g
	c.sessions.mu.Unlock()
	var chunked json.RawMessage = json.RawMessage("true")
	if !sent {
		chunked = json.RawMessage(fmt.Sprint((info.Size + buscalls.ChunkBytes - 1) / buscalls.ChunkBytes))
	}
	return buscalls.ObjectsChunked{Chunked: chunked, Size: info.Size, Digest: digest, Call: call}, nil
}

func (c *consumers) dropGet(key string, g *getSession) {
	g.timer.Stop()
	c.sessions.mu.Lock()
	delete(c.sessions.gets, key)
	c.sessions.mu.Unlock()
	_ = g.closer.Close()
}

// getChunk answers the next chunk of a get session: {seq, last, data}, or {error}.
func (c *consumers) getChunk(module string, params json.RawMessage) (any, error) {
	var r buscalls.ObjectsGetChunkRequest
	if err := decode(buscalls.ObjectsGetChunk, params, &r); err != nil {
		return nil, err
	}
	key := sessionKey(module, r.Call)
	c.sessions.mu.Lock()
	g := c.sessions.gets[key]
	c.sessions.mu.Unlock()
	if g == nil {
		return buscalls.ObjectsChunk{Error: fmt.Sprintf("no get %s is open at chunk %d: it was dropped at its bound, or never begun", r.Call, r.Seq)}, nil
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	if r.Seq != g.next {
		c.dropGet(key, g)
		return buscalls.ObjectsChunk{Error: fmt.Sprintf("chunk %d was asked where chunk %d is next", r.Seq, g.next)}, nil
	}
	buf := make([]byte, buscalls.ChunkBytes)
	n, err := io.ReadFull(g.rd, buf)
	last := errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF)
	if err != nil && !last {
		c.dropGet(key, g)
		return buscalls.ObjectsChunk{Error: fmt.Sprintf("reading chunk %d: %v", r.Seq, err)}, nil
	}
	if !last {
		if _, perr := g.rd.Peek(1); errors.Is(perr, io.EOF) {
			last = true
		}
	}
	g.next++
	if last {
		c.dropGet(key, g)
	}
	return buscalls.ObjectsChunk{Seq: r.Seq, Last: last, Data: buf[:n]}, nil
}

// dropModuleSessions ends every session of a module whose bundle is gone.
func (c *consumers) dropModuleSessions(module string) {
	c.sessions.mu.Lock()
	var puts []*putSession
	var putKeys []string
	var gets []*getSession
	var getKeys []string
	for k, s := range c.sessions.puts {
		if strings.HasPrefix(k, module+" ") {
			puts, putKeys = append(puts, s), append(putKeys, k)
		}
	}
	for k, g := range c.sessions.gets {
		if strings.HasPrefix(k, module+" ") {
			gets, getKeys = append(gets, g), append(getKeys, k)
		}
	}
	c.sessions.mu.Unlock()
	for i, s := range puts {
		c.dropPut(putKeys[i], s, errors.New("the bundle's channel closed"))
	}
	for i, g := range gets {
		c.dropGet(getKeys[i], g)
	}
}

// sessionDigest is the `sha256:` form of an object's digest as the bus store keeps it (`SHA-256=` and base64), or as a
// reference carries it (hex).
func sessionDigest(d string) string {
	if raw, ok := strings.CutPrefix(d, "SHA-256="); ok {
		if b, err := base64.URLEncoding.DecodeString(raw); err == nil {
			return "sha256:" + hex.EncodeToString(b)
		}
	}
	return "sha256:" + strings.TrimPrefix(d, "sha256:")
}

func newCall() string {
	var b [12]byte
	_, _ = rand.Read(b[:])
	return hex.EncodeToString(b[:])
}

// CapabilityWords are the capability words this runner states, answered to a bundle's `mesh/capabilities` (§576).
func (b *moduleBus) CapabilityWords() []string { return CallCapabilities() }

// DropSessions ends the sessions of the module whose bundle's channel closed (§518).
func (b *moduleBus) DropSessions() { b.all.dropModuleSessions(b.module) }