libfabintercept.cpp

/*
 * libfab-intercepter — implementation.
 *
 * Interposes the exported setup symbols (fi_getinfo / fi_fabric / fi_fabric2),
 * walks the fabric -> domain -> {endpoint, cq} object graph, and patches each
 * object's ops table in place so selected data-path slots dispatch through
 * consumer hooks. See libfabintercept.h for the model and the contract for
 * adding a new op.
 *
 * Only slots a consumer hooked are overridden; the object handed back to the
 * app is always the provider's real fid. Setup-path work (cloning tables,
 * re-patching on enable) is off the hot path; the only per-call cost on a
 * hooked op is one indirect call through the shim plus the wrap.
 */

/* Required on glibc for RTLD_NEXT; must precede any system header. */
#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#include "libfabintercept.h"

#include <rdma/fi_errno.h>

#include <cstdarg>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <dlfcn.h>
#include <unistd.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace
{

	// The registered hook table (a copy) and whether anything is registered.
	// Until lfi_register() runs, g_registered is false and the interposed
	// entry points forward straight to the provider.
	static struct lfi_hooks g_hooks;
	static bool g_registered = false;

	// Setup-path tracing, off unless LFI_TRACE=1. Never on the hot path.
	static bool g_trace = false;

	static void trace(const char *fmt, ...) __attribute__((format(printf, 1, 2)));
	static void trace(const char *fmt, ...)
	{
		if (!g_trace)
			return;
		va_list ap;
		va_start(ap, fmt);
		fprintf(stderr, "[libfabintercept pid=%d] ", (int)getpid());
		vfprintf(stderr, fmt, ap);
		fputc('\n', stderr);
		fflush(stderr);
		va_end(ap);
	}

	// Which object families a consumer actually asked for. Used to skip
	// patching objects nobody hooked (e.g. don't touch CQs for a send-only
	// consumer).
	static bool want_ep()
	{
		return g_hooks.on_ep_create || g_hooks.on_ep_close ||
		       g_hooks.on_ep_bind || g_hooks.on_ep_info ||
		       g_hooks.rma_read ||
		       g_hooks.rma_readv ||
		       g_hooks.rma_readmsg || g_hooks.rma_write ||
		       g_hooks.rma_writev || g_hooks.rma_writemsg ||
		       g_hooks.rma_inject || g_hooks.rma_writedata ||
		       g_hooks.rma_injectdata || g_hooks.msg_recv ||
		       g_hooks.msg_recvv || g_hooks.msg_recvmsg ||
		       g_hooks.msg_send || g_hooks.msg_sendv ||
		       g_hooks.msg_sendmsg || g_hooks.msg_inject ||
		       g_hooks.msg_senddata || g_hooks.msg_injectdata ||
		       g_hooks.tagged_recv || g_hooks.tagged_recvv ||
		       g_hooks.tagged_recvmsg || g_hooks.tagged_send ||
		       g_hooks.tagged_sendv || g_hooks.tagged_sendmsg ||
		       g_hooks.tagged_inject || g_hooks.tagged_senddata ||
		       g_hooks.tagged_injectdata;
	}

	static bool want_cq()
	{
		return g_hooks.on_cq_create || g_hooks.on_cq_close ||
		       g_hooks.cq_open || g_hooks.cq_read ||
		       g_hooks.cq_readfrom || g_hooks.cq_readerr ||
		       g_hooks.cq_sread || g_hooks.cq_sreadfrom;
	}

	// ---- Generic ops-table patching --------------------------------------
	//
	// Each fid family's ops table begins with a size_t `size`; we clone the
	// provider's real layout into a zeroed buffer at least as large as our
	// own view of the struct, so overriding fields known to our headers is
	// always in bounds even if the provider table is older/newer.

	template <typename T> static T *clone_ops(const T *orig)
	{
		size_t sz = orig->size;
		size_t alloc = sz < sizeof(T) ? sizeof(T) : sz;
		T *copy = static_cast<T *>(calloc(1, alloc));
		if (copy)
			memcpy(copy, orig, sz);
		return copy;
	}

	// True if the provider table advertises the member (offset within its
	// size) and the member is non-null.
#define HAS_OP(orig, type, member)                                             \
	((orig)->size > offsetof(type, member) && (orig)->member != nullptr)

	template <typename T> struct patched_ops {
		const T *orig;
		T *spec;
	};

	static std::vector<patched_ops<struct fi_ops_rma>> g_rma_ops;
	static std::vector<patched_ops<struct fi_ops_msg>> g_msg_ops;
	static std::vector<patched_ops<struct fi_ops_tagged>> g_tagged_ops;
	static std::vector<patched_ops<struct fi_ops_cq>> g_cq_ops;

	template <typename T>
	static T *existing_spec_for(const std::vector<patched_ops<T>> &ops,
				    const T *table)
	{
		for (const auto &p : ops)
			if (p.spec == table || p.orig == table)
				return p.spec;
		return nullptr;
	}

	template <typename T>
	static const T *orig_for_spec(const std::vector<patched_ops<T>> &ops,
				      const T *spec)
	{
		for (const auto &p : ops)
			if (p.spec == spec)
				return p.orig;
		return nullptr;
	}

	static const struct fi_ops_rma *orig_rma_for(const struct fid_ep *ep)
	{
		return ep ? orig_for_spec(g_rma_ops, ep->rma) : nullptr;
	}
	static const struct fi_ops_msg *orig_msg_for(const struct fid_ep *ep)
	{
		return ep ? orig_for_spec(g_msg_ops, ep->msg) : nullptr;
	}
	static const struct fi_ops_tagged *
	orig_tagged_for(const struct fid_ep *ep)
	{
		return ep ? orig_for_spec(g_tagged_ops, ep->tagged) : nullptr;
	}
	static const struct fi_ops_cq *orig_cq_for(const struct fid_cq *cq)
	{
		return cq ? orig_for_spec(g_cq_ops, cq->ops) : nullptr;
	}

	// ---- Shims: forward through the consumer hook, else call provider ----
	//
	// Each shim has the provider op's exact signature. It resolves the
	// provider op for this object, then: if the consumer registered a hook,
	// call it with (args..., provider_op, user); else call the provider op
	// directly. A shim is only installed on a table when its hook is set
	// (see the patched_*_ops installers), so the hook check is belt-and-
	// suspenders, not a per-op gate on an unhooked call.

#define FWD_OR_HOOK(orig, member, hook, ...)                                   \
	do {                                                                   \
		if (!orig || !orig->member)                                    \
			return -FI_EINVAL;                                     \
		if (g_hooks.hook)                                              \
			return g_hooks.hook(__VA_ARGS__, orig->member,         \
					    g_hooks.user);                     \
		return orig->member(__VA_ARGS__);                              \
	} while (0)

	// RMA -----------------------------------------------------------------
	static ssize_t sh_rma_read(struct fid_ep *ep, void *buf, size_t len,
				   void *desc, fi_addr_t src_addr, uint64_t addr,
				   uint64_t key, void *context)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, read, rma_read, ep, buf, len, desc, src_addr,
			    addr, key, context);
	}
	static ssize_t sh_rma_readv(struct fid_ep *ep, const struct iovec *iov,
				    void **desc, size_t count,
				    fi_addr_t src_addr, uint64_t addr,
				    uint64_t key, void *context)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, readv, rma_readv, ep, iov, desc, count, src_addr,
			    addr, key, context);
	}
	static ssize_t sh_rma_readmsg(struct fid_ep *ep,
				      const struct fi_msg_rma *msg,
				      uint64_t flags)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, readmsg, rma_readmsg, ep, msg, flags);
	}
	static ssize_t sh_rma_write(struct fid_ep *ep, const void *buf,
				    size_t len, void *desc, fi_addr_t dest_addr,
				    uint64_t addr, uint64_t key, void *context)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, write, rma_write, ep, buf, len, desc, dest_addr,
			    addr, key, context);
	}
	static ssize_t sh_rma_writev(struct fid_ep *ep, const struct iovec *iov,
				     void **desc, size_t count,
				     fi_addr_t dest_addr, uint64_t addr,
				     uint64_t key, void *context)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, writev, rma_writev, ep, iov, desc, count,
			    dest_addr, addr, key, context);
	}
	static ssize_t sh_rma_writemsg(struct fid_ep *ep,
				       const struct fi_msg_rma *msg,
				       uint64_t flags)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, writemsg, rma_writemsg, ep, msg, flags);
	}
	static ssize_t sh_rma_inject(struct fid_ep *ep, const void *buf,
				     size_t len, fi_addr_t dest_addr,
				     uint64_t addr, uint64_t key)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, inject, rma_inject, ep, buf, len, dest_addr, addr,
			    key);
	}
	static ssize_t sh_rma_writedata(struct fid_ep *ep, const void *buf,
					size_t len, void *desc, uint64_t data,
					fi_addr_t dest_addr, uint64_t addr,
					uint64_t key, void *context)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, writedata, rma_writedata, ep, buf, len, desc,
			    data, dest_addr, addr, key, context);
	}
	static ssize_t sh_rma_injectdata(struct fid_ep *ep, const void *buf,
					 size_t len, uint64_t data,
					 fi_addr_t dest_addr, uint64_t addr,
					 uint64_t key)
	{
		const struct fi_ops_rma *o = orig_rma_for(ep);
		FWD_OR_HOOK(o, injectdata, rma_injectdata, ep, buf, len, data,
			    dest_addr, addr, key);
	}

	// MSG -----------------------------------------------------------------
	static ssize_t sh_msg_recv(struct fid_ep *ep, void *buf, size_t len,
				   void *desc, fi_addr_t src_addr, void *context)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, recv, msg_recv, ep, buf, len, desc, src_addr,
			    context);
	}
	static ssize_t sh_msg_recvv(struct fid_ep *ep, const struct iovec *iov,
				    void **desc, size_t count,
				    fi_addr_t src_addr, void *context)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, recvv, msg_recvv, ep, iov, desc, count, src_addr,
			    context);
	}
	static ssize_t sh_msg_recvmsg(struct fid_ep *ep, const struct fi_msg *msg,
				      uint64_t flags)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, recvmsg, msg_recvmsg, ep, msg, flags);
	}
	static ssize_t sh_msg_send(struct fid_ep *ep, const void *buf, size_t len,
				   void *desc, fi_addr_t dest_addr, void *context)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, send, msg_send, ep, buf, len, desc, dest_addr,
			    context);
	}
	static ssize_t sh_msg_sendv(struct fid_ep *ep, const struct iovec *iov,
				    void **desc, size_t count,
				    fi_addr_t dest_addr, void *context)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, sendv, msg_sendv, ep, iov, desc, count, dest_addr,
			    context);
	}
	static ssize_t sh_msg_sendmsg(struct fid_ep *ep, const struct fi_msg *msg,
				      uint64_t flags)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, sendmsg, msg_sendmsg, ep, msg, flags);
	}
	static ssize_t sh_msg_inject(struct fid_ep *ep, const void *buf,
				     size_t len, fi_addr_t dest_addr)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, inject, msg_inject, ep, buf, len, dest_addr);
	}
	static ssize_t sh_msg_senddata(struct fid_ep *ep, const void *buf,
				       size_t len, void *desc, uint64_t data,
				       fi_addr_t dest_addr, void *context)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, senddata, msg_senddata, ep, buf, len, desc, data,
			    dest_addr, context);
	}
	static ssize_t sh_msg_injectdata(struct fid_ep *ep, const void *buf,
					 size_t len, uint64_t data,
					 fi_addr_t dest_addr)
	{
		const struct fi_ops_msg *o = orig_msg_for(ep);
		FWD_OR_HOOK(o, injectdata, msg_injectdata, ep, buf, len, data,
			    dest_addr);
	}

	// TAGGED --------------------------------------------------------------
	static ssize_t sh_tagged_recv(struct fid_ep *ep, void *buf, size_t len,
				      void *desc, fi_addr_t src_addr,
				      uint64_t tag, uint64_t ignore,
				      void *context)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, recv, tagged_recv, ep, buf, len, desc, src_addr,
			    tag, ignore, context);
	}
	static ssize_t sh_tagged_recvv(struct fid_ep *ep, const struct iovec *iov,
				       void **desc, size_t count,
				       fi_addr_t src_addr, uint64_t tag,
				       uint64_t ignore, void *context)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, recvv, tagged_recvv, ep, iov, desc, count,
			    src_addr, tag, ignore, context);
	}
	static ssize_t sh_tagged_recvmsg(struct fid_ep *ep,
					 const struct fi_msg_tagged *msg,
					 uint64_t flags)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, recvmsg, tagged_recvmsg, ep, msg, flags);
	}
	static ssize_t sh_tagged_send(struct fid_ep *ep, const void *buf,
				      size_t len, void *desc,
				      fi_addr_t dest_addr, uint64_t tag,
				      void *context)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, send, tagged_send, ep, buf, len, desc, dest_addr,
			    tag, context);
	}
	static ssize_t sh_tagged_sendv(struct fid_ep *ep, const struct iovec *iov,
				       void **desc, size_t count,
				       fi_addr_t dest_addr, uint64_t tag,
				       void *context)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, sendv, tagged_sendv, ep, iov, desc, count,
			    dest_addr, tag, context);
	}
	static ssize_t sh_tagged_sendmsg(struct fid_ep *ep,
					 const struct fi_msg_tagged *msg,
					 uint64_t flags)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, sendmsg, tagged_sendmsg, ep, msg, flags);
	}
	static ssize_t sh_tagged_inject(struct fid_ep *ep, const void *buf,
					size_t len, fi_addr_t dest_addr,
					uint64_t tag)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, inject, tagged_inject, ep, buf, len, dest_addr,
			    tag);
	}
	static ssize_t sh_tagged_senddata(struct fid_ep *ep, const void *buf,
					  size_t len, void *desc, uint64_t data,
					  fi_addr_t dest_addr, uint64_t tag,
					  void *context)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, senddata, tagged_senddata, ep, buf, len, desc,
			    data, dest_addr, tag, context);
	}
	static ssize_t sh_tagged_injectdata(struct fid_ep *ep, const void *buf,
					    size_t len, uint64_t data,
					    fi_addr_t dest_addr, uint64_t tag)
	{
		const struct fi_ops_tagged *o = orig_tagged_for(ep);
		FWD_OR_HOOK(o, injectdata, tagged_injectdata, ep, buf, len, data,
			    dest_addr, tag);
	}

	// CQ ------------------------------------------------------------------
	static ssize_t sh_cq_read(struct fid_cq *cq, void *buf, size_t count)
	{
		const struct fi_ops_cq *o = orig_cq_for(cq);
		FWD_OR_HOOK(o, read, cq_read, cq, buf, count);
	}
	static ssize_t sh_cq_readfrom(struct fid_cq *cq, void *buf, size_t count,
				      fi_addr_t *src_addr)
	{
		const struct fi_ops_cq *o = orig_cq_for(cq);
		FWD_OR_HOOK(o, readfrom, cq_readfrom, cq, buf, count, src_addr);
	}
	static ssize_t sh_cq_readerr(struct fid_cq *cq,
				     struct fi_cq_err_entry *buf, uint64_t flags)
	{
		const struct fi_ops_cq *o = orig_cq_for(cq);
		FWD_OR_HOOK(o, readerr, cq_readerr, cq, buf, flags);
	}
	static ssize_t sh_cq_sread(struct fid_cq *cq, void *buf, size_t count,
				   const void *cond, int timeout)
	{
		const struct fi_ops_cq *o = orig_cq_for(cq);
		FWD_OR_HOOK(o, sread, cq_sread, cq, buf, count, cond, timeout);
	}
	static ssize_t sh_cq_sreadfrom(struct fid_cq *cq, void *buf, size_t count,
				       fi_addr_t *src_addr, const void *cond,
				       int timeout)
	{
		const struct fi_ops_cq *o = orig_cq_for(cq);
		FWD_OR_HOOK(o, sreadfrom, cq_sreadfrom, cq, buf, count, src_addr,
			    cond, timeout);
	}

