make it importable

ba0f3 · Sep 30, 2018 · 3c5f254 · 3c5f254
1 parent b2fd9a6
commit 3c5f254
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -26,13  26,17 @@ Installation
 git clone --depth=1 https://github.com/ba0f3/proxyproto.nim.git proxyproto
 cd proxyproto
 nim c -d:release src/proxyproto
-
 ```
 
 Usage
 -----
 
-Just add `src/libproxyproto.so` to `LD_PRELOAD`
 Use directly inside a nim program that accept incoming connections
 ```nim
 import proxyproto
 ```
 
 For existing programs, use LD_PRELOAD
 
 ```shell
 $LD_PRELOAD=./src/libproxyproto.so nc -vkl -p 4444

diff --git a/src/proxyproto.nim b/src/proxyproto.nim
@@ -52,8  52,6 @@ proc handshake(fd: SocketHandle, sa: ptr SockAddr, sl: ptr Socklen): int =
     size: int
     hdr: Header
 
-  echo "[PROXY] connection 0x", toHex(fd.int), " handshaking"
-
   while true:
     result = recv(fd, addr hdr, sizeof(hdr), MSG_PEEK)
     if not (result == -1 and errno == EINTR):
@@ -127,7  125,7 @@ var
   RTLD_NEXT {.importc: "RTLD_NEXT", header: "<dlfcn.h>".}: LibHandle
   real_accept: AcceptProc
 
-proc MY_accept(a1: SocketHandle, a2: ptr SockAddr, a3: ptr Socklen): SocketHandle {.exportc:"accept",cdecl.} =
 proc pp_accept*(a1: SocketHandle, a2: ptr SockAddr, a3: ptr Socklen): SocketHandle {.exportc:"accept",cdecl.} =
   result = real_accept(a1, a2, a3)
   if result.int != -1:
     if handshake(result, a2, a3) <= 0:
@@ -136,15  134,12 @@ proc MY_accept(a1: SocketHandle, a2: ptr SockAddr, a3: ptr Socklen): SocketHandl
   return result
 
 
-proc main() =
-  echo "[PROXY] initializing"
-
 proc init() =
   let accept_ptr = symAddr(RTLD_NEXT, "accept")
   if accept_ptr == nil:
     quit "[PROXY] cannot find accept proc"
 
   real_accept = cast[AcceptProc](accept_ptr)
   echo "[PROXY] hook accept OK"
 
-when isMainModule:
-  main()
 init()
diff --git a/tests/server.nim b/tests/server.nim
@@ -0,0  1,29 @@
 import posix, asyncnet, asyncdispatch
 import ../src/proxyproto
 
 proc processClient(client: AsyncSocket) {.async.} =
   while true:
     let line = await client.recvLine()
     await client.send(line & "\c\L")
     var
       slen: SockLen
       src: SockAddr_in
 
     if getpeername(client.getFd(), cast[ptr SockAddr](addr src), addr slen) == 0:
       echo "remote ip ", inet_ntoa(src.sin_addr)
 
 proc serve() {.async.} =
   var server = newAsyncSocket()
   server.bindAddr(Port(6001))
   server.listen()
 
   while true:
     try:
       let client = await server.accept()
       discard processClient(client)
     except:
       discard
 
 
 asyncCheck serve()
 runForever()