src/third_party/tlslite/tlslite/integration/AsyncStateMachine.py

   1 """
   2 A state machine for using TLS Lite with asynchronous I/O.
   3 """
   4
   5 class AsyncStateMachine:
   6     """
   7     This is an abstract class that's used to integrate TLS Lite with
   8     asyncore and Twisted.
   9
  10     This class signals wantsReadsEvent() and wantsWriteEvent().  When
  11     the underlying socket has become readable or writeable, the event
  12     should be passed to this class by calling inReadEvent() or
  13     inWriteEvent().  This class will then try to read or write through
  14     the socket, and will update its state appropriately.
  15
  16     This class will forward higher-level events to its subclass.  For
  17     example, when a complete TLS record has been received,
  18     outReadEvent() will be called with the decrypted data.
  19     """
  20
  21     def __init__(self):
  22         self._clear()
  23
  24     def _clear(self):
  25         #These store the various asynchronous operations (i.e.
  26         #generators).  Only one of them, at most, is ever active at a
  27         #time.
  28         self.handshaker = None
  29         self.closer = None
  30         self.reader = None
  31         self.writer = None
  32
  33         #This stores the result from the last call to the
  34         #currently active operation.  If 0 it indicates that the
  35         #operation wants to read, if 1 it indicates that the
  36         #operation wants to write.  If None, there is no active
  37         #operation.
  38         self.result = None
  39
  40     def _checkAssert(self, maxActive=1):
  41         #This checks that only one operation, at most, is
  42         #active, and that self.result is set appropriately.
  43         activeOps = 0
  44         if self.handshaker:
  45             activeOps += 1
  46         if self.closer:
  47             activeOps += 1
  48         if self.reader:
  49             activeOps += 1
  50         if self.writer:
  51             activeOps += 1
  52
  53         if self.result == None:
  54             if activeOps != 0:
  55                 raise AssertionError()
  56         elif self.result in (0,1):
  57             if activeOps != 1:
  58                 raise AssertionError()
  59         else:
  60             raise AssertionError()
  61         if activeOps > maxActive:
  62             raise AssertionError()
  63
  64     def wantsReadEvent(self):
  65         """If the state machine wants to read.
  66
  67         If an operation is active, this returns whether or not the
  68         operation wants to read from the socket.  If an operation is
  69         not active, this returns None.
  70
  71         @rtype: bool or None
  72         @return: If the state machine wants to read.
  73         """
  74         if self.result != None:
  75             return self.result == 0
  76         return None
  77
  78     def wantsWriteEvent(self):
  79         """If the state machine wants to write.
  80
  81         If an operation is active, this returns whether or not the
  82         operation wants to write to the socket.  If an operation is
  83         not active, this returns None.
  84
  85         @rtype: bool or None
  86         @return: If the state machine wants to write.
  87         """
  88         if self.result != None:
  89             return self.result == 1
  90         return None
  91
  92     def outConnectEvent(self):
  93         """Called when a handshake operation completes.
  94
  95         May be overridden in subclass.
  96         """
  97         pass
  98
  99     def outCloseEvent(self):
 100         """Called when a close operation completes.
 101
 102         May be overridden in subclass.
 103         """
 104         pass
 105
 106     def outReadEvent(self, readBuffer):
 107         """Called when a read operation completes.
 108
 109         May be overridden in subclass."""
 110         pass
 111
 112     def outWriteEvent(self):
 113         """Called when a write operation completes.
 114
 115         May be overridden in subclass."""
 116         pass
 117
 118     def inReadEvent(self):
 119         """Tell the state machine it can read from the socket."""
 120         try:
 121             self._checkAssert()
 122             if self.handshaker:
 123                 self._doHandshakeOp()
 124             elif self.closer:
 125                 self._doCloseOp()
 126             elif self.reader:
 127                 self._doReadOp()
 128             elif self.writer:
 129                 self._doWriteOp()
 130             else:
 131                 self.reader = self.tlsConnection.readAsync(16384)
 132                 self._doReadOp()
 133         except:
 134             self._clear()
 135             raise
 136
 137     def inWriteEvent(self):
 138         """Tell the state machine it can write to the socket."""
 139         try:
 140             self._checkAssert()
 141             if self.handshaker:
 142                 self._doHandshakeOp()
 143             elif self.closer:
 144                 self._doCloseOp()
 145             elif self.reader:
 146                 self._doReadOp()
 147             elif self.writer:
 148                 self._doWriteOp()
 149             else:
 150                 self.outWriteEvent()
 151         except:
 152             self._clear()
 153             raise
 154
 155     def _doHandshakeOp(self):
 156         try:
 157             self.result = self.handshaker.next()
 158         except StopIteration:
 159             self.handshaker = None
 160             self.result = None
 161             self.outConnectEvent()
 162
 163     def _doCloseOp(self):
 164         try:
 165             self.result = self.closer.next()
 166         except StopIteration:
 167             self.closer = None
 168             self.result = None
 169             self.outCloseEvent()
 170
 171     def _doReadOp(self):
 172         self.result = self.reader.next()
 173         if not self.result in (0,1):
 174             readBuffer = self.result
 175             self.reader = None
 176             self.result = None
 177             self.outReadEvent(readBuffer)
 178
 179     def _doWriteOp(self):
 180         try:
 181             self.result = self.writer.next()
 182         except StopIteration:
 183             self.writer = None
 184             self.result = None
 185
 186     def setHandshakeOp(self, handshaker):
 187         """Start a handshake operation.
 188
 189         @type handshaker: generator
 190         @param handshaker: A generator created by using one of the
 191         asynchronous handshake functions (i.e. handshakeServerAsync, or
 192         handshakeClientxxx(..., async=True).
 193         """
 194         try:
 195             self._checkAssert(0)
 196             self.handshaker = handshaker
 197             self._doHandshakeOp()
 198         except:
 199             self._clear()
 200             raise
 201
 202     def setServerHandshakeOp(self, **args):
 203         """Start a handshake operation.
 204
 205         The arguments passed to this function will be forwarded to
 206         L{tlslite.TLSConnection.TLSConnection.handshakeServerAsync}.
 207         """
 208         handshaker = self.tlsConnection.handshakeServerAsync(**args)
 209         self.setHandshakeOp(handshaker)
 210
 211     def setCloseOp(self):
 212         """Start a close operation.
 213         """
 214         try:
 215             self._checkAssert(0)
 216             self.closer = self.tlsConnection.closeAsync()
 217             self._doCloseOp()
 218         except:
 219             self._clear()
 220             raise
 221
 222     def setWriteOp(self, writeBuffer):
 223         """Start a write operation.
 224
 225         @type writeBuffer: str
 226         @param writeBuffer: The string to transmit.
 227         """
 228         try:
 229             self._checkAssert(0)
 230             self.writer = self.tlsConnection.writeAsync(writeBuffer)
 231             self._doWriteOp()
 232         except:
 233             self._clear()
 234             raise
 235