/gdata/tlslite/integration/AsyncStateMachine.py

  1"""
  2A state machine for using TLS Lite with asynchronous I/O.
  3"""
  4
  5class 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