def test_handle_error(self):
sock = self.mox.CreateMock(socket.socket)
sock.fileno = lambda: -1
pp = PPV2Edge(self, sock, (None, None))
self.mox.StubOutWithMock(ProxyProtocolV2, 'process_pp_v2')
ProxyProtocolV2.process_pp_v2(sock, b'').AndRaise(AssertionError)
self.mox.ReplayAll()
pp.handle(sock, None)
self.assertTrue(pp.handle_called)
def test_handle_error(self):
sock = self.mox.CreateMock(socket.socket)
sock.fileno = lambda: -1
pp = PPV2Edge(self, sock, (None, None))
self.mox.StubOutWithMock(ProxyProtocolV2, 'process_pp_v2')
ProxyProtocolV2.process_pp_v2(sock, b'').AndRaise(AssertionError)
self.mox.ReplayAll()
pp.handle(sock, None)
self.assertTrue(pp.handle_called)
def test_handle_pp_v2(self):
sock = self.mox.CreateMock(socket.socket)
sock.fileno = lambda: -1
pp = PPEdge(self, sock, 13)
self.mox.StubOutWithMock(ProxyProtocol, '_ProxyProtocol__read_pp_initial')
self.mox.StubOutWithMock(ProxyProtocolV2, 'process_pp_v2')
pp._ProxyProtocol__read_pp_initial(sock).AndReturn(b'\r\n\r\n\x00\r\nQ')
ProxyProtocolV2.process_pp_v2(sock, b'\r\n\r\n\x00\r\nQ').AndReturn((13, 14))
self.mox.ReplayAll()
pp.handle(sock, None)
self.assertTrue(pp.handle_called)
def test_handle_pp_v2(self):
sock = self.mox.CreateMock(socket.socket)
sock.fileno = lambda: -1
pp = PPEdge(self, sock, 13)
self.mox.StubOutWithMock(ProxyProtocol,
'_ProxyProtocol__read_pp_initial')
self.mox.StubOutWithMock(ProxyProtocolV2, 'process_pp_v2')
pp._ProxyProtocol__read_pp_initial(sock).AndReturn(
b'\r\n\r\n\x00\r\nQ')
ProxyProtocolV2.process_pp_v2(sock, b'\r\n\r\n\x00\r\nQ').AndReturn(
(13, 14))
self.mox.ReplayAll()
pp.handle(sock, None)
self.assertTrue(pp.handle_called)
def handle(self, socket, address):
"""
This method is mostly copy-paste from Slimta to be able to handle
remote ip retrieved before calling TransactionalSmtpEdge.handle method
"""
try:
initial = self._ProxyProtocol__read_pp_initial(socket)
if initial.startswith(b'PROXY '):
address, _ = ProxyProtocolV1.process_pp_v1(socket, initial)
elif initial == b'\r\n\r\n\x00\r\nQ':
address, _ = ProxyProtocolV2.process_pp_v2(socket, initial)
else:
raise AssertionError('Invalid proxy protocol signature')
except LocalConnection:
log.error(
"Parsed proxy protocol header si for local "
"connection, and should not be proxied.")
return
except AssertionError as exc:
log.error(
'Error while retrieving ProxyProto '
'address: {} ({})'.format(socket, str(exc)))
address = invalid_pp_source_address
else:
log.debug(
'Successfully retrieve ProxyProto address: {}:{}'.format(
address[0], address[1]))
request_started.send(
sender='transactional-edge-{}-{}'.format(address[0], address[1]))
# Don't call super() to work around python MRO
TransactionalSmtpEdge.handle(self, socket, address)
def test_process_pp_v2(self):
sock = self.mox.CreateMock(socket.socket)
self.mox.StubOutWithMock(ProxyProtocolV2, '_ProxyProtocolV2__read_pp_data')
self.mox.StubOutWithMock(ProxyProtocolV2, '_ProxyProtocolV2__parse_pp_data')
self.mox.StubOutWithMock(ProxyProtocolV2, '_ProxyProtocolV2__parse_pp_addresses')
ProxyProtocolV2._ProxyProtocolV2__read_pp_data(sock, 16, b'init').AndReturn(b'data1')
ProxyProtocolV2._ProxyProtocolV2__parse_pp_data(b'data1').AndReturn((1, 2, 3, 4))
ProxyProtocolV2._ProxyProtocolV2__read_pp_data(sock, 4, b'').AndReturn(b'data2')
ProxyProtocolV2._ProxyProtocolV2__parse_pp_addresses(2, b'data2').AndReturn((13, 14))
self.mox.ReplayAll()
src_addr, dst_addr = ProxyProtocolV2.process_pp_v2(sock, b'init')
self.assertEqual(13, src_addr)
self.assertEqual(14, dst_addr)
def test_process_pp_v2(self):
sock = self.mox.CreateMock(socket.socket)
self.mox.StubOutWithMock(ProxyProtocolV2,
'_ProxyProtocolV2__read_pp_data')
self.mox.StubOutWithMock(ProxyProtocolV2,
'_ProxyProtocolV2__parse_pp_data')
self.mox.StubOutWithMock(ProxyProtocolV2,
'_ProxyProtocolV2__parse_pp_addresses')
ProxyProtocolV2._ProxyProtocolV2__read_pp_data(
sock, 16, b'init').AndReturn(b'data1')
ProxyProtocolV2._ProxyProtocolV2__parse_pp_data(b'data1').AndReturn(
(1, 2, 3, 4))
ProxyProtocolV2._ProxyProtocolV2__read_pp_data(sock, 4,
b'').AndReturn(b'data2')
ProxyProtocolV2._ProxyProtocolV2__parse_pp_addresses(
2, b'data2').AndReturn((13, 14))
self.mox.ReplayAll()
src_addr, dst_addr = ProxyProtocolV2.process_pp_v2(sock, b'init')
self.assertEqual(13, src_addr)
self.assertEqual(14, dst_addr)
def setUp(self):
super(TestProxyProtocolV2, self).setUp()
self.pp = ProxyProtocolV2()