#undef FWD_OR_HOOK

	// ---- ops-table installers --------------------------------------------
	//
	// Clone the provider table once (cached by orig pointer), override only
	// the slots whose hook is registered and whose provider slot exists.

#define INSTALL(orig, spec, type, member, hook, shim)                          \
	do {                                                                   \
		if (g_hooks.hook && HAS_OP(orig, type, member))                \
			(spec)->member = shim;                                 \
	} while (0)

	static struct fi_ops_rma *patched_rma_ops(const struct fi_ops_rma *orig)
	{
		if (!orig)
			return nullptr;
		if (struct fi_ops_rma *s = existing_spec_for(g_rma_ops, orig))
			return s;
		struct fi_ops_rma *spec = clone_ops(orig);
		if (!spec)
			return nullptr;
		INSTALL(orig, spec, struct fi_ops_rma, read, rma_read,
			sh_rma_read);
		INSTALL(orig, spec, struct fi_ops_rma, readv, rma_readv,
			sh_rma_readv);
		INSTALL(orig, spec, struct fi_ops_rma, readmsg, rma_readmsg,
			sh_rma_readmsg);
		INSTALL(orig, spec, struct fi_ops_rma, write, rma_write,
			sh_rma_write);
		INSTALL(orig, spec, struct fi_ops_rma, writev, rma_writev,
			sh_rma_writev);
		INSTALL(orig, spec, struct fi_ops_rma, writemsg, rma_writemsg,
			sh_rma_writemsg);
		INSTALL(orig, spec, struct fi_ops_rma, inject, rma_inject,
			sh_rma_inject);
		INSTALL(orig, spec, struct fi_ops_rma, writedata, rma_writedata,
			sh_rma_writedata);
		INSTALL(orig, spec, struct fi_ops_rma, injectdata,
			rma_injectdata, sh_rma_injectdata);
		g_rma_ops.push_back({orig, spec});
		trace("installed rma shims (orig=%p spec=%p)", (void *)orig,
		      (void *)spec);
		return spec;
	}

	static struct fi_ops_msg *patched_msg_ops(const struct fi_ops_msg *orig)
	{
		if (!orig)
			return nullptr;
		if (struct fi_ops_msg *s = existing_spec_for(g_msg_ops, orig))
			return s;
		struct fi_ops_msg *spec = clone_ops(orig);
		if (!spec)
			return nullptr;
		INSTALL(orig, spec, struct fi_ops_msg, recv, msg_recv,
			sh_msg_recv);
		INSTALL(orig, spec, struct fi_ops_msg, recvv, msg_recvv,
			sh_msg_recvv);
		INSTALL(orig, spec, struct fi_ops_msg, recvmsg, msg_recvmsg,
			sh_msg_recvmsg);
		INSTALL(orig, spec, struct fi_ops_msg, send, msg_send,
			sh_msg_send);
		INSTALL(orig, spec, struct fi_ops_msg, sendv, msg_sendv,
			sh_msg_sendv);
		INSTALL(orig, spec, struct fi_ops_msg, sendmsg, msg_sendmsg,
			sh_msg_sendmsg);
		INSTALL(orig, spec, struct fi_ops_msg, inject, msg_inject,
			sh_msg_inject);
		INSTALL(orig, spec, struct fi_ops_msg, senddata, msg_senddata,
			sh_msg_senddata);
		INSTALL(orig, spec, struct fi_ops_msg, injectdata,
			msg_injectdata, sh_msg_injectdata);
		g_msg_ops.push_back({orig, spec});
		trace("installed msg shims (orig=%p spec=%p)", (void *)orig,
		      (void *)spec);
		return spec;
	}

	static struct fi_ops_tagged *
	patched_tagged_ops(const struct fi_ops_tagged *orig)
	{
		if (!orig)
			return nullptr;
		if (struct fi_ops_tagged *s =
			    existing_spec_for(g_tagged_ops, orig))
			return s;
		struct fi_ops_tagged *spec = clone_ops(orig);
		if (!spec)
			return nullptr;
		INSTALL(orig, spec, struct fi_ops_tagged, recv, tagged_recv,
			sh_tagged_recv);
		INSTALL(orig, spec, struct fi_ops_tagged, recvv, tagged_recvv,
			sh_tagged_recvv);
		INSTALL(orig, spec, struct fi_ops_tagged, recvmsg,
			tagged_recvmsg, sh_tagged_recvmsg);
		INSTALL(orig, spec, struct fi_ops_tagged, send, tagged_send,
			sh_tagged_send);
		INSTALL(orig, spec, struct fi_ops_tagged, sendv, tagged_sendv,
			sh_tagged_sendv);
		INSTALL(orig, spec, struct fi_ops_tagged, sendmsg,
			tagged_sendmsg, sh_tagged_sendmsg);
		INSTALL(orig, spec, struct fi_ops_tagged, inject, tagged_inject,
			sh_tagged_inject);
		INSTALL(orig, spec, struct fi_ops_tagged, senddata,
			tagged_senddata, sh_tagged_senddata);
		INSTALL(orig, spec, struct fi_ops_tagged, injectdata,
			tagged_injectdata, sh_tagged_injectdata);
		g_tagged_ops.push_back({orig, spec});
		trace("installed tagged shims (orig=%p spec=%p)", (void *)orig,
		      (void *)spec);
		return spec;
	}

	static struct fi_ops_cq *patched_cq_ops(const struct fi_ops_cq *orig)
	{
		if (!orig)
			return nullptr;
		if (struct fi_ops_cq *s = existing_spec_for(g_cq_ops, orig))
			return s;
		struct fi_ops_cq *spec = clone_ops(orig);
		if (!spec)
			return nullptr;
		INSTALL(orig, spec, struct fi_ops_cq, read, cq_read, sh_cq_read);
		INSTALL(orig, spec, struct fi_ops_cq, readfrom, cq_readfrom,
			sh_cq_readfrom);
		INSTALL(orig, spec, struct fi_ops_cq, readerr, cq_readerr,
			sh_cq_readerr);
		INSTALL(orig, spec, struct fi_ops_cq, sread, cq_sread,
			sh_cq_sread);
		INSTALL(orig, spec, struct fi_ops_cq, sreadfrom, cq_sreadfrom,
			sh_cq_sreadfrom);
		g_cq_ops.push_back({orig, spec});
		trace("installed cq shims (orig=%p spec=%p)", (void *)orig,
		      (void *)spec);
		return spec;
	}

