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

package indexer

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

	oclib "cloud.o-forge.io/core/oc-lib"
	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"
	"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,
	}
	b, err := json.Marshal(p)
	if err != nil {
		return pp.SELF == p.Relation, nil, err
	}

	if time.Now().UTC().After(pr.ExpiryDate) {
		return pp.SELF == p.Relation, nil, errors.New("peer " + key + " is offline")
	}
	go tools.NewNATSCaller().SetNATSPub(tools.CREATE_RESOURCE, tools.NATSResponse{
		FromApp:    "oc-discovery",
		Datatype:   tools.PEER,
		Method:     int(tools.CREATE_RESOURCE),
		SearchAttr: "peer_id",
		Payload:    b,
	})

	return pp.SELF == p.Relation, p, nil
}

type GetValue struct {
	Key    string  `json:"key"`
	PeerID peer.ID `json:"peer_id"`
	Name   string  `json:"name,omitempty"`
	Search bool    `json:"search,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) genNameKey(name string) string {
	return "/name/" + name
}

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

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(pid peer.ID) {
		ctx1, cancel1 := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel1()
		res, err := ix.DHT.GetValue(ctx1, ix.genPIDKey(pid.String()))
		if err != nil {
			logger.Warn().Err(err)
			return
		}
		did := string(res)
		ctx2, cancel2 := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel2()
		res, err = ix.DHT.GetValue(ctx2, ix.genKey(did))
		if err != nil {
			logger.Warn().Err(err)
			return
		}
		var rec PeerRecord
		if err := json.Unmarshal(res, &rec); err != nil {
			logger.Warn().Err(err).Str("peer", pid.String()).Msg("indexer: heartbeat record unmarshal failed")
			return
		}
		if _, err := rec.Verify(); err != nil {
			logger.Warn().Err(err).Str("peer", pid.String()).Msg("indexer: heartbeat record signature invalid")
			return
		}
		data, err := json.Marshal(rec)
		if err != nil {
			return
		}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		logger.Info().Msg("REFRESH PutValue " + ix.genKey(rec.DID))
		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 failed")
			return
		}
		if rec.Name != "" {
			ctx2, cancel2 := context.WithTimeout(context.Background(), 10*time.Second)
			ix.DHT.PutValue(ctx2, ix.genNameKey(rec.Name), []byte(rec.DID))
			cancel2()
		}
		if rec.PeerID != "" {
			ctx3, cancel3 := context.WithTimeout(context.Background(), 10*time.Second)
			ix.DHT.PutValue(ctx3, ix.genPIDKey(rec.PeerID), []byte(rec.DID))
			cancel3()
		}
	}
	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.ProtocolIndexerGetNatives, ix.handleGetNatives)
}

func (ix *IndexerService) handleNodePublish(s network.Stream) {
	defer s.Close()
	logger := oclib.GetLogger()

	var rec PeerRecord
	if err := json.NewDecoder(s).Decode(&rec); err != nil {
		logger.Err(err)
		return
	}
	if _, err := rec.Verify(); err != nil {
		logger.Err(err)
		return
	}
	if rec.PeerID == "" || rec.ExpiryDate.Before(time.Now().UTC()) {
		logger.Err(errors.New(rec.PeerID + " is expired."))
		return
	}
	pid, err := peer.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[peer.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: /name/<name> → DID, so peers can resolve by human-readable name.
	if rec.Name != "" {
		ctx2, cancel2 := context.WithTimeout(context.Background(), 10*time.Second)
		if err := ix.DHT.PutValue(ctx2, ix.genNameKey(rec.Name), []byte(rec.DID)); err != nil {
			logger.Err(err).Str("name", rec.Name).Msg("indexer: failed to write name index")
		}
		cancel2()
	}
	// Secondary index: /pid/<peerID> → DID, so peers can resolve by libp2p PeerID.
	if rec.PeerID != "" {
		ctx3, cancel3 := context.WithTimeout(context.Background(), 10*time.Second)
		if err := ix.DHT.PutValue(ctx3, ix.genPIDKey(rec.PeerID), []byte(rec.DID)); err != nil {
			logger.Err(err).Str("pid", rec.PeerID).Msg("indexer: failed to write pid index")
		}
		cancel3()
	}
}

func (ix *IndexerService) handleNodeGet(s network.Stream) {

	defer s.Close()
	logger := oclib.GetLogger()

	var req GetValue
	if err := json.NewDecoder(s).Decode(&req); err != nil {
		logger.Err(err)
		return
	}

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

	keys := []string{}
	// Name substring search — scan in-memory connected nodes first, then DHT exact match.
	if req.Name != "" {
		if req.Search {
			for _, did := range ix.LookupNameIndex(strings.ToLower(req.Name)) {
				keys = append(keys, did)
			}
		} else {
			// 2. DHT exact-name lookup: covers nodes that published but aren't currently connected.
			nameCtx, nameCancel := context.WithTimeout(context.Background(), 5*time.Second)
			if ch, err := ix.DHT.SearchValue(nameCtx, ix.genNameKey(req.Name)); err == nil {
				for did := range ch {
					keys = append(keys, string(did))
					break
				}
			}
			nameCancel()
		}
	} else if req.PeerID != "" {
		pidCtx, pidCancel := context.WithTimeout(context.Background(), 5*time.Second)
		if did, err := ix.DHT.GetValue(pidCtx, ix.genPIDKey(req.PeerID.String())); err == nil {
			keys = append(keys, string(did))
		}
		pidCancel()
	} else {
		keys = append(keys, req.Key)
	}

	// DHT record fetch by DID key (covers exact-name and PeerID paths).
	if len(keys) > 0 {
		for _, k := range keys {
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			c, err := ix.DHT.GetValue(ctx, ix.genKey(k))
			cancel()
			if err == nil {
				var rec PeerRecord
				if json.Unmarshal(c, &rec) == nil {
					// Filter by PeerID only when one was explicitly specified.
					if req.PeerID == "" || rec.PeerID == req.PeerID.String() {
						resp.Records[rec.PeerID] = rec
					}
				}
			} else if req.Name == "" && req.PeerID == "" {
				logger.Err(err).Msg("Failed to fetch PeerRecord from DHT " + req.Key)
			}
		}
	}

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

// handleGetNatives returns this indexer's configured native addresses,
// excluding any in the request's Exclude list.
func (ix *IndexerService) handleGetNatives(s network.Stream) {
	defer s.Close()
	logger := oclib.GetLogger()

	var req common.GetIndexerNativesRequest
	if err := json.NewDecoder(s).Decode(&req); err != nil {
		logger.Err(err).Msg("indexer get natives: decode")
		return
	}

	excludeSet := make(map[string]struct{}, len(req.Exclude))
	for _, e := range req.Exclude {
		excludeSet[e] = struct{}{}
	}

	resp := common.GetIndexerNativesResponse{}
	for _, addr := range strings.Split(conf.GetConfig().NativeIndexerAddresses, ",") {
		addr = strings.TrimSpace(addr)
		if addr == "" {
			continue
		}
		if _, excluded := excludeSet[addr]; !excluded {
			resp.Natives = append(resp.Natives, addr)
		}
	}

	if err := json.NewEncoder(s).Encode(resp); err != nil {
		logger.Err(err).Msg("indexer get natives: encode response")
	}
}