How to write a minimal-overhead proxy to localhost

Question

Update: question now contains the final edited answer!

I now use the following (final answer):

module Main where

import Control.Concurrent        (forkIO)
import Control.Monad             (when,forever,void)
import Network                   (PortID(PortNumber),listenOn)
import Network.Socket hiding     (listen,recv,send)
import Network.Socket.ByteString (recv,sendAll)
import qualified Data.ByteString as B
import System

type Host = String
type Port = PortNumber

main :: IO ()
main = do
  [lp,h,p] <- getArgs  
  start (port lp) h (port p)
  where
    port = fromInteger . read

start :: Port -> Host -> Port -> IO ()
start lp rh rp = withSocketsDo $ do
  proxy <- listenOn $ PortNumber lp
  forever $ do
    (client,_) <- accept proxy
    void . forkIO $ (client >-<) =<< rh .@. rp

(.@.) :: Host -> Port -> IO Socket
host .@. port = do
  addr:_ <- getAddrInfo Nothing (Just host) (Just $ show port)
  server <- socket (addrFamily  addr) Stream defaultProtocol
  connect server   (addrAddress addr)
  return  server

(>-<) :: Socket -> Socket -> IO ()
x >-< y = do x >- y; y >- x

(>-) :: Socket -> Socket -> IO ()
s >- r = void . forkIO . handle $ forever stream
  where
    stream = recv s (64 * 1024) >>= ifNot0 >>= sendAll r
    ifNot0  = \c -> do when (B.null c) $ handle (error "0"); return c
    handle = flip catch $ \e -> print e >> sClose s >> sClose r

which can be run like this:

proxy 2000 localhost 3389

Using mRemote, if I connect to localhost:2000, I do see the login screen of the local machine! :)

*If I find a way to improve (>-) even further, I will update this answer!

Answer 1

It seems that you came to this tcp proxy gist when looking for information. At this time, is is broken and a bit messy. In such case, please don't hesitate to ping the author (in this case, me) so that he can fix the gist for future references :)

I'll fix it ASAP and link to this SO question. The fixed version will include sendAll as well as all nice suggestions coming from this SO question, so please share your best thoughts. As a side note, this branch of throttle already had the sendAll fix, in case of interest.

EDIT : the gist is fixed now

Answer 2

Found this gist a few months ago when I was getting started with Haskell.

It's really simple and easy to understand.

EDIT: Based on the gist above, here is a tested RDP proxy. Difference is replacing send with sendAll to make sure all data is delivered. Found this problem when testing through the linux rdp server (large payload disconnects).

module Main where

import Control.Concurrent      (forkIO)
import Control.Monad           (forever,unless)
import Network                 (PortID(PortNumber),listenOn)
import qualified Data.ByteString as S
import Network.Socket hiding (listen,recv,send)
import Network.Socket.ByteString (recv,sendAll)
import System.Posix            (Handler(Ignore),installHandler,sigPIPE)


localPort :: PortNumber
localPort = 3390

remoteHost :: String
remoteHost = "localhost"

remotePort :: Integer
remotePort = 3389

main :: IO ()
main = do
  ignore $ installHandler sigPIPE Ignore Nothing
  start

start :: IO ()
start = withSocketsDo $ do
  listener <- listenOn $ PortNumber localPort
  forever $ do
    (client,_) <- accept listener
    ignore $ forkIO $ do
      server <- connectToServer
      client `proxyTo` server
      server `proxyTo` client
    return ()
  where
    connectToServer = do
      addrinfos <- getAddrInfo Nothing (Just remoteHost) (Just $ show remotePort)
      let serveraddr = head addrinfos
      server <- socket (addrFamily serveraddr) Stream defaultProtocol
      connect server (addrAddress serveraddr)
      return server
    proxyTo from to = do
      ignore $ forkIO $ flip catch (close from to) $ forever $ do
        content <- recv from 1024
        unless (S.null content) $ sendAll to content
      return ()
    close a b _ = do
      sClose a
      sClose b

-- | Run an action and ignore the result.
ignore :: Monad m => m a -> m ()
ignore m = m >> return ()