verae-mcp-timestamping-server-go/pkg/tsnats/client.go

package tsnats

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

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

	"git.georgelambert.org/marchon/verae-mcp-timestamping-server-go/pkg/digest"
)

const defaultTimeout = 2 * time.Second

// Requester is the request/reply surface used by Stamp/Verify/HashReport.
// Inject a mock in unit tests; Connect wraps *nats.Conn.
type Requester interface {
	Request(ctx context.Context, subj string, data []byte) ([]byte, error)
}

// MessageHandler receives raw message payloads from a core NATS subscription.
type MessageHandler func(data []byte)

// SubConn is the subscribe surface for SubscribeCompleted.
type SubConn interface {
	Subscribe(subj string, cb MessageHandler) (Subscription, error)
}

// Subscription is returned by SubscribeCompleted.
type Subscription interface {
	Unsubscribe() error
}

// Receipt is a timestamp receipt from request/reply or hash.report.
type Receipt struct {
	ReceiptID string `json:"receipt_id"`
	Hash      string `json:"hash"`
	Scope     string `json:"scope"`
	CubeID    string `json:"cube_id,omitempty"`
	OrgID     string `json:"org_id,omitempty"`
	CreatedAt string `json:"created_at"`
	Merkle    string `json:"merkle"`
	RestampOf string `json:"restamp_of,omitempty"`
}

// VerifyResult is the outcome of Verify.
type VerifyResult struct {
	Valid   bool
	Error   string
	Receipt *Receipt
}

// CompletedEvent is a fire-and-forget job-complete notification on SubjectCompleted.
type CompletedEvent struct {
	JobID     string `json:"jobId"`
	ReceiptID string `json:"receipt_id"`
	Hash      string `json:"hash,omitempty"`
}

// StampOptions configures Stamp. Authorization/AccessToken are optional
// passthrough fields for embedded handlers (JSON keys authorization / access_token).
type StampOptions struct {
	Scope         string
	CubeID        string
	OrgID         string
	Authorization string
	AccessToken   string
}

// VerifyOptions configures Verify.
type VerifyOptions struct {
	WantScope     string
	Authorization string
	AccessToken   string
}

// ReportOptions configures HashReport.
type ReportOptions struct {
	OrgID         string
	WantScope     string
	Authorization string
	AccessToken   string
}

// Option configures Client construction.
type Option func(*Client)

// WithTimeout sets the request timeout (default 2s).
func WithTimeout(d time.Duration) Option {
	return func(c *Client) {
		if d > 0 {
			c.timeout = d
		}
	}
}

// WithRequester injects a Requester (tests or custom transports).
func WithRequester(r Requester) Option {
	return func(c *Client) { c.req = r }
}

// WithSubscriber injects a SubConn for SubscribeCompleted.
func WithSubscriber(s SubConn) Option {
	return func(c *Client) { c.sub = s }
}

// Client is a digests-only NATS timestamping client (core NATS only; no JetStream).
type Client struct {
	req     Requester
	sub     SubConn
	timeout time.Duration
	nc      *nats.Conn
}

