oc-discovery/daemons/node/indexer/handler.go

package indexer

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"math/rand"
	"oc-discovery/daemons/node/common"
	"strings"
	"time"

	oclib "cloud.o-forge.io/core/oc-lib"
	"cloud.o-forge.io/core/oc-lib/dbs"
	pp "cloud.o-forge.io/core/oc-lib/models/peer"
	"cloud.o-forge.io/core/oc-lib/models/utils"
	"cloud.o-forge.io/core/oc-lib/tools"
	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/network"
	lpp "github.com/libp2p/go-libp2p/core/peer"
)

type PeerRecordPayload struct {
	Name       string    `json:"name"`
	DID        string    `json:"did"`
	PubKey     []byte    `json:"pub_key"`
	ExpiryDate time.Time `json:"expiry_date"`
}

type PeerRecord struct {
	PeerRecordPayload
	PeerID        string `json:"peer_id"`
	APIUrl        string `json:"api_url"`
	StreamAddress string `json:"stream_address"`
	NATSAddress   string `json:"nats_address"`
	WalletAddress string `json:"wallet_address"`
	Signature     []byte `json:"signature"`
}

func (p *PeerRecord) Sign() error {
	priv, err := tools.LoadKeyFromFilePrivate()
	if err != nil {
		return err
	}
	payload, _ := json.Marshal(p.PeerRecordPayload)
	b, err := common.Sign(priv, payload)
	p.Signature = b
	return err
}

func (p *PeerRecord) Verify() (crypto.PubKey, error) {
	pubKey, err := crypto.UnmarshalPublicKey(p.PubKey) // retrieve pub key in message
	if err != nil {
		return pubKey, err
	}
	payload, _ := json.Marshal(p.PeerRecordPayload)

	if ok, _ := pubKey.Verify(payload, p.Signature); !ok { // verify minimal message was sign per pubKey
		return pubKey, errors.New("invalid signature")
	}
	return pubKey, nil
}

func (pr *PeerRecord) ExtractPeer(ourkey string, key string, pubKey crypto.PubKey) (bool, *pp.Peer, error) {
	pubBytes, err := crypto.MarshalPublicKey(pubKey)
	if err != nil {
		return false, nil, err
	}
	rel := pp.NONE
	if ourkey == key { // at this point is PeerID is same as our... we are... thats our peer INFO
		rel = pp.SELF
	}

	p := &pp.Peer{
		AbstractObject: utils.AbstractObject{
			UUID: pr.DID,
			Name: pr.Name,
		},
		Relation:      rel, // VERIFY.... it crush nothing
		PeerID:        pr.PeerID,
		PublicKey:     base64.StdEncoding.EncodeToString(pubBytes),
		APIUrl:        pr.APIUrl,
		StreamAddress: pr.StreamAddress,
		NATSAddress:   pr.NATSAddress,
		WalletAddress: pr.WalletAddress,
	}
	if time.Now().UTC().After(pr.ExpiryDate) {
		return pp.SELF == p.Relation, nil, errors.New("peer " + key + " is offline")
	}
	return pp.SELF == p.Relation, p, nil
}

type GetValue struct {
	Key    string `json:"key"`
	PeerID string `json:"peer_id,omitempty"`
}

type GetResponse struct {
	Found   bool                  `json:"found"`
	Records map[string]PeerRecord `json:"records,omitempty"`
}

func (ix *IndexerService) genKey(did string) string {
	return "/node/" + did
}

func (ix *IndexerService) genPIDKey(peerID string) string {
	return "/pid/" + peerID
}

// isPeerKnown is the stream-level gate: returns true if pid is allowed.
// Check order (fast → slow):
//  1. In-memory stream records — currently heartbeating to this indexer.
//  2. Local DB by peer_id — known peer, blacklist enforced here.
//  3. DHT /pid/{peerID} → /node/{DID} — registered on any indexer.
//
// ProtocolHeartbeat and ProtocolPublish handlers do NOT call this — they are
// the streams through which a node first makes itself known.
func (ix *IndexerService) isPeerKnown(pid lpp.ID) bool {
	// 1. Fast path: active heartbeat session.
	ix.StreamMU.RLock()
	_, active := ix.StreamRecords[common.ProtocolHeartbeat][pid]
	ix.StreamMU.RUnlock()
	if active {
		return true
	}
	// 2. Local DB: known peer (handles blacklist).
	access := oclib.NewRequestAdmin(oclib.LibDataEnum(oclib.PEER), nil)
	results := access.Search(&dbs.Filters{
		And: map[string][]dbs.Filter{
			"peer_id": {{Operator: dbs.EQUAL.String(), Value: pid.String()}},
		},
	}, pid.String(), false)
	for _, item := range results.Data {
		p, ok := item.(*pp.Peer)
		if !ok || p.PeerID != pid.String() {
			continue
		}
		return p.Relation != pp.BLACKLIST
	}
	// 3. DHT lookup by peer_id.
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	did, err := ix.DHT.GetValue(ctx, ix.genPIDKey(pid.String()))
	cancel()
	if err != nil || len(did) == 0 {
		return false
	}
	ctx2, cancel2 := context.WithTimeout(context.Background(), 3*time.Second)
	_, err = ix.DHT.GetValue(ctx2, ix.genKey(string(did)))
	cancel2()
	return err == nil
}

