didplcbft / abciapp / import.go

package abciapp

import (
	"bufio"
	"context"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/did-method-plc/go-didplc"
	"github.com/ipfs/go-cid"
	cbornode "github.com/ipfs/go-ipld-cbor"
	"github.com/palantir/stacktrace"
	"tangled.org/gbl08ma.com/didplcbft/store"
)

const EagerFetchMaxOps = 10000
const OpsPerImportTx = 1000
const OpsPerEagerFetch = 1000

type authoritativeOperationsCache struct {
	mu sync.Mutex

	plcURL               string
	operations           map[uint64]logEntryWithSeq
	highestFetchedHeight uint64
}

type logEntryWithSeq struct {
	didplc.LogEntry
	Seq uint64 `json:"seq"`
}

func newAuthoritativeOperationsCache(ctx context.Context, plc string) *authoritativeOperationsCache {
	aoc := &authoritativeOperationsCache{
		plcURL:     plc,
		operations: make(map[uint64]logEntryWithSeq),
	}

	go func() {
		ticker := time.NewTicker(500 * time.Millisecond)
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				aoc.eagerlyFetch(ctx)
			}
		}
	}()

	return aoc
}

func getOrCreateAuthoritativeOperationsCache(ctx context.Context, aocsByPLC map[string]*authoritativeOperationsCache, plc string) *authoritativeOperationsCache {
	aoc, ok := aocsByPLC[plc]
	if !ok {
		aoc = newAuthoritativeOperationsCache(ctx, plc)
		aocsByPLC[plc] = aoc
	}
	return aoc
}

func (a *authoritativeOperationsCache) eagerlyFetch(ctx context.Context) {
	a.mu.Lock()
	defer a.mu.Unlock()

	curOps := len(a.operations)
	if curOps >= EagerFetchMaxOps {
		return
	}
	_, _ = a.fetchInMutex(ctx, a.highestFetchedHeight, OpsPerEagerFetch)
}

func (a *authoritativeOperationsCache) dropSeqBelowOrEqual(highestCommittedSeq uint64) {
	a.mu.Lock()
	defer a.mu.Unlock()

	for i := range a.operations {
		if a.operations[i].Seq <= highestCommittedSeq {
			delete(a.operations, i)
		}
	}
}

func (a *authoritativeOperationsCache) fetchInMutex(ctx context.Context, after, count uint64) (bool, error) {
	entries, _, err := fetchExportedBatchFromAuthoritativeSource(ctx, a.plcURL, after, count)
	if err != nil {
		return false, stacktrace.Propagate(err, "")
	}

	for _, entry := range entries {
		a.operations[entry.Seq] = entry
		a.highestFetchedHeight = max(a.highestFetchedHeight, entry.Seq)
	}
	return uint64(len(entries)) < count, nil
}

func (a *authoritativeOperationsCache) get(ctx context.Context, after, count uint64) ([]logEntryWithSeq, error) {
	a.mu.Lock()
	defer a.mu.Unlock()

	result := make([]logEntryWithSeq, 0, count)
	reachedEnd := false
	for i := uint64(0); uint64(len(result)) < count; i++ {
		opSeq := after + i + 1
		op, ok := a.operations[opSeq]
		if !ok {
			if reachedEnd {
				// it's because we are asking about ops that don't exist yet, return
				break
			}

			re, err := a.fetchInMutex(ctx, after+i, count)
			if err != nil {
				return nil, stacktrace.Propagate(err, "")
			}

			reachedEnd = reachedEnd || re

			op, ok = a.operations[opSeq]
			if !ok {
				// still not present even after fetching
				// the authoritative source probably skipped this seq?
				continue
			}
		}

		result = append(result, op)
	}

	return result, nil
}

