handlers.go

package zikade

import (
	"bytes"
	"context"
	"errors"
	"fmt"

	ds "github.com/ipfs/go-datastore"
	record "github.com/libp2p/go-libp2p-record"
	recpb "github.com/libp2p/go-libp2p-record/pb"
	"github.com/libp2p/go-libp2p/core/peer"
	"go.opentelemetry.io/otel/attribute"
	otel "go.opentelemetry.io/otel/trace"
	"golang.org/x/exp/slog"

	"github.com/plprobelab/zikade/kadt"
	"github.com/plprobelab/zikade/pb"
	"github.com/plprobelab/zikade/private_routing"
)

// handleFindPeer handles FIND_NODE requests from remote peers.
func (d *DHT) handleFindPeer(ctx context.Context, remote peer.ID, req *pb.Message) (*pb.Message, error) {
	if len(req.GetKey()) == 0 {
		return nil, fmt.Errorf("handleFindPeer with empty key")
	}

	// tell the coordinator that this peer supports finding closer nodes
	d.kad.NotifyConnectivity(ctx, kadt.PeerID(remote))

	// "parse" requested peer ID from the key field
	target := peer.ID(req.GetKey())

	// initialize the response message
	resp := &pb.Message{
		Type: pb.Message_FIND_NODE,
		Key:  req.GetKey(),
	}

	// get reference to peer store
	pstore := d.host.Peerstore()

	// if the remote is asking for us, short-circuit and return us only
	if target == d.host.ID() {
		resp.CloserPeers = []*pb.Message_Peer{pb.FromAddrInfo(pstore.PeerInfo(d.host.ID()))}
		return resp, nil
	}

	// gather closer peers that we know
	resp.CloserPeers = d.closerPeers(ctx, remote, kadt.PeerID(target).Key())

	// if we happen to know the target peers addresses (e.g., although we are
	// far away in the keyspace), we add the peer to the result set. This means
	// we potentially return bucketSize + 1 peers. We won't add the peer to the
	// response if it's already contained in CloserPeers.
	targetInfo := pstore.PeerInfo(target)
	if len(targetInfo.Addrs) > 0 && !resp.ContainsCloserPeer(target) && target != remote {
		resp.CloserPeers = append(resp.CloserPeers, pb.FromAddrInfo(targetInfo))
	}

	return resp, nil
}

// handlePing handles PING requests from remote peers.
func (d *DHT) handlePing(ctx context.Context, remote peer.ID, req *pb.Message) (*pb.Message, error) {
	d.log.LogAttrs(ctx, slog.LevelDebug, "Responding to ping", slog.String("remote", remote.String()))
	return &pb.Message{Type: pb.Message_PING}, nil
}

// handleGetValue handles PUT_VALUE RPCs from remote peers.
func (d *DHT) handlePutValue(ctx context.Context, remote peer.ID, req *pb.Message) (*pb.Message, error) {
	// validate incoming request -> key and record must not be empty/nil
	k := string(req.GetKey())
	if len(k) == 0 {
		return nil, fmt.Errorf("no key was provided")
	}

	rec := req.GetRecord()
	if rec == nil {
		return nil, fmt.Errorf("nil record")
	}

	if !bytes.Equal(req.GetKey(), rec.GetKey()) {
		return nil, fmt.Errorf("key doesn't match record key")
	}

	// TODO: use putValueLocal?

	// key is /$namespace/$binary_id
	ns, path, err := record.SplitKey(k) // get namespace (prefix of the key)
	if err != nil || len(path) == 0 {
		return nil, fmt.Errorf("invalid key %s: %w", k, err)
	}

	backend, found := d.backends[ns]
	if !found {
		return nil, fmt.Errorf("unsupported record type: %s", ns)
	}

	_, err = backend.Store(ctx, path, rec)

	return nil, err
}

