From 2f48c7b95dfa3ee04222df1f6ea1c0ecab1b88fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Er=C3=A8be=20-=20Romain=20Gerard?= <r.gerard@criteo.com>
Date: Wed, 22 Jan 2020 19:04:07 +0100
Subject: [PATCH] Add SO_MARK support (#34)

* Add SO_MARK support

* SO_MARK is now parametable


Former-commit-id: d3a8db61378688b10741e37e9750962250d5b413
Former-commit-id: 3f2cfb1cf7bf25b7e98b908d7962c0d04c87b173 [formerly 4d7672c8b097390de2d4789a3612a7acdb251563] [formerly 84f02bd794a524c46a3303851b9090178c07e315 [formerly cc09d795159786b36d9e9f5ab48c92260cd0cdf5 [formerly cc09d795159786b36d9e9f5ab48c92260cd0cdf5 [formerly cc09d795159786b36d9e9f5ab48c92260cd0cdf5 [formerly 603b65e7d9fdd0aa421ef2cf5f8684eaca044dc0]]]]]
Former-commit-id: cf75e5c7a77085f4c4bfce4da8c01ab9dcd94edf [formerly 4033012e5f39470e3c08b8918f686afe3dfb0200]
Former-commit-id: a903f60697af67855408554260028dbea6fade14
Former-commit-id: 3115a5bc330ac8708eda9430b6f85ef5a0a8d1ce
Former-commit-id: 6b2ba8627004dd0c916ecdaf9aea9cebb6de312e
Former-commit-id: 5f65ff28217d015ac6a76a0803e677ac728563e4 [formerly 677120e2a339cc83c420f12354fbc4b1e3f7d91e]
Former-commit-id: 0a91ab53293d5383b251f470962b640126dea526
---
 app/Main.hs   | 4 ++++
 src/Tunnel.hs | 2 ++
 src/Types.hs  | 7 +++++++
 3 files changed, 13 insertions(+)

diff --git a/app/Main.hs b/app/Main.hs
index 572de7a..a91bfac 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -25,6 +25,7 @@ data WsTunnel = WsTunnel
   , udpMode         :: Bool
   , udpTimeout      :: Int
   , proxy           :: String
+  , soMark          :: Int
   , serverMode      :: Bool
   , restrictTo      :: String
   , verbose         :: Bool
@@ -62,6 +63,8 @@ cmdLine = WsTunnel
                          &= typ "String" &= groupname "Client options"
   , proxy          = def &= explicit &= name "p" &= name "httpProxy"
                          &= help "If set, will use this proxy to connect to the server" &= typ "USER:PASS@HOST:PORT"
+  , soMark         = def &= explicit &= name "soMark"
+                         &= help "(linux only) Mark network packet with SO_MARK sockoption with the specified value" &= typ "int"
   , wsTunnelServer = def &= argPos 0 &= typ "ws[s]://wstunnelServer[:port]"
 
   , serverMode     = def &= explicit &= name "server"
@@ -173,6 +176,7 @@ main = do
                             then Logger.VERBOSE
                             else Logger.NORMAL)
 
+  _ <- writeIORef sO_MARK_Value (soMark cfg)
   runApp cfg serverInfo
   putStrLn "Goodbye !"
   return ()
diff --git a/src/Tunnel.hs b/src/Tunnel.hs
index cf4b951..e0fd81a 100644
--- a/src/Tunnel.hs
+++ b/src/Tunnel.hs
@@ -44,6 +44,8 @@ rrunTCPClient cfg app = bracket
       (s,addr) <- N.getSocketFamilyTCP (N.getHost cfg) (N.getPort cfg) (N.getAddrFamily cfg)
       N.setSocketOption s N.RecvBuffer defaultRecvBufferSize
       N.setSocketOption s N.SendBuffer defaultSendBufferSize
+      so_mark_val <- readIORef sO_MARK_Value
+      when (N.isSupportedSocketOption sO_MARK) (N.setSocketOption s sO_MARK so_mark_val)
       return (s,addr)
     )
     (\r -> catch (N.close $ fst r) (\(_ :: SomeException) -> return ()))
diff --git a/src/Types.hs b/src/Types.hs
index 839ea63..ae6c49c 100644
--- a/src/Types.hs
+++ b/src/Types.hs
@@ -35,6 +35,13 @@ defaultRecvBufferSize = unsafeDupablePerformIO $
 defaultSendBufferSize :: Int
 defaultSendBufferSize = defaultRecvBufferSize
 
+sO_MARK :: N.SocketOption
+sO_MARK = N.CustomSockOpt (fromIntegral 1, fromIntegral 36) -- https://elixir.bootlin.com/linux/latest/source/arch/alpha/include/uapi/asm/socket.h#L64
+
+{-# NOINLINE sO_MARK_Value #-}
+sO_MARK_Value :: IORef Int
+sO_MARK_Value = unsafeDupablePerformIO $ (newIORef 131072)
+
 data Protocol = UDP | TCP | STDIO | SOCKS5 deriving (Show, Read, Eq)
 
 data StdioAppData = StdioAppData