From 0340dc49f11c57d65d1fe359ad8b1e48eb3ae0e1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Er=C3=A8be?= <romain.gerard@erebe.eu>
Date: Mon, 16 May 2016 23:55:06 +0200
Subject: [PATCH] Maj restricted option

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

diff --git a/app/Main.hs b/app/Main.hs
index c10f7cf..459f2b5 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns       #-}
 {-# LANGUAGE DeriveDataTypeable #-}
 {-# OPTIONS_GHC -fno-cse #-}
 
@@ -5,7 +6,7 @@ module Main where
 
 import           Lib
 
-import           ClassyPrelude          (guard, readMay)
+import           ClassyPrelude          (ByteString, guard, readMay)
 import qualified Data.ByteString.Char8  as BC
 import           Data.Maybe             (fromMaybe)
 import           System.Console.CmdArgs
@@ -78,16 +79,16 @@ 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 :: String -> ((ByteString, Int)-> Bool)
 parseRestrictTo "" = const True
-parseRestrictTo str = let (h, p) = fromMaybe (error "Invalid Parameter restart") parse
-  in (\(hst, port) -> hst == h && port == p)
+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)
+              return (BC.pack (head ret), portNumber)
 
 main :: IO ()
 main = do
diff --git a/src/Lib.hs b/src/Lib.hs
index 02f23a7..39c90b8 100644
--- a/src/Lib.hs
+++ b/src/Lib.hs
@@ -136,14 +136,15 @@ runTunnelingServer (host, port) isAllowed = do
     case path of
       Nothing -> putStrLn "Rejecting connection" >> WS.rejectRequest pendingConn "Invalid tunneling information"
       Just (!proto, !rhost, !rport) ->
-        if isAllowed (rhost, rport)
+        if not $ isAllowed (rhost, rport)
         then do
+          putStrLn "Rejecting tunneling"
+          WS.rejectRequest pendingConn "Restriction is on, You cannot request this tunneling"
+        else 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"
 
@@ -180,9 +181,8 @@ runClient useTls proto local wsServer remote = do
         TCP -> runTCPServer local (\hOther -> out (`propagateRW` hOther))
 
 
-runServer :: (HostName, PortNumber) -> ((String, Int) -> Bool) -> IO ()
-runServer wsInfo isAllowed = let isAllowed' (str, port) = isAllowed (BC.unpack str, fromIntegral port)
-                             in runTunnelingServer wsInfo isAllowed'
+runServer :: (HostName, PortNumber) -> ((ByteString, Int) -> Bool) -> IO ()
+runServer = runTunnelingServer
 
 
 runTlsTunnelingClient :: Proto -> (HostName, PortNumber) -> (HostName, PortNumber) -> (WS.Connection -> IO ()) -> IO ()