// handleGetValue handles GET_VALUE RPCs from remote peers.
func (d *DHT) handleGetValue(ctx context.Context, remote peer.ID, req *pb.Message) (*pb.Message, error) {
	k := string(req.GetKey())
	if len(k) == 0 {
		return nil, fmt.Errorf("handleGetValue but no key in request")
	}

	// prepare the response message
	resp := &pb.Message{
		Type:        pb.Message_GET_VALUE,
		Key:         req.GetKey(),
		CloserPeers: d.closerPeers(ctx, remote, kadt.NewKey(req.GetKey())),
	}

	ns, path, err := record.SplitKey(k) // get namespace (prefix of the key)
	if err != nil || path == "" {
		return nil, fmt.Errorf("invalid key %s: %w", k, err)
	}

	backend, found := d.backends[ns]
	if !found {
		return nil, fmt.Errorf("unsupported record type: %s", ns)
	}

	fetched, err := backend.Fetch(ctx, path)
	if err != nil {
		if errors.Is(err, ds.ErrNotFound) {
			return resp, nil
		}
		return nil, fmt.Errorf("fetch record for key %s: %w", k, err)
	} else if fetched == nil {
		return resp, nil
	}

	rec, ok := fetched.(*recpb.Record)
	if ok {
		resp.Record = rec
		return resp, nil
	}
	// the returned value wasn't a record, which could be the case if the
	// key was prefixed with "providers."

	pset, ok := fetched.(*providerSet)
	if ok {
		resp.ProviderPeers = make([]*pb.Message_Peer, len(pset.providers))
		for i, p := range pset.providers {
			resp.ProviderPeers[i] = pb.FromAddrInfo(p)
		}

		return resp, nil
	}

	return nil, fmt.Errorf("expected *recpb.Record or *providerSet value type, got: %T", pset)
}

// handleAddProvider handles ADD_PROVIDER RPCs from remote peers.
func (d *DHT) handleAddProvider(ctx context.Context, remote peer.ID, req *pb.Message) (*pb.Message, error) {
	k := string(req.GetKey())
	if len(k) > 80 {
		return nil, fmt.Errorf("key size too large")
	} else if len(k) == 0 {
		return nil, fmt.Errorf("key is empty")
	} else if len(req.GetProviderPeers()) == 0 {
		return nil, fmt.Errorf("no provider peers given")
	}

	var addrInfos []peer.AddrInfo
	for _, addrInfo := range req.ProviderAddrInfos() {
		addrInfo := addrInfo // TODO: remove after go.mod was updated to go 1.21

		if addrInfo.ID != remote {
			return nil, fmt.Errorf("attempted to store provider record for other peer %s", addrInfo.ID)
		}

		if len(addrInfo.Addrs) == 0 {
			return nil, fmt.Errorf("no addresses for provider")
		}

		addrInfos = append(addrInfos, addrInfo)
	}

	backend, ok := d.backends[namespaceProviders]
	if !ok {
		return nil, fmt.Errorf("unsupported record type: %s", namespaceProviders)
	}

	for _, addrInfo := range addrInfos {
		if _, err := backend.Store(ctx, k, addrInfo); err != nil {
			return nil, fmt.Errorf("storing provider record: %w", err)
		}
	}

	return nil, nil
}

// handleGetProviders handles GET_PROVIDERS RPCs from remote peers.
func (d *DHT) handleGetProviders(ctx context.Context, remote peer.ID, req *pb.Message) (*pb.Message, error) {
	k := req.GetKey()
	if len(k) > 80 {
		return nil, fmt.Errorf("handleGetProviders key size too large")
	} else if len(k) == 0 {
		return nil, fmt.Errorf("handleGetProviders key is empty")
	}

	backend, ok := d.backends[namespaceProviders]
	if !ok {
		return nil, fmt.Errorf("unsupported record type: %s", namespaceProviders)
	}

	resp := &pb.Message{
		Type:        pb.Message_GET_PROVIDERS,
		Key:         k,
		CloserPeers: d.closerPeers(ctx, remote, kadt.NewKey(k)),
	}

	fetched, err := backend.Fetch(ctx, string(req.GetKey()))
	if err != nil {
		if errors.Is(err, ds.ErrNotFound) {
			return resp, nil
		}
		return nil, fmt.Errorf("fetch providers from datastore: %w", err)
	}

	pset, ok := fetched.(*providerSet)
	if !ok {
		return nil, fmt.Errorf("expected *providerSet value type, got: %T", pset)
	}

	pbProviders := make([]*pb.Message_Peer, len(pset.providers))
	for i, p := range pset.providers {
		pbProviders[i] = pb.FromAddrInfo(p)
	}

	resp.ProviderPeers = pbProviders

	return resp, nil
}