func (ix *IndexerService) initNodeHandler() {
	logger := oclib.GetLogger()
	logger.Info().Msg("Init Node Handler")

	// Each heartbeat from a node carries a freshly signed PeerRecord.
	// Republish it to the DHT so the record never expires as long as the node
	// is alive — no separate publish stream needed from the node side.
	ix.AfterHeartbeat = func(hb *common.Heartbeat) {
		// Priority 1: use the fresh signed PeerRecord embedded in the heartbeat.
		// Each heartbeat tick, the node re-signs with ExpiryDate = now+2min, so
		// this record is always fresh. Fetching from DHT would give a stale expiry.
		var rec PeerRecord
		if len(hb.Record) > 0 {
			if err := json.Unmarshal(hb.Record, &rec); err != nil {
				logger.Warn().Err(err).Msg("indexer: heartbeat embedded record unmarshal failed")
				return
			}
		} else {
			// Fallback: node didn't embed a record yet (first heartbeat before claimInfo).
			// Fetch from DHT using the DID resolved by HandleHeartbeat.
			ctx2, cancel2 := context.WithTimeout(context.Background(), 10*time.Second)
			res, err := ix.DHT.GetValue(ctx2, ix.genKey(hb.DID))
			cancel2()
			if err != nil {
				logger.Warn().Err(err).Str("did", hb.DID).Msg("indexer: DHT fetch for refresh failed")
				return
			}
			if err := json.Unmarshal(res, &rec); err != nil {
				logger.Warn().Err(err).Str("did", hb.DID).Msg("indexer: heartbeat record unmarshal failed")
				return
			}
		}
		if _, err := rec.Verify(); err != nil {
			logger.Warn().Err(err).Str("did", rec.DID).Msg("indexer: heartbeat record signature invalid")
			return
		}
		// Keep StreamRecord.Record in sync so BuildHeartbeatResponse always
		// sees a populated PeerRecord (Name, DID, etc.) regardless of whether
		// handleNodePublish ran before or after the heartbeat stream was opened.
		if pid, err := lpp.Decode(rec.PeerID); err == nil {
			ix.StreamMU.Lock()
			if srec, ok := ix.StreamRecords[common.ProtocolHeartbeat][pid]; ok {
				srec.Record = rec
			}
			ix.StreamMU.Unlock()
		}
		data, err := json.Marshal(rec)
		if err != nil {
			return
		}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		if err := ix.DHT.PutValue(ctx, ix.genKey(rec.DID), data); err != nil {
			logger.Warn().Err(err).Str("did", rec.DID).Msg("indexer: DHT refresh /node/ failed")
		}
		cancel()
		// /pid/ is written unconditionally — the gater queries by PeerID and this
		// index must stay fresh regardless of whether the /node/ write succeeded.
		if rec.PeerID != "" {
			ctx2, cancel2 := context.WithTimeout(context.Background(), 10*time.Second)
			if err := ix.DHT.PutValue(ctx2, ix.genPIDKey(rec.PeerID), []byte(rec.DID)); err != nil {
				logger.Warn().Err(err).Str("pid", rec.PeerID).Msg("indexer: DHT refresh /pid/ failed")
			}
			cancel2()
		}
	}
	ix.Host.SetStreamHandler(common.ProtocolHeartbeat, ix.HandleHeartbeat)
	ix.Host.SetStreamHandler(common.ProtocolPublish, ix.handleNodePublish)
	ix.Host.SetStreamHandler(common.ProtocolGet, ix.handleNodeGet)
	ix.Host.SetStreamHandler(common.ProtocolIndexerCandidates, ix.handleCandidateRequest)
	ix.initSearchHandlers()
}

// handleCandidateRequest responds to a node's consensus candidate request.
// Returns a random sample of indexers from the local DHT cache.
func (ix *IndexerService) handleCandidateRequest(s network.Stream) {
	defer s.Close()
	if !ix.isPeerKnown(s.Conn().RemotePeer()) {
		logger := oclib.GetLogger()
		logger.Warn().Str("peer", s.Conn().RemotePeer().String()).Msg("[candidates] unknown peer, rejecting stream")
		s.Reset()
		return
	}
	s.SetDeadline(time.Now().Add(5 * time.Second))
	var req common.IndexerCandidatesRequest
	if err := json.NewDecoder(s).Decode(&req); err != nil {
		return
	}
	if req.Count <= 0 || req.Count > 10 {
		req.Count = 3
	}
	ix.dhtCacheMu.RLock()
	cache := make([]dhtCacheEntry, len(ix.dhtCache))
	copy(cache, ix.dhtCache)
	ix.dhtCacheMu.RUnlock()

	// Shuffle for randomness: each voter offers a different subset.
	rand.Shuffle(len(cache), func(i, j int) { cache[i], cache[j] = cache[j], cache[i] })
	candidates := make([]lpp.AddrInfo, 0, req.Count)
	for _, e := range cache {
		if len(candidates) >= req.Count {
			break
		}
		candidates = append(candidates, e.AI)
	}
	json.NewEncoder(s).Encode(common.IndexerCandidatesResponse{Candidates: candidates})
}

