pause/resume socket properly

closes issue #2

Author: mrluanma

README.md

@@ -54,8 +54,8 @@ ws@0.6.4 node_modules/ws
 Then run:
 
 ```
-$ node local.js -s still-tor-8707.herokuapp.com -l 1080 -m rc4 -k foobar
-shadowsocks-heroku v0.9.6
+$ node local.js -s still-tor-8707.herokuapp.com -l 1080 -m rc4 -k foobar -r 80
+server listening at { address: '0.0.0.0', family: 'IPv4', port: 1080 }
 ```
 
 Change proxy settings of your browser into:

config.json

@@ -1,5 +1,6 @@
 {
   "server": "127.0.0.1",
+  "local_address": "127.0.0.1",
   "local_port": 1080,
   "remote_port": 8080,
   "password": "`try*(^^$some^$%^complex>:<>?~password/",

local.js

@@ -63,7 +63,9 @@ server = net.createServer (connection) ->
     if stage is 5
       # pipe sockets
       data = encryptor.encrypt data
-      ws.send data, { binary: true }
+      if ws.readyState is WebSocket.OPEN
+        ws.send data, { binary: true }
+        connection.pause() if ws.bufferedAmount > 0
       return
     if stage is 0
       tempBuf = new Buffer(2)
@@ -130,11 +132,15 @@ server = net.createServer (connection) ->
           ping = setInterval(->
             ws.ping "", null, true
           , 50 * 1000)
+
+          ws._socket.on "drain", ->
+            connection.resume()
+
           return
 
         ws.on "message", (data, flags) ->
           data = encryptor.decrypt data
-          connection.write(data)
+          ws._socket.pause() unless connection.write(data)
 
         ws.on "close", ->
           clearInterval ping
@@ -165,22 +171,25 @@ server = net.createServer (connection) ->
 
   connection.on "end", ->
     console.log "local disconnected"
-    ws.close() if ws
+    ws.terminate() if ws
     server.getConnections (err, count) ->
       console.log "concurrent connections:", count
       return
 
   connection.on "error", (e)->
     console.log "local error: #{e}"
-    ws.close() if ws
+    ws.terminate() if ws
     server.getConnections (err, count) ->
       console.log "concurrent connections:", count
       return
 
+  connection.on "drain", ->
+    ws.resume()
+
   connection.setTimeout timeout, ->
     console.log "local timeout"
     connection.destroy()
-    ws.close() if ws
+    ws.terminate() if ws
 
 server.listen PORT, LOCAL_ADDRESS, ->
   address = server.address()

server.js

@@ -57,7 +57,7 @@ wss.on "connection", (ws) ->
   ws.on "message", (data, flags) ->
     data = encryptor.decrypt data
     if stage is 5
-      remote.write(data)
+      ws.pause() unless remote.write(data)
       return
     if stage is 0
       try
@@ -92,14 +92,20 @@ wss.on "connection", (ws) ->
         )
         remote.on "data", (data) ->
           data = encryptor.encrypt data
-          ws.send data, { binary: true } if ws.readyState is WebSocket.OPEN
+          if ws.readyState is WebSocket.OPEN
+            ws.send data, { binary: true }
+            remote.pause() if ws.bufferedAmount > 0
+          return
 
         remote.on "end", ->
-          ws.emit "close"
+          ws.close()
           console.log "remote disconnected"
 
+        remote.on "drain", ->
+          ws.resume()
+
         remote.on "error", (e)->
-          ws.emit "close"
+          ws.close()
           console.log "remote: #{e}"
 
         remote.setTimeout timeout, ->
@@ -127,6 +133,9 @@ wss.on "connection", (ws) ->
   ws.on "ping", ->
     ws.pong '', null, true
 
+  ws._socket.on "drain", ->
+    remote.resume() if remote
+
   ws.on "close", ->
     console.log "server disconnected"
     console.log "concurrent connections:", wss.clients.length