func fetchExportedBatchFromAuthoritativeSource(ctx context.Context, plcURL string, startAt, maxCount uint64) ([]logEntryWithSeq, uint64, error) {
	baseURL, err := url.JoinPath(plcURL, "/export")
	if err != nil {
		return nil, 0, stacktrace.Propagate(err, "")
	}

	client := &http.Client{Timeout: 30 * time.Second}

	entries := make([]logEntryWithSeq, 0, maxCount)
	for {
		req, err := http.NewRequestWithContext(ctx, "GET", baseURL, nil)
		if err != nil {
			return nil, 0, stacktrace.Propagate(err, "")
		}

		req.Header.Set("User-Agent", "didplcbft")

		requestCount := min(1000, maxCount-uint64(len(entries)))

		q := req.URL.Query()
		q.Add("count", fmt.Sprint(requestCount))
		q.Add("after", fmt.Sprint(startAt))
		req.URL.RawQuery = q.Encode()

		resp, err := client.Do(req)
		if err != nil {
			return nil, 0, stacktrace.Propagate(err, "")
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			return nil, 0, stacktrace.NewError("non-200 status code")
		}

		// Read response body
		s := bufio.NewScanner(resp.Body)
		numEntriesThisResponse := 0
		for s.Scan() && len(entries) < int(maxCount) {
			var entry logEntryWithSeq
			if err := json.Unmarshal(s.Bytes(), &entry); err != nil {
				return nil, 0, stacktrace.Propagate(err, "")
			}
			entries = append(entries, entry)
			numEntriesThisResponse++
			startAt = entry.Seq
		}
		if s.Err() != nil {
			return nil, 0, stacktrace.Propagate(s.Err(), "")
		}

		if uint64(numEntriesThisResponse) < requestCount || len(entries) >= int(maxCount) {
			break
		}
	}

	return entries, startAt, nil
}

func computeLogEntriesHash(logEntries []logEntryWithSeq) ([]byte, error) {
	// let's _not_ rely on the specifics of the JSON representation
	// (instead let's rely on specifics of our implementation, heh)

	hash := sha256.New()
	for i, entry := range logEntries {
		// Write DID
		didBytes, err := store.DIDToBytes(entry.DID)
		if err != nil {
			return nil, stacktrace.Propagate(err, "invalid DID in entry index %d", i)
		}

		_, err = hash.Write(didBytes)
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}

		// Write CID
		// (We trust that the authoritative source computed CIDs properly, so we use theirs rather than decoding the operation and recomputing)
		cid, err := cid.Decode(entry.CID)
		if err != nil {
			return nil, stacktrace.Propagate(err, "invalid CID in entry index %d", i)
		}

		_, err = hash.Write(cid.Bytes())
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}

		// Write CreatedAt
		createdAt, err := syntax.ParseDatetime(entry.CreatedAt)
		if err != nil {
			return nil, stacktrace.Propagate(err, "invalid CreatedAt in entry index %d", i)
		}

		tsBytes := make([]byte, 8)
		binary.BigEndian.PutUint64(tsBytes, uint64(createdAt.Time().Truncate(1*time.Millisecond).UTC().UnixNano()))

		_, err = hash.Write(tsBytes)
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}

		// Nullified can't be part of the hash as it can change on the authoritative source at any moment,
		// we always import operations as if they weren't nullified and recompute the nullification status as needed
	}

	return hash.Sum(nil), nil
}

func (d *DIDPLCApplication) maybeCreateAuthoritativeImportTx(ctx context.Context) ([]byte, error) {
	// use WorkingTreeVersion so we take into account any import operation that may have been processed in this block
	readTx := d.txFactory.ReadWorking(time.Now())

	plcURL, err := store.Tree.AuthoritativePLC(readTx)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}

	if plcURL == "" {
		// we're not doing imports
		return nil, nil
	}

	cursor, err := store.Tree.AuthoritativeImportProgress(readTx)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}

	aoc := getOrCreateAuthoritativeOperationsCache(d.runnerContext, d.aocsByPLC, plcURL)

	entries, err := aoc.get(ctx, cursor, OpsPerImportTx)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}

	if len(entries) == 0 {
		// nothing to import at the moment
		return nil, nil
	}

	hashBytes, err := computeLogEntriesHash(entries)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}

	tx := Transaction[AuthoritativeImportArguments]{
		Action: TransactionActionAuthoritativeImport,
		Arguments: AuthoritativeImportArguments{
			PLCURL: plcURL,
			Hash:   hex.EncodeToString(hashBytes),
			Cursor: cursor,
			Count:  uint64(len(entries)),
		},
	}

	out, err := cbornode.DumpObject(tx)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}
	return out, nil
}