improve plugin

test/integration/targets/setup_proxy/files/hamsandwich.py

@@ -1,12 +1,31 @@
 from __future__ import annotations
 
 from proxy.http.proxy import HttpProxyBasePlugin
+from proxy.http.parser import HttpParser, httpParserStates
 
 
 class HamSandwichPlugin(HttpProxyBasePlugin):
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.parser = None
+        self.parsable = True
+
     def handle_upstream_chunk(self, chunk):
-        headers, sep, body = chunk.tobytes().partition(b'\r\n\r\n')
-        if not sep:
+        if not self.parsable:
+            return chunk
+
+        if not self.parser:
+            self.parser = HttpParser.response(chunk)
+        else:
+            self.parser.parse(chunk)
+
+        if self.parser.state == httpParserStates.INITIALIZED:
+            self.parsable = False
             return chunk
-        return memoryview(bytearray(headers + b'\r\nX-Sandwich: ham' + sep + body))
+
+        if not self.parser.is_complete:
+            return None
+
+        self.parser.add_header(b'X-Sandwich', b'ham')
+        return self.parser.build_response()