Fix CI, make COPY_META mode work

oremanj · oremanj · commit 305b258a482a · 2022-01-13T17:19:28.000-07:00
diff --git a/netfilterqueue.pxd b/netfilterqueue.pxd
@@ -169,6 +169,7 @@ cdef enum:
     NF_MAX_VERDICT = NF_STOP
 
 cdef class NetfilterQueue:
+    cdef object __weakref__
     cdef object user_callback # User callback
     cdef nfq_handle *h # Handle to NFQueue library
     cdef nfq_q_handle *qh # A handle to the queue
diff --git a/netfilterqueue.pyx b/netfilterqueue.pyx
@@ -82,7 +82,9 @@ cdef class Packet:
 
         self.payload_len = nfq_get_payload(nfa, &self.payload)
         if self.payload_len < 0:
-            raise OSError("Failed to get payload of packet.")
+            # Probably using a mode that doesn't provide the payload
+            self.payload = NULL
+            self.payload_len = 0
 
         nfq_get_timestamp(nfa, &self.timestamp)
         self.mark = nfq_get_nfmark(nfa)
@@ -142,6 +144,10 @@ cdef class Packet:
             return self._owned_payload
         elif self.payload != NULL:
             return self.payload[:self.payload_len]
+        elif self.payload_len == 0:
+            raise RuntimeError(
+                "Packet has no payload -- perhaps you're using COPY_META mode?"
+            )
         else:
             raise RuntimeError(
                 "Payload data is no longer available. You must call "
@@ -191,7 +197,6 @@ cdef class NetfilterQueue:
         cdef u_int16_t af # Address family
         af = kwargs.get("af", PF_INET)
 
-        self.unbinding = False
         self.h = nfq_open()
         if self.h == NULL:
             raise OSError("Failed to open NFQueue.")
diff --git a/setup.py b/setup.py
@@ -3,11 +3,11 @@
 
 VERSION = "0.9.0"  # Remember to change CHANGES.txt and netfilterqueue.pyx when version changes.
 
+setup_requires = []
 try:
     # Use Cython
     from Cython.Build import cythonize
 
-    setup_requires = []
     ext_modules = cythonize(
         Extension(
             "netfilterqueue", ["netfilterqueue.pyx"], libraries=["netfilter_queue"]
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -191,7 +191,8 @@ async def run(self):
             # Tell each peer about the other one's port
             for idx in (1, 2):
                 self.dest_addr[idx] = (
-                    PEER_IP[idx], int(await self._received[idx].receive())
+                    PEER_IP[idx],
+                    int(await self._received[idx].receive()),
                 )
                 await self._conn[3 - idx].send(b"%d" % self.dest_addr[idx][1])
             yield
diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -9,7 +9,7 @@
 import time
 import weakref
 
-from netfilterqueue import NetfilterQueue
+from netfilterqueue import NetfilterQueue, COPY_META
 
 
 async def test_comms_without_queue(harness):
@@ -94,6 +94,8 @@ async def test_mark_repeat(harness):
 
     def cb(chan, pkt):
         nonlocal counter
+        with pytest.raises(RuntimeError, match="Packet has no payload"):
+            pkt.get_payload()
         assert pkt.get_mark() == counter
         timestamps.append(pkt.get_timestamp())
         if counter < 5:
@@ -104,7 +106,7 @@ def cb(chan, pkt):
         else:
             pkt.accept()
 
-    async with harness.capture_packets_to(2, cb):
+    async with harness.capture_packets_to(2, cb, mode=COPY_META):
         t0 = time.time()
         await harness.send(2, b"testing")
         await harness.expect(2, b"testing")
@@ -155,7 +157,7 @@ async def listen_for_packets():
         (mac1, FORWARD, b"one"),
         (mac1, FORWARD, b"two"),
         (None, OUTPUT, b"three"),
-        (None, OUTPUT, b"four")
+        (None, OUTPUT, b"four"),
     ]
 
 
