From 30108a5ff3ef8e9a9013685e1a4e212539fbcd61 Mon Sep 17 00:00:00 2001
From: Florian Mounier <florian.mounier@kozea.fr>
Date: Mon, 3 Mar 2014 19:44:54 +0100
Subject: [PATCH] Fix uid management and websocket closure before shell exit

---
 butterfly/__init__.py |  2 +-
 butterfly/routes.py   | 28 ++++++++++++----------------
 butterfly/utils.py    |  7 ++-----
 3 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/butterfly/__init__.py b/butterfly/__init__.py
index f2b29a7..30e1772 100644
--- a/butterfly/__init__.py
+++ b/butterfly/__init__.py
@@ -14,7 +14,7 @@
 #
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
-__version__ = '1.2.7'
+__version__ = '1.2.8'
 
 
 import os
diff --git a/butterfly/routes.py b/butterfly/routes.py
index 491cf25..55ba3e3 100644
--- a/butterfly/routes.py
+++ b/butterfly/routes.py
@@ -196,7 +196,6 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
             return
 
         self.socket = utils.Socket(self.ws_connection.stream.socket)
-        self.errors = 0
         self.set_nodelay(True)
         self.log.info('Websocket opened %r' % self.socket)
         self.path = path
@@ -209,7 +208,7 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
                     if field[0][0] == 'commonName':
                         self.user = self.callee = field[0][1]
 
-        if self.socket.local:
+        if self.socket.local and self.socket.uid is not None:
             self.caller = utils.User(uid=self.socket.uid)
         else:
             # We don't know uid is on the other machine
@@ -231,6 +230,7 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
 
     def on_message(self, message):
         if not hasattr(self, 'writer'):
+            self.on_close()
             self.close()
             return
         if message[0] == 'R':
@@ -245,25 +245,21 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
 
     def shell_handler(self, fd, events):
         if events & ioloop.READ:
-            self.errors = 0
             try:
                 read = self.reader.read()
             except IOError:
+                pass
+            else:
                 self.log.info('READ>%r' % read)
-                self.write_message('DIE')
-                return
-
-            self.log.info('READ>%r' % read)
-            self.write_message(read.decode('utf-8', 'replace'))
-
+                if self.ws_connection:
+                    self.write_message(read.decode('utf-8', 'replace'))
+                else:
+                    events = ioloop.ERROR
         if events & ioloop.ERROR:
-            self.errors += 1
-            if self.errors >= 10:
-                self.log.info(
-                    'Too many errors, the shell must have been closed')
-                # Terminated
-                self.on_close()
-                self.close()
+            self.log.info('Error on fd, closing')
+            # Terminated
+            self.on_close()
+            self.close()
 
     def on_close(self):
         if getattr(self, 'pid', 0) == 0:
diff --git a/butterfly/utils.py b/butterfly/utils.py
index a38e4bf..82dc34d 100644
--- a/butterfly/utils.py
+++ b/butterfly/utils.py
@@ -26,15 +26,12 @@ log = getLogger('butterfly')
 
 class User(object):
     def __init__(self, uid=None, name=None):
-        if not uid and not name:
+        if uid is None and not name:
             uid = os.getuid()
         if uid is not None:
             self.pw = pwd.getpwuid(uid)
         else:
-            try:
-                self.pw = pwd.getpwnam(name)
-            except:
-                self.pw = pwd.getpwnam('daemon')
+            self.pw = pwd.getpwnam(name)
         if self.pw is None:
             raise LookupError('Unknown user')