From b40e0e677c23fd28e8e19f2a5dafc6e1ce9c041e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Er=C3=A8be?= <romain.gerard@erebe.eu>
Date: Mon, 16 May 2016 23:33:00 +0200
Subject: [PATCH] Add restric to for the server

---
 app/Main.hs | 16 ++++++++++++++--
 src/Lib.hs  | 25 +++++++++++++++----------
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/app/Main.hs b/app/Main.hs
index 72c8826..c10f7cf 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -5,8 +5,9 @@ module Main where
 
 import           Lib
 
-import           ClassyPrelude          (readMay)
+import           ClassyPrelude          (guard, readMay)
 import qualified Data.ByteString.Char8  as BC
+import           Data.Maybe             (fromMaybe)
 import           System.Console.CmdArgs
 import           System.Environment     (getArgs, withArgs)
 
@@ -77,6 +78,17 @@ parseTunnelInfo str = mk $ BC.unpack <$> BC.split ':' (BC.pack str)
     mk _                        = error $  "Invalid tunneling information `" ++ str ++ "`, please use format [BIND:]PORT:HOST:PORT"
 
 
+parseRestrictTo :: String -> ((String, Int)-> Bool)
+parseRestrictTo "" = const True
+parseRestrictTo str = let (h, p) = fromMaybe (error "Invalid Parameter restart") parse
+  in (\(hst, port) -> hst == h && port == p)
+  where
+    parse = do
+              let ret = BC.unpack <$> BC.split ':' (BC.pack str)
+              guard (length ret == 2)
+              portNumber <- readMay $ ret !! 1 :: Maybe Int
+              return (head ret, portNumber)
+
 main :: IO ()
 main = do
   args <- getArgs
@@ -87,7 +99,7 @@ main = do
 
   if serverMode cfg
     then putStrLn ("Starting server with opts " ++ show serverInfo )
-         >> runServer (host serverInfo, fromIntegral $ port serverInfo)
+         >> runServer (host serverInfo, fromIntegral $ port serverInfo) (parseRestrictTo $ restrictTo cfg)
     else if not $ null (localToRemote cfg)
                then let (TunnelInfo lHost lPort rHost rPort) = parseTunnelInfo (localToRemote cfg)
                     in runClient (useTls serverInfo) (if udpMode cfg then UDP else TCP) (lHost, (fromIntegral lPort))
diff --git a/src/Lib.hs b/src/Lib.hs
index b68fc0b..02f23a7 100644
--- a/src/Lib.hs
+++ b/src/Lib.hs
@@ -128,18 +128,22 @@ runTunnelingClient proto (wsHost, wsPort) (remoteHost, remotePort) app = do
   putStrLn $ "CLOSE connection to " <> tshow remoteHost <> ":" <> tshow remotePort
 
 
-runTunnelingServer :: (HostName, PortNumber) -> IO ()
-runTunnelingServer (host, port) = do
+runTunnelingServer :: (HostName, PortNumber) -> ((ByteString, Int) -> Bool) -> IO ()
+runTunnelingServer (host, port) isAllowed = do
   putStrLn $ "WAIT for connection on " <> tshow host <> ":" <> tshow port
   WS.runServer host (fromIntegral port) $ \pendingConn -> do
-    let path = parsePath . WS.requestPath $ WS.pendingRequest pendingConn
+    let path =  parsePath . WS.requestPath $ WS.pendingRequest pendingConn
     case path of
       Nothing -> putStrLn "Rejecting connection" >> WS.rejectRequest pendingConn "Invalid tunneling information"
-      Just (!proto, !rhost, !rport) -> do
-        conn <- WS.acceptRequest pendingConn
-        case proto of
-          UDP -> runUDPClient (BC.unpack rhost, fromIntegral rport) (propagateRW conn)
-          TCP -> runTCPClient (BC.unpack rhost, fromIntegral rport) (propagateRW conn)
+      Just (!proto, !rhost, !rport) ->
+        if isAllowed (rhost, rport)
+        then do
+          conn <- WS.acceptRequest pendingConn
+          case proto of
+            UDP -> runUDPClient (BC.unpack rhost, fromIntegral rport) (propagateRW conn)
+            TCP -> runTCPClient (BC.unpack rhost, fromIntegral rport) (propagateRW conn)
+        else
+          putStrLn "Rejecting tunneling" >> WS.rejectRequest pendingConn "Restriction is on, You cannot request this tunneling"
 
   putStrLn "CLOSE server"
 
@@ -176,8 +180,9 @@ runClient useTls proto local wsServer remote = do
         TCP -> runTCPServer local (\hOther -> out (`propagateRW` hOther))
 
 
-runServer :: (HostName, PortNumber) -> IO ()
-runServer = runTunnelingServer
+runServer :: (HostName, PortNumber) -> ((String, Int) -> Bool) -> IO ()
+runServer wsInfo isAllowed = let isAllowed' (str, port) = isAllowed (BC.unpack str, fromIntegral port)
+                             in runTunnelingServer wsInfo isAllowed'
 
 
 runTlsTunnelingClient :: Proto -> (HostName, PortNumber) -> (HostName, PortNumber) -> (WS.Connection -> IO ()) -> IO ()