#undef INSTALL

	// ---- Endpoint patching -----------------------------------------------
	//
	// The data-op tables may be null at fi_endpoint() time and only get
	// populated by the provider at fi_enable(); re-apply the swaps after a
	// successful bind/control/setopt so we don't miss them. fid.ops and the
	// fi_ops_ep table are assumed shared across endpoints of a provider, so
	// they are cloned once into single globals.

	static const struct fi_ops *g_orig_ep_fid_ops = nullptr;
	static const struct fi_ops_ep *g_orig_ep_ops = nullptr;
	static struct fi_ops *g_spec_ep_fid_ops = nullptr;
	static struct fi_ops_ep *g_spec_ep_ops = nullptr;

	static void patch_endpoint(struct fid_ep *ep);

	static int ep_fid_bind(struct fid *fid, struct fid *bfid, uint64_t flags)
	{
		int ret = g_orig_ep_fid_ops->bind(fid, bfid, flags);
		if (!ret) {
			struct fid_ep *ep =
				reinterpret_cast<struct fid_ep *>(fid);
			if (g_hooks.on_ep_bind)
				g_hooks.on_ep_bind(ep, bfid, flags,
						   g_hooks.user);
			patch_endpoint(ep);
		}
		return ret;
	}

	static int ep_fid_control(struct fid *fid, int command, void *arg)
	{
		int ret = g_orig_ep_fid_ops->control(fid, command, arg);
		if (!ret)
			patch_endpoint(reinterpret_cast<struct fid_ep *>(fid));
		return ret;
	}

	// Notify the consumer before the provider tears the endpoint down, while
	// it (and its domain) are still valid, then close for real.
	static int ep_fid_close(struct fid *fid)
	{
		if (g_hooks.on_ep_close)
			g_hooks.on_ep_close(
				reinterpret_cast<struct fid_ep *>(fid),
				g_hooks.user);
		return g_orig_ep_fid_ops->close(fid);
	}

	static int ep_setopt(fid_t fid, int level, int optname,
			     const void *optval, size_t optlen)
	{
		int ret = g_orig_ep_ops->setopt(fid, level, optname, optval,
						optlen);
		if (!ret)
			patch_endpoint(reinterpret_cast<struct fid_ep *>(fid));
		return ret;
	}

	// ep -> domain, so on_ep_create can hand the consumer the owning domain.
	static std::unordered_map<const void *, struct fid_domain *> g_ep_domain;
	static std::unordered_set<const void *> g_ep_announced;

	static void announce_ep(struct fid_ep *ep, struct fid_domain *dom)
	{
		if (ep && dom)
			g_ep_domain[ep] = dom;
		if (!g_hooks.on_ep_create || !ep)
			return;
		if (!g_ep_announced.insert(ep).second)
			return;  // once per ep
		auto it = g_ep_domain.find(ep);
		g_hooks.on_ep_create(
			ep, it != g_ep_domain.end() ? it->second : nullptr,
			g_hooks.user);
	}

	static int ep_tx_ctx(struct fid_ep *sep, int index,
			     struct fi_tx_attr *attr, struct fid_ep **tx_ep,
			     void *context)
	{
		int ret = g_orig_ep_ops->tx_ctx(sep, index, attr, tx_ep,
						context);
		if (!ret && tx_ep && *tx_ep) {
			patch_endpoint(*tx_ep);
			auto it = g_ep_domain.find(sep);
			announce_ep(*tx_ep, it != g_ep_domain.end() ? it->second
								    : nullptr);
		}
		return ret;
	}

	static int ep_rx_ctx(struct fid_ep *sep, int index,
			     struct fi_rx_attr *attr, struct fid_ep **rx_ep,
			     void *context)
	{
		int ret = g_orig_ep_ops->rx_ctx(sep, index, attr, rx_ep,
						context);
		if (!ret && rx_ep && *rx_ep) {
			patch_endpoint(*rx_ep);
			auto it = g_ep_domain.find(sep);
			announce_ep(*rx_ep, it != g_ep_domain.end() ? it->second
								    : nullptr);
		}
		return ret;
	}

	static void patch_endpoint(struct fid_ep *ep)
	{
		trace("patch_endpoint ep=%p rma=%p msg=%p tagged=%p", (void *)ep,
		      (void *)ep->rma, (void *)ep->msg, (void *)ep->tagged);
		if (ep->rma) {
			struct fi_ops_rma *s = patched_rma_ops(ep->rma);
			if (s)
				ep->rma = s;
		}
		if (ep->msg) {
			struct fi_ops_msg *s = patched_msg_ops(ep->msg);
			if (s)
				ep->msg = s;
		}
		if (ep->tagged) {
			struct fi_ops_tagged *s = patched_tagged_ops(ep->tagged);
			if (s)
				ep->tagged = s;
		}

		if (ep->fid.ops) {
			if (!g_spec_ep_fid_ops) {
				g_orig_ep_fid_ops = ep->fid.ops;
				g_spec_ep_fid_ops = clone_ops(ep->fid.ops);
				if (g_spec_ep_fid_ops) {
					if (HAS_OP(g_orig_ep_fid_ops,
						   struct fi_ops, bind))
						g_spec_ep_fid_ops->bind =
							ep_fid_bind;
					if (HAS_OP(g_orig_ep_fid_ops,
						   struct fi_ops, close))
						g_spec_ep_fid_ops->close =
							ep_fid_close;
					g_spec_ep_fid_ops->control =
						ep_fid_control;
				}
			}
			if (g_spec_ep_fid_ops &&
			    ep->fid.ops == g_orig_ep_fid_ops)
				ep->fid.ops = g_spec_ep_fid_ops;
		}

		if (ep->ops) {
			if (!g_spec_ep_ops) {
				g_orig_ep_ops = ep->ops;
				g_spec_ep_ops = clone_ops(ep->ops);
				if (g_spec_ep_ops) {
					if (HAS_OP(g_orig_ep_ops,
						   struct fi_ops_ep, setopt))
						g_spec_ep_ops->setopt =
							ep_setopt;
					if (HAS_OP(g_orig_ep_ops,
						   struct fi_ops_ep, tx_ctx))
						g_spec_ep_ops->tx_ctx =
							ep_tx_ctx;
					if (HAS_OP(g_orig_ep_ops,
						   struct fi_ops_ep, rx_ctx))
						g_spec_ep_ops->rx_ctx =
							ep_rx_ctx;
				}
			}
			if (g_spec_ep_ops && ep->ops == g_orig_ep_ops)
				ep->ops = g_spec_ep_ops;
		}
	}

	// ---- CQ patching -----------------------------------------------------

	static const struct fi_ops *g_orig_cq_fid_ops = nullptr;
	static struct fi_ops *g_spec_cq_fid_ops = nullptr;

	static int cq_fid_close(struct fid *fid)
	{
		if (g_hooks.on_cq_close)
			g_hooks.on_cq_close(
				reinterpret_cast<struct fid_cq *>(fid),
				g_hooks.user);
		return g_orig_cq_fid_ops->close(fid);
	}

	static void patch_cq(struct fid_cq *cq)
	{
		trace("patch_cq cq=%p ops=%p", (void *)cq, (void *)cq->ops);
		if (cq->ops) {
			struct fi_ops_cq *s = patched_cq_ops(cq->ops);
			if (s)
				cq->ops = s;
		}
		if (cq->fid.ops) {
			if (!g_spec_cq_fid_ops) {
				g_orig_cq_fid_ops = cq->fid.ops;
				g_spec_cq_fid_ops = clone_ops(cq->fid.ops);
				if (g_spec_cq_fid_ops &&
				    HAS_OP(g_orig_cq_fid_ops, struct fi_ops,
					   close))
					g_spec_cq_fid_ops->close = cq_fid_close;
			}
			if (g_spec_cq_fid_ops &&
			    cq->fid.ops == g_orig_cq_fid_ops)
				cq->fid.ops = g_spec_cq_fid_ops;
		}
	}

	// ---- Domain patching: intercept endpoint / cq creation ---------------

	static const struct fi_ops_domain *g_orig_domain_ops = nullptr;
	static struct fi_ops_domain *g_spec_domain_ops = nullptr;

	static int domain_endpoint(struct fid_domain *domain,
				   struct fi_info *info, struct fid_ep **ep,
				   void *context)
	{
		int ret = g_orig_domain_ops->endpoint(domain, info, ep,
						       context);
		if (!ret && ep && *ep) {
			patch_endpoint(*ep);
			announce_ep(*ep, domain);
			if (g_hooks.on_ep_info)
				g_hooks.on_ep_info(*ep, info, g_hooks.user);
		}
		return ret;
	}

	static int domain_endpoint2(struct fid_domain *domain,
				    struct fi_info *info, struct fid_ep **ep,
				    uint64_t flags, void *context)
	{
		int ret = g_orig_domain_ops->endpoint2(domain, info, ep, flags,
						       context);
		if (!ret && ep && *ep) {
			patch_endpoint(*ep);
			announce_ep(*ep, domain);
			if (g_hooks.on_ep_info)
				g_hooks.on_ep_info(*ep, info, g_hooks.user);
		}
		return ret;
	}

	static int domain_scalable_ep(struct fid_domain *domain,
				      struct fi_info *info, struct fid_ep **sep,
				      void *context)
	{
		int ret = g_orig_domain_ops->scalable_ep(domain, info, sep,
							 context);
		if (!ret && sep && *sep) {
			patch_endpoint(*sep);
			announce_ep(*sep, domain);
			if (g_hooks.on_ep_info)
				g_hooks.on_ep_info(*sep, info, g_hooks.user);
		}
		return ret;
	}

	static int domain_srx_ctx(struct fid_domain *domain,
				  struct fi_rx_attr *attr,
				  struct fid_ep **rx_ep, void *context)
	{
		int ret = g_orig_domain_ops->srx_ctx(domain, attr, rx_ep,
						     context);
		if (!ret && rx_ep && *rx_ep) {
			patch_endpoint(*rx_ep);
			announce_ep(*rx_ep, domain);
		}
		return ret;
	}

	// Real cq_open plus wrap: create the provider CQ, then patch it and fire
	// on_cq_create. This is what the cq_open hook receives as `next`, so a
	// hook that forwards still gets the CQ wrapped and the lifecycle fired;
	// a hook that returns without calling it suppresses the open and nothing
	// is patched (no CQ exists).
	static int cq_open_forward(struct fid_domain *domain,
				   struct fi_cq_attr *attr, struct fid_cq **cq,
				   void *context)
	{
		int ret = g_orig_domain_ops->cq_open(domain, attr, cq, context);
		if (!ret && cq && *cq) {
			patch_cq(*cq);
			if (g_hooks.on_cq_create)
				g_hooks.on_cq_create(*cq, domain, g_hooks.user);
		}
		return ret;
	}

	static int domain_cq_open(struct fid_domain *domain,
				  struct fi_cq_attr *attr, struct fid_cq **cq,
				  void *context)
	{
		if (g_hooks.cq_open)
			return g_hooks.cq_open(domain, attr, cq, context,
					       cq_open_forward, g_hooks.user);
		return cq_open_forward(domain, attr, cq, context);
	}

	static void patch_domain(struct fid_domain *dom)
	{
		trace("patch_domain dom=%p ops=%p", (void *)dom,
		      (void *)(dom ? dom->ops : nullptr));
		if (!dom->ops)
			return;
		if (g_spec_domain_ops && dom->ops == g_spec_domain_ops)
			return;
		if (!g_orig_domain_ops) {
			g_orig_domain_ops = dom->ops;
			g_spec_domain_ops = clone_ops(dom->ops);
			if (g_spec_domain_ops) {
				const struct fi_ops_domain *o = g_orig_domain_ops;
				if (want_ep()) {
					if (HAS_OP(o, struct fi_ops_domain,
						   endpoint))
						g_spec_domain_ops->endpoint =
							domain_endpoint;
					if (HAS_OP(o, struct fi_ops_domain,
						   endpoint2))
						g_spec_domain_ops->endpoint2 =
							domain_endpoint2;
					if (HAS_OP(o, struct fi_ops_domain,
						   scalable_ep))
						g_spec_domain_ops->scalable_ep =
							domain_scalable_ep;
					if (HAS_OP(o, struct fi_ops_domain,
						   srx_ctx))
						g_spec_domain_ops->srx_ctx =
							domain_srx_ctx;
				}
				if (want_cq() &&
				    HAS_OP(o, struct fi_ops_domain, cq_open))
					g_spec_domain_ops->cq_open =
						domain_cq_open;
			}
		}
		if (g_spec_domain_ops && dom->ops == g_orig_domain_ops)
			dom->ops = g_spec_domain_ops;
	}

	// ---- Fabric patching: intercept domain creation ----------------------

	static const struct fi_ops_fabric *g_orig_fabric_ops = nullptr;
	static struct fi_ops_fabric *g_spec_fabric_ops = nullptr;

	static int fabric_domain(struct fid_fabric *fabric, struct fi_info *info,
				 struct fid_domain **domain, void *context)
	{
		int ret = g_orig_fabric_ops->domain(fabric, info, domain,
						    context);
		if (!ret && domain && *domain) {
			patch_domain(*domain);
			if (g_hooks.on_domain_info)
				g_hooks.on_domain_info(*domain, info,
						       g_hooks.user);
		}
		return ret;
	}

	static int fabric_domain2(struct fid_fabric *fabric,
				  struct fi_info *info,
				  struct fid_domain **domain, uint64_t flags,
				  void *context)
	{
		int ret = g_orig_fabric_ops->domain2(fabric, info, domain, flags,
						     context);
		if (!ret && domain && *domain) {
			patch_domain(*domain);
			if (g_hooks.on_domain_info)
				g_hooks.on_domain_info(*domain, info,
						       g_hooks.user);
		}
		return ret;
	}

	static void patch_fabric(struct fid_fabric *fab)
	{
		trace("patch_fabric fab=%p ops=%p", (void *)fab,
		      (void *)(fab ? fab->ops : nullptr));
		if (!fab->ops)
			return;
		if (g_spec_fabric_ops && fab->ops == g_spec_fabric_ops)
			return;
		if (!g_orig_fabric_ops) {
			g_orig_fabric_ops = fab->ops;
			g_spec_fabric_ops = clone_ops(fab->ops);
			if (g_spec_fabric_ops) {
				const struct fi_ops_fabric *o = g_orig_fabric_ops;
				if (HAS_OP(o, struct fi_ops_fabric, domain))
					g_spec_fabric_ops->domain =
						fabric_domain;
				if (HAS_OP(o, struct fi_ops_fabric, domain2))
					g_spec_fabric_ops->domain2 =
						fabric_domain2;
			}
		}
		if (g_spec_fabric_ops && fab->ops == g_orig_fabric_ops)
			fab->ops = g_spec_fabric_ops;
	}