// NewClient builds a Client with an injected Requester (unit tests).
func NewClient(r Requester, opts ...Option) *Client {
	c := &Client{req: r, timeout: defaultTimeout}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

// Connect dials a NATS URL and returns a Client using core request/reply and subscribe.
func Connect(url string, opts ...Option) (*Client, error) {
	nc, err := nats.Connect(url)
	if err != nil {
		return nil, err
	}
	wrap := &natsAdapter{nc: nc}
	c := &Client{req: wrap, sub: wrap, timeout: defaultTimeout, nc: nc}
	for _, opt := range opts {
		opt(c)
	}
	if c.req == nil {
		c.req = wrap
	}
	if c.sub == nil {
		c.sub = wrap
	}
	return c, nil
}

// Close closes the underlying NATS connection when opened via Connect.
func (c *Client) Close() {
	if c.nc != nil {
		c.nc.Close()
	}
}

// Timeout returns the request timeout.
func (c *Client) Timeout() time.Duration {
	return c.timeout
}

// Stamp publishes a digests-only stamp request on SubjectRequest.
func (c *Client) Stamp(ctx context.Context, sha256 string, opts StampOptions) ([]Receipt, error) {
	h, err := digest.NormalizeSHA256(sha256)
	if err != nil {
		return nil, err
	}
	scope := opts.Scope
	if scope == "" {
		scope = "local"
	}
	payload := map[string]any{
		"hash":    h,
		"scope":   scope,
		"cube_id": opts.CubeID,
		"org_id":  opts.OrgID,
	}
	putAuth(payload, opts.Authorization, opts.AccessToken)
	raw, err := c.request(ctx, SubjectRequest, payload)
	if err != nil {
		return nil, err
	}
	return parseStampReply(raw)
}

// Verify asks SubjectVerify for a receipt by id.
func (c *Client) Verify(ctx context.Context, receiptID string, opts VerifyOptions) (VerifyResult, error) {
	if receiptID == "" {
		return VerifyResult{Valid: false, Error: "receipt_id required"}, nil
	}
	payload := map[string]any{
		"receipt_id": receiptID,
		"want_scope": opts.WantScope,
	}
	putAuth(payload, opts.Authorization, opts.AccessToken)
	raw, err := c.request(ctx, SubjectVerify, payload)
	if err != nil {
		return VerifyResult{}, err
	}
	return parseVerifyReply(raw)
}

// HashReport asks SubjectReport for receipts for a digest.
func (c *Client) HashReport(ctx context.Context, sha256 string, opts ReportOptions) ([]Receipt, error) {
	h, err := digest.NormalizeSHA256(sha256)
	if err != nil {
		return nil, err
	}
	payload := map[string]any{
		"hash":       h,
		"org_id":     opts.OrgID,
		"want_scope": opts.WantScope,
	}
	putAuth(payload, opts.Authorization, opts.AccessToken)
	raw, err := c.request(ctx, SubjectReport, payload)
	if err != nil {
		return nil, err
	}
	return parseReportReply(raw)
}

// SubscribeCompleted core-subscribes to SubjectCompleted (not JetStream).
func (c *Client) SubscribeCompleted(ctx context.Context, handler func(CompletedEvent)) (Subscription, error) {
	if c.sub == nil {
		return nil, fmt.Errorf("tsnats: no subscriber configured")
	}
	if handler == nil {
		return nil, fmt.Errorf("tsnats: handler required")
	}
	_ = ctx
	return c.sub.Subscribe(SubjectCompleted, func(data []byte) {
		var ev CompletedEvent
		if err := json.Unmarshal(data, &ev); err != nil {
			return
		}
		handler(ev)
	})
}

func (c *Client) request(ctx context.Context, subj string, payload map[string]any) ([]byte, error) {
	if c.req == nil {
		return nil, fmt.Errorf("tsnats: no requester configured")
	}
	data, err := json.Marshal(payload)
	if err != nil {
		return nil, err
	}
	ctx, cancel := withTimeout(ctx, c.timeout)
	defer cancel()
	return c.req.Request(ctx, subj, data)
}

func putAuth(payload map[string]any, authorization, accessToken string) {
	if authorization != "" {
		payload["authorization"] = authorization
	}
	if accessToken != "" {
		payload["access_token"] = accessToken
	}
}

func withTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
	if _, ok := ctx.Deadline(); ok {
		return context.WithCancel(ctx)
	}
	return context.WithTimeout(ctx, d)
}

