diff --git a/moesocks.cabal b/moesocks.cabal --- a/moesocks.cabal +++ b/moesocks.cabal @@ -1,11 +1,8 @@ --- Initial moesocks.cabal generated by cabal init. For further --- documentation, see http://haskell.org/cabal/users-guide/ - name: moesocks category: Network -version: 0.1.0.10 +version: 0.1.0.11 license: Apache-2.0 -synopsis: moe for all +synopsis: A functional firewall killer description: A socks5 proxy using the client / server architecture. homepage: https://github.com/nfjinjing/moesocks license-file: LICENSE @@ -25,7 +22,8 @@ executable moesocks main-is: Main.hs - ghc-options: -Wall -fno-warn-unused-do-bind -threaded + ghc-options: -Wall -fno-warn-unused-do-bind -threaded + -rtsopts "-with-rtsopts=-N" build-depends: base > 4 && <= 5 , lens , network diff --git a/src/Network/MoeSocks/App.hs b/src/Network/MoeSocks/App.hs --- a/src/Network/MoeSocks/App.hs +++ b/src/Network/MoeSocks/App.hs @@ -137,10 +137,11 @@ let _header = shadowSocksRequestBuilder _clientRequest - sendChannel <- newTBQueueIO _TBQueue_Size - receiveChannel <- newTBQueueIO _TBQueue_Size + sendChannel <- newTBQueueIO - aConfig ^. tcpBufferSizeInPacket + receiveChannel <- newTBQueueIO - aConfig ^. tcpBufferSizeInPacket let _logId x = x <> " " <> _msg + _timeout = aConfig ^. timeout * 1000 * 1000 let sendThread = do sendBuilderEncrypted @@ -150,14 +151,19 @@ atomically . writeTBQueue sendChannel . Just =<< _encrypt _partialBytesAfterClientRequest - let _produce = produceLoop (_logId "L --> +Loop") - aSocket - sendChannel - _encrypt - let _consume = consumeLoop (_logId "L --> -Loop") - __remoteSocket - sendChannel + let _produce = do + produceLoop (_logId "L --> + Loop") + _timeout + aSocket + sendChannel + _encrypt + + let _consume = do + consumeLoop (_logId "L --> - Loop") + _timeout + __remoteSocket + sendChannel finally ( connectProduction (Just - _logId "L --> +", _produce) @@ -166,14 +172,17 @@ pure () let receiveThread = do - let _produce = produceLoop (_logId "L <-- +Loop") + let _produce = produceLoop (_logId "L <-- + Loop") + _timeout __remoteSocket receiveChannel _decrypt - let _consume = consumeLoop (_logId "L <-- -Loop") - aSocket - receiveChannel + let _consume = do + consumeLoop (_logId "L <-- - Loop") + _timeout + aSocket + receiveChannel finally ( connectProduction (Just - _logId "L <-- +", _produce) @@ -244,21 +253,25 @@ let handleTarget __leftOverBytes __targetSocket = do - sendChannel <- newTBQueueIO _TBQueue_Size - receiveChannel <- newTBQueueIO _TBQueue_Size + sendChannel <- newTBQueueIO - aConfig ^. tcpBufferSizeInPacket + receiveChannel <- newTBQueueIO - aConfig ^. tcpBufferSizeInPacket let _logId x = x <> " " <> _msg + _timeout = aConfig ^. timeout * 1000 * 1000 let sendThread = do when (_leftOverBytes & isn't _Empty) - atomically - writeTBQueue sendChannel - Just _leftOverBytes - let _produce = produceLoop (_logId "R --> +Loop") - aSocket - sendChannel - _decrypt + let _produce = do + produceLoop (_logId "R --> + Loop") + _timeout + aSocket + sendChannel + _decrypt - let _consume = consumeLoop (_logId "R --> -Loop") + let _consume = consumeLoop (_logId "R --> - Loop") + _timeout __targetSocket sendChannel @@ -270,15 +283,20 @@ pure () let receiveThread = do - let _produce = produceLoop (_logId "R --> +Loop") - __targetSocket - receiveChannel - _encrypt + let _produce = do + produceLoop (_logId "R --> + Loop") + _timeout + __targetSocket + receiveChannel + _encrypt - let _consume = consumeLoop (_logId "R --> -Loop") - aSocket - receiveChannel + let _consume = do + consumeLoop (_logId "R --> - Loop") + _timeout + aSocket + receiveChannel + finally ( connectProduction (Just - _logId "R <-- +", _produce) @@ -319,22 +337,62 @@ fixConfig :: Value -> Value fixConfig (Object _obj) = Object - - _obj & H.toList & fromSS & - over (mapped . _1) (T.cons '_') & H.fromList + _obj + & H.toList + & fromSS + & over (mapped . _1) (T.cons '_') + & H.fromList fixConfig _ = Null - _maybeConfig = _v >>= decode . encode . fixConfig :: Maybe MoeConfig + formatConfig :: Value -> Value formatConfig (Object _obj) = Object - - _obj & H.toList & - over (mapped . _1) T.tail & H.fromList + _obj + & H.toList + & over (mapped . _1) T.tail + & H.fromList formatConfig _ = Null + filterEssentialConfig :: Value -> Value + filterEssentialConfig (Object _obj) = + Object - + foldl (flip H.delete) _obj - + [ + "_remote" + , "_remotePort" + , "_local" + , "_localPort" + , "_password" + ] + + filterEssentialConfig _ = Null + mergeConfigObject :: Value -> Value -> Value + mergeConfigObject (Object _x) (Object _y) = + Object - _x `H.union` _y + mergeConfigObject _ _ = Null + + optionalConfig = filterEssentialConfig - toJSON defaultMoeConfig + + _maybeConfig = _v + >>= decode + . encode + . flip mergeConfigObject optionalConfig + . fixConfig + + let + showConfig :: MoeConfig -> ByteString + showConfig = + toStrict + . encode + . formatConfig + . toJSON + + case _maybeConfig of Nothing -> do let _r = @@ -344,16 +402,17 @@ tell "Example: \n" let configBS :: ByteString - configBS = toStrict .encode . formatConfig . toJSON - - defaultMoeConfig + configBS = showConfig defaultMoeConfig tell - configBS ^. utf8 <> "\n" throwError - _r ^. _Text Just _config -> do + let configStr = showConfig _config ^. utf8 . _Text :: String + io - puts - "Using config: " <> configStr pure - _config - + initLogger :: Priority -> IO () initLogger aLevel = do @@ -397,7 +456,7 @@ setSocketOption _localSocket ReuseAddr 1 bindSocket _localSocket _localAddr - listen _localSocket 1 + listen _localSocket maxListenQueue let handleLocal _socket = do (_newSocket, _) <- accept _socket @@ -427,9 +486,9 @@ setSocketOption _remoteSocket ReuseAddr 1 bindSocket _remoteSocket _remoteAddr - let _maximum_number_of_queued_connection = 1 + {-let _maximum_number_of_queued_connection = 1 :: Int-} - listen _remoteSocket _maximum_number_of_queued_connection + listen _remoteSocket maxListenQueue let handleRemote _socket = do (_newSocket, _) <- accept _socket diff --git a/src/Network/MoeSocks/Config.hs b/src/Network/MoeSocks/Config.hs --- a/src/Network/MoeSocks/Config.hs +++ b/src/Network/MoeSocks/Config.hs @@ -15,6 +15,8 @@ , _localPort = 1090 , _password = "moesocks" , _method = _DefaultMethod + , _timeout = 300 + , _tcpBufferSizeInPacket = 128 } diff --git a/src/Network/MoeSocks/Helper.hs b/src/Network/MoeSocks/Helper.hs --- a/src/Network/MoeSocks/Helper.hs +++ b/src/Network/MoeSocks/Helper.hs @@ -28,6 +28,7 @@ import System.IO.Unsafe (unsafePerformIO) import System.Log.Logger import System.Posix.IO (FdOption(CloseOnExec), setFdOption) +import System.Timeout (timeout) import qualified Data.ByteString as S import qualified Data.ByteString.Builder as B import qualified Data.ByteString.Lazy as LB @@ -40,8 +41,6 @@ -- END backports -_TBQueue_Size :: Int -_TBQueue_Size = 32 io :: (MonadIO m) => IO a -> m a io = liftIO @@ -58,8 +57,9 @@ sync :: IO a -> IO a sync aIO = do - putMVar syncLock () - aIO <* takeMVar syncLock + {-putMVar syncLock ()-} + {-aIO <* takeMVar syncLock-} + aIO puts :: String -> IO () puts = sync . debugM "moe" . ("😽 " <>) @@ -94,7 +94,7 @@ logSocketWithAddress aID _init f = do catch (bracket _init (logClose aID . fst) f) - \(e :: SomeException) -> do - pute - "logSocket: Exception in " <> aID <> ": " <> show e + puts - "logSocket: Exception in " <> aID <> ": " <> show e throw e logSA:: String -> IO (Socket, SockAddr) -> @@ -104,7 +104,7 @@ logSocket :: String -> IO Socket -> (Socket -> IO a) -> IO a logSocket aID _init f = catch (bracket _init (logClose aID) f) - \e -> do - pute - "Exception in " <> aID <> ": " <> show (e :: SomeException) + puts - "Exception in " <> aID <> ": " <> show (e :: SomeException) throw e catchExceptAsyncLog :: String -> IO a -> IO () @@ -129,7 +129,7 @@ logException :: String -> IO a -> IO () logException aID aIO = catch (() <$ aIO) - \e -> do - putw - "Logged error " <> aID <> ": " + puts - "Error in " <> aID <> ": " <> show (e :: SomeException) throw e @@ -149,12 +149,12 @@ waitBoth x y = do waitBothDebug (Nothing, x) (Nothing, y) -data WaitException = WaitException String +data TimeoutException = TimeoutException String -instance Show WaitException where - show (WaitException s) = "Wait exception: " ++ s +instance Show TimeoutException where + show (TimeoutException s) = "Timeout exception: " ++ s -instance Exception WaitException +instance Exception TimeoutException waitBothDebug :: (Maybe String, IO ()) -> (Maybe String, IO ()) -> IO () waitBothDebug x y = do @@ -272,31 +272,47 @@ "Failed to parse " <> _id <> ": " <> msg Partial _p -> parseSocketWith _id _p _socket -produceLoop :: String -> Socket -> TBQueue (Maybe ByteString) -> +type Timeout = Int + +timeoutFor :: String -> Timeout -> IO a -> IO a +timeoutFor aID aTimeout aIO = do + _r <- timeout aTimeout aIO + case _r of + Nothing -> throw - TimeoutException aID + Just _r -> pure _r + +produceLoop :: String -> Timeout -> Socket -> TBQueue (Maybe ByteString) -> (ByteString -> IO ByteString) -> IO () -produceLoop aID aSocket aTBQueue f = do - let _shutdown = tryIO aID - shutdown aSocket ShutdownReceive +produceLoop aID aTimeout aSocket aTBQueue f = do + let _shutdown = do + tryIO aID - shutdown aSocket ShutdownReceive + {-tryIO aID - close aSocket-} + _produce = do - _r <- recv_ aSocket + _r <- timeoutFor aID aTimeout - recv_ aSocket if (_r & isn't _Empty) then do f _r >>= atomically . writeTBQueue aTBQueue . Just _produce else do + puts - "Half closed: " <> aID atomically - writeTBQueue aTBQueue Nothing _produce `onException` _shutdown pure () -consumeLoop :: String -> Socket -> TBQueue (Maybe ByteString) -> IO () -consumeLoop aID aSocket aTBQueue = do - let _shutdown = tryIO aID - shutdown aSocket ShutdownSend +consumeLoop :: String -> Timeout -> Socket -> TBQueue (Maybe ByteString) -> IO () +consumeLoop aID aTimeout aSocket aTBQueue = do + let _shutdown = do + tryIO aID - shutdown aSocket ShutdownSend + {-tryIO aID - close aSocket-} _consume = do - _r <- atomically - readTBQueue aTBQueue + _r <- atomically - readAll aTBQueue case _r of Nothing -> _shutdown Just _data -> do - send_ aSocket _data >> _consume + timeoutFor aID aTimeout - + sendMany aSocket _data >> _consume _consume `onException` _shutdown pure () @@ -308,6 +324,32 @@ setSocketCloseOnExec :: Socket -> IO () setSocketCloseOnExec aSocket = setFdOption (fromIntegral $ fdSocket aSocket) CloseOnExec True + + +readAll :: TBQueue (Maybe ByteString) -> STM (Maybe [ByteString]) +readAll aTBQueue = do + _r <- readTBQueue aTBQueue + case _r of + Nothing -> pure Nothing + Just _b -> do + fmap (Just . reverse) - readMore [_b] aTBQueue + + where + readMore :: [ByteString] -> TBQueue (Maybe ByteString) -> STM [ByteString] + readMore _acc _q = do + _empty <- isEmptyTBQueue _q + if _empty + then pure - _acc + else + do + _more <- readTBQueue _q + case _more of + Nothing -> do + unGetTBQueue _q Nothing + pure - _acc + Just _r -> readMore (_r : _acc) _q + + tryIO :: String -> IO a -> IO (Either IOException a) diff --git a/src/Network/MoeSocks/Options.hs b/src/Network/MoeSocks/Options.hs --- a/src/Network/MoeSocks/Options.hs +++ b/src/Network/MoeSocks/Options.hs @@ -19,7 +19,7 @@ short 'm' <> long "mode" <> metavar "MODE" - <> help "local | remote" + <> help "Tell moesocks to runs in local or remote mode" ) <|> pure "local" parseMode :: String -> RunningMode @@ -35,18 +35,16 @@ short 'c' <> long "config" <> metavar "CONFIG" - <> help "path to the configuration file" + <> help "Point to the path of the configuration file" ) <|> pure "config.json" in let __verbosity :: O.Parser Priority - __verbosity = ( option auto - + __verbosity = flag INFO DEBUG - short 'v' <> long "verbose" - <> metavar "PRIORITY" - <> help "DEBUG|INFO|NOTICE" - ) <|> pure INFO + <> help "Turn on logging" in @@ -55,7 +53,7 @@ short 'L' <> metavar "port:host:hostport" <> help ("" - <> "Specifies that the given port on the local" + <> "Specify that the given port on the local" <> "(client) host is to be forwarded to the " <> "given host and port on the remote side." ) diff --git a/src/Network/MoeSocks/Type.hs b/src/Network/MoeSocks/Type.hs --- a/src/Network/MoeSocks/Type.hs +++ b/src/Network/MoeSocks/Type.hs @@ -44,7 +44,8 @@ deriving (Show, Eq) makeLenses ''ClientRequest - + + data MoeConfig = MoeConfig { _remote :: Text @@ -53,6 +54,8 @@ , _localPort :: Int , _password :: Text , _method :: Text + , _timeout :: Int + , _tcpBufferSizeInPacket :: Int } deriving (Show, Eq, Generic)