#undef HAS_OP

	// ---- Resolve the real exported setup symbols -------------------------

	using fi_fabric_fn = int (*)(struct fi_fabric_attr *,
				     struct fid_fabric **, void *);
	using fi_fabric2_fn = int (*)(struct fi_info *, struct fid_fabric **,
				      uint64_t, void *);

	// Resolve a real exported libfabric symbol behind our interposer.
	//
	// dlsym(RTLD_NEXT, ...) only searches the *global* scope after this
	// object. That suffices when the app links libfabric directly (it's in
	// the global scope). But some apps reach libfabric through a library
	// dlopen'd with RTLD_LOCAL -- e.g. MPI via a Python C-extension
	// (mpi4py/torch) -- which loads libfabric into a private namespace that
	// is *not* on the global search list. Our interposed symbol is still
	// what the app binds to (LD_PRELOAD symbols are global and win), but
	// RTLD_NEXT then sees no provider and we'd wrongly return -FI_ENOSYS.
	//
	// Fall back to dlopen'ing libfabric by soname with RTLD_GLOBAL: NOLOAD
	// first, which (when it is already loaded privately) returns the app's
	// existing map and promotes it into the global scope -- so we use the
	// very same libfabric instance, no second copy. If it isn't loaded yet,
	// load it. This is setup-path only, never the hot path.
	static void *libfabric_handle()
	{
		static void *h = []() -> void * {
			void *p = dlopen("libfabric.so.1",
					 RTLD_NOW | RTLD_GLOBAL | RTLD_NOLOAD);
			if (!p)
				p = dlopen("libfabric.so.1",
					   RTLD_NOW | RTLD_GLOBAL);
			if (!p)
				trace("dlopen libfabric.so.1 failed: %s",
				      dlerror());
			return p;
		}();
		return h;
	}

	static void *resolve_real(const char *name)
	{
		if (void *p = dlsym(RTLD_NEXT, name))
			return p;
		void *h = libfabric_handle();
		return h ? dlsym(h, name) : nullptr;
	}

	static lfi_getinfo_fn real_fi_getinfo()
	{
		static auto fn = reinterpret_cast<lfi_getinfo_fn>(
			resolve_real("fi_getinfo"));
		return fn;
	}
	static fi_fabric_fn real_fi_fabric()
	{
		static auto fn = reinterpret_cast<fi_fabric_fn>(
			resolve_real("fi_fabric"));
		return fn;
	}
	static fi_fabric2_fn real_fi_fabric2()
	{
		static auto fn = reinterpret_cast<fi_fabric2_fn>(
			resolve_real("fi_fabric2"));
		return fn;
	}

	struct lfi_env {
		lfi_env()
		{
			const char *t = getenv("LFI_TRACE");
			g_trace = t && atoi(t) != 0;
		}
	};
	static lfi_env g_lfi_env;

} // namespace

