diff --git a/src/xdp/program.c b/src/xdp/program.c
index 8c395388..167cc20d 100644
--- a/src/xdp/program.c
+++ b/src/xdp/program.c
@@ -18,6 +18,7 @@ typedef struct _EBPF_PROG_TEST_RUN_CONTEXT {
 } EBPF_PROG_TEST_RUN_CONTEXT;
 
 typedef struct _EBPF_XDP_MD {
+    EBPF_CONTEXT_HEADER;
     xdp_md_t Base;
     EBPF_PROG_TEST_RUN_CONTEXT* ProgTestRunContext;
 } EBPF_XDP_MD;
@@ -120,7 +121,7 @@ XdpCreateContext(
         XdpMd->Base.ingress_ifindex = xdp_context->ingress_ifindex;
     }
 
-    *Context = XdpMd;
+    *Context = &XdpMd->Base;
     XdpMd = NULL;
 
     EbpfResult = EBPF_SUCCESS;
@@ -154,7 +155,7 @@ XdpDeleteContext(
         goto Exit;
     }
 
-    XdpMd = (EBPF_XDP_MD*)Context;
+    XdpMd = CONTAINING_RECORD(Context, EBPF_XDP_MD, Base);
 
     // Copy the packet data to the output buffer.
     if (DataOut != NULL && DataSizeOut != NULL && XdpMd->Base.data != NULL) {
@@ -653,6 +654,7 @@ static const ebpf_program_data_t EbpfXdpProgramData = {
     .context_create = XdpCreateContext,
     .context_destroy = XdpDeleteContext,
     .required_irql = DISPATCH_LEVEL,
+    .capabilities = {.supports_context_header = TRUE},
 };
 
 static const NPI_MODULEID EbpfXdpProgramInfoProviderModuleId = {