func (ix *IndexerService) handleNodePublish(s network.Stream) {
	defer s.Close()
	logger := oclib.GetLogger()
	remotePeer := s.Conn().RemotePeer()
	if err := ix.behavior.RecordPublish(remotePeer); err != nil {
		logger.Warn().Err(err).Str("peer", remotePeer.String()).Msg("publish refused")
		s.Reset()
		return
	}
	for {
		var rec PeerRecord
		if err := json.NewDecoder(s).Decode(&rec); err != nil {
			logger.Err(err)
			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) ||
				strings.Contains(err.Error(), "reset") ||
				strings.Contains(err.Error(), "closed") ||
				strings.Contains(err.Error(), "too many connections") {
				return
			}
			continue
		}
		if _, err := rec.Verify(); err != nil {
			ix.behavior.RecordBadSignature(remotePeer)
			logger.Warn().Err(err).Str("peer", remotePeer.String()).Msg("bad signature on publish")
			return
		}
		if err := ix.behavior.CheckIdentity(remotePeer, rec.DID); err != nil {
			logger.Warn().Err(err).Msg("identity mismatch on publish")
			s.Reset()
			return
		}
		if rec.PeerID == "" || rec.ExpiryDate.Before(time.Now().UTC()) {
			logger.Err(errors.New(rec.PeerID + " is expired."))
			return
		}
		pid, err := lpp.Decode(rec.PeerID)
		if err != nil {
			return
		}

		ix.StreamMU.Lock()
		defer ix.StreamMU.Unlock()
		if ix.StreamRecords[common.ProtocolHeartbeat] == nil {
			ix.StreamRecords[common.ProtocolHeartbeat] = map[lpp.ID]*common.StreamRecord[PeerRecord]{}
		}
		streams := ix.StreamRecords[common.ProtocolHeartbeat]
		if srec, ok := streams[pid]; ok {
			srec.DID = rec.DID
			srec.Record = rec
			srec.HeartbeatStream.UptimeTracker.LastSeen = time.Now().UTC()
		}

		key := ix.genKey(rec.DID)
		data, err := json.Marshal(rec)
		if err != nil {
			logger.Err(err)
			return
		}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		if err := ix.DHT.PutValue(ctx, key, data); err != nil {
			logger.Err(err)
			cancel()
			return
		}
		cancel()

		// Secondary index: /pid/<peerID> → DID, so peers can resolve by libp2p PeerID.
		if rec.PeerID != "" {
			ctx2, cancel2 := context.WithTimeout(context.Background(), 10*time.Second)
			if err := ix.DHT.PutValue(ctx2, ix.genPIDKey(rec.PeerID), []byte(rec.DID)); err != nil {
				logger.Err(err).Str("pid", rec.PeerID).Msg("indexer: failed to write pid index")
			}
			cancel2()
		}
		return
	}
}

func (ix *IndexerService) handleNodeGet(s network.Stream) {
	defer s.Close()
	logger := oclib.GetLogger()
	remotePeer := s.Conn().RemotePeer()
	if !ix.isPeerKnown(remotePeer) {
		logger.Warn().Str("peer", remotePeer.String()).Msg("[get] unknown peer, rejecting stream")
		s.Reset()
		return
	}
	if err := ix.behavior.RecordGet(remotePeer); err != nil {
		logger.Warn().Err(err).Str("peer", remotePeer.String()).Msg("get refused")
		s.Reset()
		return
	}
	for {
		var req GetValue
		if err := json.NewDecoder(s).Decode(&req); err != nil {
			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) ||
				strings.Contains(err.Error(), "reset") ||
				strings.Contains(err.Error(), "closed") ||
				strings.Contains(err.Error(), "too many connections") {
				return
			}
			logger.Err(err)
			continue
		}

		resp := GetResponse{Found: false, Records: map[string]PeerRecord{}}

		// Resolve DID key: by PeerID (secondary /pid/ index) or direct DID key.
		var key string
		if req.PeerID != "" {
			pidCtx, pidCancel := context.WithTimeout(context.Background(), 5*time.Second)
			did, err := ix.DHT.GetValue(pidCtx, ix.genPIDKey(req.PeerID))
			pidCancel()
			if err == nil {
				key = string(did)
			}
		} else {
			key = req.Key
		}

		if key != "" {
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			c, err := ix.DHT.GetValue(ctx, ix.genKey(key))
			cancel()
			if err == nil {
				var rec PeerRecord
				if json.Unmarshal(c, &rec) == nil {
					resp.Records[rec.PeerID] = rec
				}
			} else {
				logger.Err(err).Msg("Failed to fetch PeerRecord from DHT " + key)
			}
		}

		resp.Found = len(resp.Records) > 0
		_ = json.NewEncoder(s).Encode(resp)
		break
	}
}