// ---- Public registration ---------------------------------------------------

extern "C" int lfi_register(const struct lfi_hooks *hooks)
{
	if (!hooks)
		return -FI_EINVAL;
	g_hooks = *hooks;
	g_registered = true;
	trace("registered hooks (want_ep=%d want_cq=%d)", (int)want_ep(),
	      (int)want_cq());
	return 0;
}

// ---- Interposed setup symbols -----------------------------------------------

extern "C" int fi_getinfo(uint32_t version, const char *node,
			  const char *service, uint64_t flags,
			  const struct fi_info *hints, struct fi_info **info)
{
	auto real_fn = real_fi_getinfo();
	if (!real_fn)
		return -FI_ENOSYS;
	if (g_registered && g_hooks.getinfo)
		return g_hooks.getinfo(version, node, service, flags, hints,
				       info, real_fn, g_hooks.user);
	return real_fn(version, node, service, flags, hints, info);
}

extern "C" int fi_fabric(struct fi_fabric_attr *attr,
			 struct fid_fabric **fabric, void *context)
{
	auto real_fn = real_fi_fabric();
	if (!real_fn)
		return -FI_ENOSYS;
	int ret = real_fn(attr, fabric, context);
	if (!ret && g_registered && fabric && *fabric)
		patch_fabric(*fabric);
	return ret;
}

extern "C" int fi_fabric2(struct fi_info *info, struct fid_fabric **fabric,
			  uint64_t flags, void *context)
{
	auto real_fn = real_fi_fabric2();
	if (!real_fn)
		return -FI_ENOSYS;
	int ret = real_fn(info, fabric, flags, context);
	if (!ret && g_registered && fabric && *fabric)
		patch_fabric(*fabric);
	return ret;
}