func parseStampReply(raw []byte) ([]Receipt, error) {
	var asList []Receipt
	if err := json.Unmarshal(raw, &asList); err == nil && (len(asList) > 0 || isJSONArray(raw)) {
		return asList, nil
	}
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil {
		return nil, fmt.Errorf("tsnats: stamp reply: %w", err)
	}
	if errRaw, ok := obj["error"]; ok {
		var msg string
		_ = json.Unmarshal(errRaw, &msg)
		if msg == "" {
			msg = string(errRaw)
		}
		return nil, fmt.Errorf("%s", msg)
	}
	if _, hasReceipt := obj["receipt_id"]; hasReceipt {
		var rec Receipt
		if err := json.Unmarshal(raw, &rec); err != nil {
			return nil, err
		}
		return []Receipt{rec}, nil
	}
	if jobRaw, ok := obj["jobId"]; ok {
		var jobID string
		_ = json.Unmarshal(jobRaw, &jobID)
		return nil, fmt.Errorf("stamp pending jobId=%s; use SubscribeCompleted", jobID)
	}
	return nil, fmt.Errorf("tsnats: unrecognized stamp reply")
}

func parseVerifyReply(raw []byte) (VerifyResult, error) {
	var obj map[string]any
	if err := json.Unmarshal(raw, &obj); err != nil {
		return VerifyResult{Valid: false, Error: "invalid verify payload"}, nil
	}
	if v, ok := obj["valid"]; ok {
		if b, _ := v.(bool); !b {
			errMsg, _ := obj["error"].(string)
			return VerifyResult{Valid: false, Error: errMsg}, nil
		}
	}
	if errMsg, ok := obj["error"].(string); ok && errMsg != "" {
		if valid, _ := obj["valid"].(bool); !valid {
			return VerifyResult{Valid: false, Error: errMsg}, nil
		}
	}
	if _, ok := obj["receipt_id"]; ok {
		var rec Receipt
		if err := json.Unmarshal(raw, &rec); err != nil {
			return VerifyResult{Valid: false, Error: err.Error()}, nil
		}
		return VerifyResult{Valid: true, Receipt: &rec}, nil
	}
	return VerifyResult{Valid: false, Error: "unknown receipt"}, nil
}

func parseReportReply(raw []byte) ([]Receipt, error) {
	var asList []Receipt
	if err := json.Unmarshal(raw, &asList); err == nil && (len(asList) > 0 || isJSONArray(raw)) {
		return asList, nil
	}
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil {
		return nil, nil
	}
	if errRaw, ok := obj["error"]; ok {
		var msg string
		_ = json.Unmarshal(errRaw, &msg)
		if msg == "" {
			msg = string(errRaw)
		}
		return nil, fmt.Errorf("%s", msg)
	}
	if receiptsRaw, ok := obj["receipts"]; ok {
		var receipts []Receipt
		if err := json.Unmarshal(receiptsRaw, &receipts); err != nil {
			return nil, err
		}
		return receipts, nil
	}
	if timelineRaw, ok := obj["timeline"]; ok {
		var events []map[string]json.RawMessage
		if err := json.Unmarshal(timelineRaw, &events); err != nil {
			return nil, nil
		}
		var out []Receipt
		for _, ev := range events {
			recRaw, ok := ev["receipt"]
			if !ok {
				continue
			}
			var rec Receipt
			if err := json.Unmarshal(recRaw, &rec); err != nil {
				continue
			}
			out = append(out, rec)
		}
		return out, nil
	}
	return nil, nil
}

func isJSONArray(raw []byte) bool {
	for _, b := range raw {
		if b == ' ' || b == '\t' || b == '\n' || b == '\r' {
			continue
		}
		return b == '['
	}
	return false
}

type natsAdapter struct {
	nc *nats.Conn
}

func (a *natsAdapter) Request(ctx context.Context, subj string, data []byte) ([]byte, error) {
	msg, err := a.nc.RequestWithContext(ctx, subj, data)
	if err != nil {
		return nil, err
	}
	return msg.Data, nil
}

func (a *natsAdapter) Subscribe(subj string, cb MessageHandler) (Subscription, error) {
	sub, err := a.nc.Subscribe(subj, func(msg *nats.Msg) {
		cb(msg.Data)
	})
	if err != nil {
		return nil, err
	}
	return sub, nil
}