src/third_party/tlslite/tlslite/integration/asyncstatemachine.py

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