socks5.hs

import Data.Attoparsec as A
import Data.Attoparsec.Binary as A
import Data.Attoparsec.Combinator as A

import Data.ByteString as B (pack, ByteString)

import Control.Concurrent.Async

import Control.Proxy as P
import Control.Proxy.TCP as P
import Control.Proxy.Parse as P
import Control.Proxy.Attoparsec as P
import Control.Proxy.Trans.Either as P
import Control.Proxy.Trans.State as P

import Control.Monad
import Data.Monoid
import Data.Word (Word8)

sockVer = word8 5
reserved = word8 0

{--
        +----+----------+----------+
        |VER | NMETHODS | METHODS  |
        +----+----------+----------+
        | 1  |    1     | 1 to 255 |
        +----+----------+----------+
        Methods:
            0 -> NONE
            1 -> GSSAPI
            2 -> user/pass
--}

authMsg = do
    sockVer
    num <- anyWord8
    choice . replicate (fromIntegral num) $ word8 0
    return 0

{--
        +----+--------+
        |VER | METHOD |
        +----+--------+
        | 1  |   1    |
        +----+--------+
--}

authResp = B.pack [5, 0]
{--
        +----+-----+-------+------+----------+----------+
        |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
        +----+-----+-------+------+----------+----------+
        | 1  |  1  | X'00' |  1   | Variable |    2     |
        +----+-----+-------+------+----------+----------+
        cmd:
            1 -> CONNECT
            2 -> BIND
            3 -> UPD
        atyp:
            1 -> IPV4
            3 -> FQDN
            4 -> IPV6
--}
cmdMsg = do
    sockVer
    cmd <- satisfy $ \w -> w == 1 -- || w == 2 || w == 3
    reserved
    atyp <- satisfy $ \w -> w == 1 || w == 3 || w == 4
    addrToken <- case atyp of
        1 -> A.take 4
        4 -> A.take 16
        3 -> do
            maybeLen  <- peekWord8
            case maybeLen of
                Just len -> A.take (fromIntegral len + 1)
                Nothing -> fail "ugh"

    portToken <- A.take 2
    return (cmd, atyp, addrToken, portToken)

{--
        +----+-----+-------+------+----------+----------+
        |VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
        +----+-----+-------+------+----------+----------+
        | 1  |  1  | X'00' |  1   | Variable |    2     |
        +----+-----+-------+------+----------+----------+
        REP:
        0 -> success
        1 -> general failure
--}

type HandshakeTokens = (Word8, ByteString, ByteString)

handshake
  :: (Proxy p, Monad m)
  => P.Pipe (P.EitherP P.ParsingError (P.StateP [ByteString] p)) (Maybe ByteString) ByteString m HandshakeTokens
handshake = do
    _ <- P.parse authMsg
    P.respond $ authResp
    (cmd, atyp, addrToken, portToken) <- P.parse cmdMsg
    P.respond $ (B.pack [5, 0, 0, atyp]) <> addrToken <> portToken
    return (atyp, addrToken, portToken)

talkSocks
  :: (P.Proxy p, Monad m)
  => () -> P.Pipe (P.EitherP P.ParsingError (P.StateP [ByteString] p)) (Maybe ByteString) ByteString m ()
talkSocks () = do
    tokens <- handshake
    -- .. possibly do something with tokens

    -- Now that we have finished the handshake, we proceed to just
    -- forward downstream all the input from upstream.
    let loop = do
          ma <- P.draw
          case ma of
            Nothing -> return ()
            Just a  -> P.respond a >> loop
    P.liftP loop


main = serve (Host "127.0.0.1") "8000" $ \(cs, _) -> do
    connect addr port $ \(ss, _) -> do
        a1 <- async $ do
            let sess = P.wrap . (socketReadS 4096 cs) >-> talkSocks >-> socketWriteD ss
            ex <- P.runProxy . P.evalStateK mempty . P.runEitherK $ sess
            case ex of
              Left e  -> error "There was some error which I'm not dealing with since this is just a demonstration"
              Right _ -> return ()
        runProxy $ socketReadS 4096 ss >-> socketWriteD cs
        wait a1
        putStrLn "Everything OK!"
    where
        addr = undefined -- parse address from addrToken
        port = undefined -- parse address from portToken