proj: Combine test_preload.c and test_wrap.c into intercept.c

proj/server/src/server.py

@@ -2,6 +2,7 @@
 
 from socketserver import UnixStreamServer, StreamRequestHandler, ThreadingMixIn
 import os
+import argparse
 
 
 class Handler(StreamRequestHandler):
@@ -13,12 +14,19 @@ class Handler(StreamRequestHandler):
             msg = self.rfile.readline()
             if not msg:
                 return
-            timestamp, data = msg.split(b' ', 1)
-            if not data.startswith(b'return '):
-                print(data)
-                #self.wfile.write(b'ok\n')
+            timestamp, data = msg.rstrip(b'\n').split(b' ', 1)
+            if not data.startswith(b'return ') and not data == b'return':
+                call = data.decode('utf-8')
+                print(f'[{pid}] {call}')
+                if call.startswith('malloc('):
+                    command = 'fail ENOMEM'
+                else:
+                    command = 'ok'
+                print(f'[{pid}]  -> {command}')
+                self.wfile.write(command.encode('utf-8') + b'\n')
             else:
-                print(data)
+                ret = data.decode('utf-8')
+                print(f'[{pid}]  -> {ret}')
 
 
 class ThreadedUnixStreamServer(ThreadingMixIn, UnixStreamServer):
@@ -26,9 +34,22 @@ class ThreadedUnixStreamServer(ThreadingMixIn, UnixStreamServer):
 
 
 def main() -> None:
-    os.unlink('/tmp/test')
-    with ThreadedUnixStreamServer('/tmp/test', Handler) as server:
-        server.serve_forever()
+    parser = argparse.ArgumentParser()
+    parser.add_argument('socket', metavar='FILE')
+    args = parser.parse_args()
+
+    try:
+        with ThreadedUnixStreamServer(args.socket, Handler) as server:
+            server.serve_forever()
+    except KeyboardInterrupt:
+        print('\nBye')
+        server.shutdown()
+    finally:
+        try:
+            os.unlink(args.socket)
+        except FileNotFoundError:
+            pass
+
 
 if __name__ == '__main__':
     main()