// closerPeers returns the closest peers to the given target key this host knows
// about. It doesn't return 1) itself 2) the peer that asked for closer peers.
func (d *DHT) closerPeers(ctx context.Context, remote peer.ID, target kadt.Key) []*pb.Message_Peer {
	_, span := d.tele.Tracer.Start(ctx, "DHT.closerPeers", otel.WithAttributes(attribute.String("remote", remote.String()), attribute.String("target", target.HexString())))
	defer span.End()

	peers := d.rt.NearestNodes(target, d.cfg.BucketSize)
	if len(peers) == 0 {
		return nil
	}

	// pre-allocated the result set slice.
	filtered := make([]*pb.Message_Peer, 0, len(peers))
	for _, p := range peers {
		pid := peer.ID(p) // TODO: type cast

		// check for own peer ID
		if pid == d.host.ID() {
			d.log.Warn("routing table NearestNodes returned our own ID")
			continue
		}

		// Don't send a peer back themselves
		if pid == remote {
			continue
		}

		// extract peer information from peer store and only add it to the
		// final list if we know any addresses of that peer.
		addrInfo := d.host.Peerstore().PeerInfo(pid)
		if len(addrInfo.Addrs) == 0 {
			continue
		}

		filtered = append(filtered, pb.FromAddrInfo(addrInfo))
	}

	return filtered
}

// Responds to a PIR request in a private FindNode message with a PIR response.
func (d *DHT) handlePrivateFindPeer(ctx context.Context, remote peer.ID, msg *pb.Message) (*pb.Message, error) {
	_, span := d.tele.Tracer.Start(ctx, "DHT.handlePrivateFindPeer", otel.WithAttributes(attribute.String("remote", remote.String())))
	defer span.End()

	pirRequest := msg.GetCloserPeersRequest()
	if pirRequest == nil {
		return nil, fmt.Errorf("PIR Request for CloserPeers not sent in the message")
	}

	bucketsWithAddrInfos, err := d.NormalizeRTJoinedWithPeerStore(kadt.PeerID(remote).Key())
	if err != nil {
		return nil, fmt.Errorf("could not form normalized, joined routing table to run PIR request over")
	}

	pirResponse, err := private_routing.RunPIRforCloserPeersRecords(pirRequest, bucketsWithAddrInfos)
	if err != nil {
		return nil, err
	}
	// println(pirResponse)

	// TODO Ask Gui: handleFindPeer also looks up peerStore directly for the target key and adds it to the closerPeers.
	// This might be necessary as we may not store the node's (KadID, PeerID) if our bucket is full,
	// but we may still record the addresses of the node in the peer store?
	// So do we need to do another PIR over the peer store?
	// Or before we normalize the RT,
	// can we "fill up" our RT with kadID, peerID of records that are in the peerStore but not in the RT?

	response := &pb.Message{
		Type:                pb.Message_PRIVATE_FIND_NODE,
		PIR_Message_ID:      msg.PIR_Message_ID,
		CloserPeersResponse: pirResponse,
	}

	return response, nil
}

