diff --git a/moesocks.cabal b/moesocks.cabal --- a/moesocks.cabal +++ b/moesocks.cabal @@ -3,7 +3,7 @@ name: moesocks category: Network -version: 0.1.0.9 +version: 0.1.0.10 license: Apache-2.0 synopsis: moe for all description: A socks5 proxy using the client / server architecture. @@ -45,6 +45,8 @@ , hslogger , transformers , mtl + , stm + , async hs-source-dirs: src default-language: Haskell2010 diff --git a/readme.md b/readme.md --- a/readme.md +++ b/readme.md @@ -14,7 +14,7 @@ Features -------- * GFW compatibility -* Can be used as a drop in replacement of shadowsocks (only client mode) +* Can be used as a drop in replacement for shadowsocks (only client mode) * TCP port forwarding Planning features 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 @@ -4,6 +4,7 @@ module Network.MoeSocks.App where import Control.Concurrent +import Control.Concurrent.STM import Control.Exception import Control.Lens import Control.Monad @@ -49,6 +50,12 @@ showConnectionType TCP_IP_port_binding = "TCP_Bind " showConnectionType UDP_port = "UDP " +showRequest :: ClientRequest -> String +showRequest _r = + view _Text (showAddressType (_r ^. addressType)) + <> ":" + <> show (_r ^. portNumber) + processLocalSocks5Request :: Socket -> IO (ClientRequest, ByteString) processLocalSocks5Request aSocket = do (_partialBytesAfterGreeting, r) <- @@ -102,27 +109,24 @@ \(_remoteSocket, _remoteAddress) -> do connect _remoteSocket _remoteAddress + _localPeerAddr <- getPeerName aSocket _remoteSocketName <- getSocketName _remoteSocket - + when shouldReplyClient - do let _connectionReplyBuilder = connectionReplyBuilder _remoteSocketName send_ aSocket - builder_To_ByteString _connectionReplyBuilder - let showRequest :: ClientRequest -> String - showRequest _r = - view _Text (showAddressType (_r ^. addressType)) - <> ":" - <> show (_r ^. portNumber) + + let _msg = + concat - L.intersperse " -> " + [ + show _localPeerAddr + , showRequest _clientRequest + ] + _log - "L " -- <> showConnectionType (_clientRequest ^. connectionType) - <> ": " <> - ( - concat - L.intersperse " -> " - [ - show _localPeerAddr - , showRequest _clientRequest - ] - ) + <> ": " <> _msg let handleLocal __remoteSocket = do (_encrypt, _decrypt) <- getCipher @@ -133,49 +137,53 @@ let _header = shadowSocksRequestBuilder _clientRequest - sendChannel <- newChan - receiveChannel <- newChan + sendChannel <- newTBQueueIO _TBQueue_Size + receiveChannel <- newTBQueueIO _TBQueue_Size + let _logId x = x <> " " <> _msg + let sendThread = do sendBuilderEncrypted sendChannel _encrypt _header when (_partialBytesAfterClientRequest & isn't _Empty) - - writeChan sendChannel . Just =<< + atomically . writeTBQueue sendChannel . Just =<< _encrypt _partialBytesAfterClientRequest - let _produce = produceLoop "L send produceLoop" + let _produce = produceLoop (_logId "L --> +Loop") aSocket sendChannel _encrypt - let _consume = consumeLoop "L send consumeLoop" + let _consume = consumeLoop (_logId "L --> -Loop") __remoteSocket sendChannel - - waitBothDebug (Just "L send produce", _produce) - (Just "L send consume", _consume) - - shutdown __remoteSocket ShutdownSend + finally + ( + connectProduction (Just - _logId "L --> +", _produce) + (Just - _logId "L --> -", _consume) + ) - + pure () let receiveThread = do - let _produce = produceLoop "L receive produceLoop" + let _produce = produceLoop (_logId "L <-- +Loop") __remoteSocket receiveChannel _decrypt - let _consume = consumeLoop "L receive consumeLoop" + let _consume = consumeLoop (_logId "L <-- -Loop") aSocket receiveChannel - - waitBothDebug (Just "L receive produce", _produce) - (Just "L receive consume", _consume) - - shutdown aSocket ShutdownSend + finally + ( + connectProduction (Just - _logId "L <-- +", _produce) + (Just - _logId "L <-- -", _consume) + ) - + pure () - waitBothDebug - (Just "L -->", sendThread) - (Just "L <--", receiveThread) + connectTunnel + (Just - _logId "L -->", sendThread) + (Just - _logId "L <--", receiveThread) handleLocal _remoteSocket @@ -195,7 +203,7 @@ aSocket - puts - "Remote get: " <> show _clientRequest + {-puts - "Remote get: " <> show _clientRequest-} let initTarget :: ClientRequest -> IO (Socket, SockAddr) @@ -224,56 +232,63 @@ _remotePeerAddr <- getPeerName aSocket _targetPeerAddr <- getPeerName _targetSocket + let _msg = + concat - L.intersperse " -> " - + [ + show _remotePeerAddr + , showRequest _clientRequest + ] + _log - "R " -- <> showConnectionType (_clientRequest ^. connectionType) - <> ": " <> - ( - concat - L.intersperse " -> " - map show - [ - _remotePeerAddr - , _targetPeerAddr - ] - ) + <> ": " <> _msg + let handleTarget __leftOverBytes __targetSocket = do - sendChannel <- newChan - receiveChannel <- newChan + sendChannel <- newTBQueueIO _TBQueue_Size + receiveChannel <- newTBQueueIO _TBQueue_Size + let _logId x = x <> " " <> _msg + let sendThread = do when (_leftOverBytes & isn't _Empty) - - writeChan sendChannel - Just _leftOverBytes + atomically - writeTBQueue sendChannel - Just _leftOverBytes - let _produce = produceLoop "R send produceLoop" + let _produce = produceLoop (_logId "R --> +Loop") aSocket sendChannel _decrypt - let _consume = consumeLoop "R send consumeLoop" + let _consume = consumeLoop (_logId "R --> -Loop") __targetSocket sendChannel - waitBothDebug (Just "R send produce", _produce) - (Just "R send consume", _consume) - - shutdown __targetSocket ShutdownSend + finally + ( + connectProduction (Just - _logId "R --> +", _produce) + (Just - _logId "R --> -", _consume) + ) - + pure () let receiveThread = do - let _produce = produceLoop "R send produceLoop" + let _produce = produceLoop (_logId "R --> +Loop") __targetSocket receiveChannel _encrypt - let _consume = consumeLoop "R send consumeLoop" + let _consume = consumeLoop (_logId "R --> -Loop") aSocket receiveChannel - waitBothDebug (Just "R receive produce", _produce) - (Just "R receive consume", _consume) - - shutdown aSocket ShutdownSend + finally + ( + connectProduction (Just - _logId "R <-- +", _produce) + (Just - _logId "R <-- -", _consume) + ) - + pure () - waitBothDebug - (Just "R -->", sendThread) - (Just "R <--", receiveThread) + connectTunnel + (Just - _logId "R -->", sendThread) + (Just - _logId "R <--", receiveThread) handleTarget _leftOverBytes _targetSocket @@ -458,7 +473,9 @@ debugRun :: IO () debugRun = do catchExceptAsyncLog "Debug app" - do - waitNone localRun remoteRun + waitBothDebug + (Just "localRun", localRun) + (Just "remoteRun", remoteRun) io - case _options ^. runningMode of DebugMode -> debugRun diff --git a/src/Network/MoeSocks/Constant.hs b/src/Network/MoeSocks/Constant.hs --- a/src/Network/MoeSocks/Constant.hs +++ b/src/Network/MoeSocks/Constant.hs @@ -20,3 +20,5 @@ _DefaultMethod :: Text _DefaultMethod = "aes-256-cfb" + + 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 @@ -1,18 +1,22 @@ {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE CPP #-} +{-# LANGUAGE RecursiveDo #-} module Network.MoeSocks.Helper where import Control.Concurrent +import Control.Concurrent.Async +import Control.Concurrent.STM import Control.Exception import Control.Lens import Control.Monad import Control.Monad.IO.Class -import Data.Attoparsec.ByteString +import Data.Attoparsec.ByteString hiding (try) import Data.Binary import Data.Binary.Put import Data.ByteString (ByteString) +import Data.Maybe import Data.Monoid import Data.Text (Text) import Data.Text.Lens @@ -28,8 +32,6 @@ import qualified Data.ByteString.Builder as B import qualified Data.ByteString.Lazy as LB -import Data.Maybe - -- BEGIN backports infixr 0 - @@ -38,6 +40,9 @@ -- END backports +_TBQueue_Size :: Int +_TBQueue_Size = 32 + io :: (MonadIO m) => IO a -> m a io = liftIO @@ -62,6 +67,9 @@ pute :: String -> IO () pute = sync . errorM "moe" . ("😾 " <>) +putw :: String -> IO () +putw = sync . warningM "moe" . ("😾 " <>) + _log :: String -> IO () _log = sync . infoM "moe" . ("😺 " <>) @@ -78,7 +86,7 @@ logClose :: String -> Socket -> IO () logClose aID aSocket = do pure aID - {-puts - "Closing socket " <> aID-} + puts - "Closing socket " <> aID close aSocket logSocketWithAddress :: String -> IO (Socket, SockAddr) -> @@ -103,106 +111,65 @@ catchExceptAsyncLog aID aIO = catches (() <$ aIO) [ Handler - \(e :: AsyncException) -> do - puts - "ASyncException in " + pute - "ASyncException in " <> aID <> " : " <> show e throw e , Handler - \(e :: SomeException) -> - puts - "CatcheAll in " + pute - "CatcheAll in " <> aID <> " : " <> show e ] catchIO:: String -> IO a -> IO () catchIO aID aIO = catch (() <$ aIO) - \e -> - puts - "Catch IO in " <> aID <> ": " + pute - "IOError in " <> aID <> ": " <> show (e :: IOException) +logException :: String -> IO a -> IO () +logException aID aIO = catch (() <$ aIO) - \e -> + do + putw - "Logged error " <> aID <> ": " + <> show (e :: SomeException) + throw e +logWaitIO :: (Maybe String, IO a) -> IO () +logWaitIO x = do + let _io = wrapIO x + aID = x ^. _1 & fromMaybe "" + + puts - "waiting for : " <> aID + _io + wrapIO :: (Maybe String, IO c) -> IO () wrapIO (s, _io) = do - pure s - {-forM_ s - puts . ("+ " <>)-} - catchExceptAsyncLog (fromMaybe "" s) _io - {-catch (() <$ _io) - \(e :: IOException) -> pure ()-} - {-<* (forM_ s - puts . ("- " <>))-} - -{-waitFirst :: IO () -> IO () -> IO ()-} -{-waitFirst = runWait True False-} - -{-waitFirstDebug :: (Maybe String, IO ()) -> (Maybe String, IO ()) -> IO ()-} -{-waitFirstDebug = runWaitDebug True False-} + logException (fromMaybe "" s) _io waitBoth :: IO () -> IO () -> IO () -waitBoth = runWait True True - -waitBothDebug :: (Maybe String, IO ()) -> (Maybe String, IO ()) -> IO () -waitBothDebug = runWaitDebug True True - -waitNone :: IO () -> IO () -> IO () -waitNone = runWait False False +waitBoth x y = do + waitBothDebug (Nothing, x) (Nothing, y) -waitNoneDebug :: (Maybe String, IO ()) -> (Maybe String, IO ()) -> IO () -waitNoneDebug = runWaitDebug False False +data WaitException = WaitException String -runWait :: Bool -> Bool -> IO () -> IO () -> IO () -runWait _waitX _waitY x y = do - runWaitDebug _waitX _waitY (Nothing, x) (Nothing, y) +instance Show WaitException where + show (WaitException s) = "Wait exception: " ++ s -runWaitDebug :: Bool -> Bool -> (Maybe String, IO ()) -> - (Maybe String, IO ()) -> IO () -runWaitDebug _waitX _waitY x y = do - let _x = wrapIO x - _y = wrapIO y +instance Exception WaitException - _xID = x ^. _1 & fromMaybe "" +waitBothDebug :: (Maybe String, IO ()) -> (Maybe String, IO ()) -> IO () +waitBothDebug x y = do + concurrently (logWaitIO x) (logWaitIO y) + let _xID = x ^. _1 & fromMaybe "" _yID = y ^. _1 & fromMaybe "" _hID = _xID <> " / " <> _yID - - let _init = do - _threadXDone <- newEmptyMVar - _threadYDone <- newEmptyMVar - xThreadID <- - forkFinally _x - - const - putMVar _threadXDone () - - yThreadID <- - forkFinally _y - const - do - when (not _waitX) - do - _threadXRunning <- isEmptyMVar _threadXDone - when _threadXRunning - killThread xThreadID - puts - "killing thread X: " <> _xID - - putMVar _threadYDone () - - return ((_threadXDone, xThreadID), (_threadYDone, yThreadID)) - - let handleError ((_, xThreadID), (_, yThreadID)) = do - puts - "handleError for " <> _hID - pure xThreadID - pure yThreadID - pure () - {-killThread yThreadID-} - {-killThread xThreadID-} - - let action ((_threadXDone, _), (_threadYDone, yThreadID)) = do - catchExceptAsyncLog _hID - do - {-puts - "waiting for " <> _xID-} - takeMVar _threadXDone - - when (not _waitY) - do - _threadYRunning <- isEmptyMVar _threadYDone - when _threadYRunning - killThread yThreadID - {-puts - "killing thread Y: " <> _yID-} + puts - "All done for " <> _hID + pure () - {-puts - "waiting for " <> _yID-} - takeMVar _threadYDone - {-puts - "All done for " <> _hID-} +connectTunnel :: (Maybe String, IO ()) -> (Maybe String, IO ()) -> IO () +connectTunnel = waitBothDebug - bracketOnError - _init - handleError - action +connectProduction :: (Maybe String, IO ()) -> (Maybe String, IO ()) -> IO () +connectProduction = waitBothDebug getSocket :: (Integral i, Show i) => Text -> i -> SocketType -> IO (Socket, SockAddr) @@ -269,12 +236,14 @@ send_ :: Socket -> ByteString -> IO () send_ = sendAll -sendBuilder :: Chan (Maybe ByteString) -> B.Builder -> IO () -sendBuilder _chan = writeChan _chan . Just . builder_To_ByteString +sendBuilder :: TBQueue (Maybe ByteString) -> B.Builder -> IO () +sendBuilder _queue = + atomically . writeTBQueue _queue . Just . builder_To_ByteString -sendBuilderEncrypted :: Chan (Maybe ByteString) -> +sendBuilderEncrypted :: TBQueue (Maybe ByteString) -> (ByteString -> IO ByteString) -> B.Builder -> IO () -sendBuilderEncrypted _chan _encrypt x = writeChan _chan . Just =<< +sendBuilderEncrypted _queue _encrypt x = + atomically . writeTBQueue _queue . Just =<< _encrypt (builder_To_ByteString x) -- | An exception raised when parsing fails. @@ -303,39 +272,43 @@ "Failed to parse " <> _id <> ": " <> msg Partial _p -> parseSocketWith _id _p _socket -produceLoop :: String -> Socket -> Chan (Maybe ByteString) -> +produceLoop :: String -> Socket -> TBQueue (Maybe ByteString) -> (ByteString -> IO ByteString) -> IO () -produceLoop aID aSocket aChan f = _produce - where - _produce = do - _r <- recv_ aSocket `catch` \(e :: IOException) -> - do - puts - aID <> ": " <> show e - pure mempty +produceLoop aID aSocket aTBQueue f = do + let _shutdown = tryIO aID - shutdown aSocket ShutdownReceive + _produce = do + _r <- recv_ aSocket + if (_r & isn't _Empty) + then do + f _r >>= atomically . writeTBQueue aTBQueue . Just + _produce + else do + atomically - writeTBQueue aTBQueue Nothing - if (_r & isn't _Empty) - then do - f _r >>= writeChan aChan . Just - _produce - else do - writeChan aChan Nothing + _produce `onException` _shutdown + pure () -consumeLoop :: String -> Socket -> Chan (Maybe ByteString) -> IO () -consumeLoop aID aSocket aChan = _consume - where - _consume = do - _r <- readChan aChan - case _r of - Nothing -> do - pure () - Just _data -> do - (send_ aSocket _data >> _consume) `catch` - \(e :: IOException) -> - (puts - aID <> ": " <> show e) +consumeLoop :: String -> Socket -> TBQueue (Maybe ByteString) -> IO () +consumeLoop aID aSocket aTBQueue = do + let _shutdown = tryIO aID - shutdown aSocket ShutdownSend + _consume = do + _r <- atomically - readTBQueue aTBQueue + case _r of + Nothing -> _shutdown + Just _data -> do + send_ aSocket _data >> _consume + + _consume `onException` _shutdown + pure () + -- Copied and slightly modified from: -- https://github.com/mzero/plush/blob/master/src/Plush/Server/Warp.hs setSocketCloseOnExec :: Socket -> IO () setSocketCloseOnExec aSocket = setFdOption (fromIntegral $ fdSocket aSocket) CloseOnExec True + + +tryIO :: String -> IO a -> IO (Either IOException a) +tryIO _ = try -- . logException aID