// Responds to a PIR request in a private GetProviders message with a PIR response.
func (d *DHT) handlePrivateGetProviderRecords(ctx context.Context, remote peer.ID, msg *pb.Message) (*pb.Message, error) {
	_, span := d.tele.Tracer.Start(ctx, "DHT.handlePrivateGetProviderRecords", otel.WithAttributes(attribute.String("remote", remote.String())))
	defer span.End()

	closerPeersRequest := msg.GetCloserPeersRequest()
	if closerPeersRequest == nil {
		return nil, fmt.Errorf("PIR Request for Closer Peers not sent in the message")
	}

	bucketsWithAddrInfos, err := d.NormalizeRTJoinedWithPeerStore(kadt.PeerID(remote).Key())
	if err != nil {
		return nil, fmt.Errorf("could not form normalized, joined routing table to run PIR request over")
	}

	closerPeersResponse, err := private_routing.RunPIRforCloserPeersRecords(closerPeersRequest, bucketsWithAddrInfos)
	if err != nil {
		return nil, err
	}

	providerPeersRequest := msg.GetProviderPeersRequest()
	if providerPeersRequest == nil {
		return nil, fmt.Errorf("PIR Request for Provider Peers not sent in the message")
	}

	backend, err := typedBackend[*ProvidersBackend](d, namespaceProviders)
	// TODO: This is the bucket index length in the number of bits. It should be hardcoded to 4096.
	bucketIndexLength := 5 * 8
	if err != nil {
		panic("could not typecast backend, to run the function to prepare the DB for PIR")
	}
	mapCIDtoProviderPeers, err := backend.MapCIDBucketsToProviderPeerBytesForPIR(ctx, bucketIndexLength)
	if err != nil {
		return nil, fmt.Errorf("could not construct a map of CIDs to provider peers for PIR,  %s\n", err)
	}

	providerPeersResponse, err := private_routing.RunPIRforProviderPeersRecords(providerPeersRequest, mapCIDtoProviderPeers)
	if err != nil {
		return nil, fmt.Errorf("PIR for provider peers failed, %s\n", err)
	}

	response := &pb.Message{
		Type:                  pb.Message_PRIVATE_GET_PROVIDERS,
		PIR_Message_ID:        msg.PIR_Message_ID,
		CloserPeersResponse:   closerPeersResponse,
		ProviderPeersResponse: providerPeersResponse,
	}

	return response, nil
}

// This function first normalizes the RT --- filling up any buckets that are not full with
// nearest nodes from other buckets, given only the common prefix length for that bucket.
// The (normalized) RT consists of <kad ID, peer ID> records.
// The d.host.Peerstore() consists of <peer ID, peer address> records.
// We then join these key-value stores here, oblivious to the target.
func (d *DHT) NormalizeRTJoinedWithPeerStore(queryingPeerKadId kadt.Key) ([][]byte, error) {
	// Bucket -> [PeerID1, PeerID2, ...]
	bucketsWithPeerIDs := d.rt.NormalizeRT(queryingPeerKadId)

	// Bucket -> <Peer ID -> multiaddress array
	bucketsWithAddrInfos := make([][]byte, len(bucketsWithPeerIDs))

	// Bucket -> <Peer ID and multiaddress array>
	for bid, bucket := range bucketsWithPeerIDs {
		addrInfos := make([]*pb.Message_Peer, len(bucket))
		for i, peerID := range bucket {
			peerInfo := d.host.Peerstore().PeerInfo(peer.ID(peerID))
			messagePeer := pb.FromAddrInfo(peerInfo)
			addrInfos[i] = messagePeer
		}
		mesg := &pb.Message{
			CloserPeers: addrInfos,
		}
		plaintext, err := private_routing.MarshallPBToPlaintext(mesg)
		if err != nil {
			return nil, err
		}
		bucketsWithAddrInfos[bid] = plaintext
	}

	//for i, bucketBytes := range bucketsWithAddrInfos {
	//	println("\nBucket index", i)
	//	for _, b := range bucketBytes {
	//		print(b, ",")
	//	}
	//}
	return bucketsWithAddrInfos, nil
}

// TODO: How to extend this function to provide the functionality in Line 52 in handleFindPeer
//  obliviously to the target key.
// Line 52 in handleFindPeer looks up the peerstore with the target kademlia ID,
// even though closerPeer looks up the peerStore with the output of
// d.rt.NearestNodes(..)
// See the comments there --- the rationale is that the target may be in a bucket
// of the RT that was full, so we don't store its <kad ID, peer ID> in the RT.
// But we still store its <peer ID, multiaddress array> in the peerstore.

// So to do this obliviously to the target, we add some steps to *this function*,
//  which is run before answering the PIR request.
// We can fill up the RT with <kad ID, peer ID> of nodes that *are* in the peerstore but *not* in the RT.
// This effectively adds more entries to the RT; some other than the target which may not have been there earlier.
// So we need to ask Gui if this is acceptable. Let's suppose it is acceptable.
// Then, we would've ensured that every <peer ID, multiaddress []> in the peerstore
// has a corresponding <kad ID, peer ID> in the RT.

// Then we would run NormalizeRT as seen below.
// (Buckets might be very full already, so we might not need to normalize some buckets.)
// Then the join will ensure that *any* target which was earlier in the peerstore,
// but not in the RT, will *also* be included in the join output:
// Join: <target's kadID, target's peerID> <target's peerID, target's address>