def test_no_trace_headers(self, getrandbits):
getrandbits.return_value = 1234
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans)
client = TestService.Client(client_prot)
try:
client.example_simple()
except (TTransportException, EOFError):
pass # we don't have a test response for the client
itrans = TMemoryBuffer(client_memory_trans.getvalue())
iprot = THeaderProtocolFactory().getProtocol(itrans)
self.processor.process(iprot, self.oprot)
self.assertEqual(self.observer.on_server_span_created.call_count, 1)
context, server_span = self.observer.on_server_span_created.call_args[
0]
self.assertEqual(server_span.trace_id, 1234)
self.assertEqual(server_span.parent_id, None)
self.assertEqual(server_span.id, 1234)
self.assertEqual(self.server_observer.on_start.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_args[0], (None, ))
def test_b3_trace_headers(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans)
client_header_trans = client_prot.trans
client_header_trans.set_header(b"B3-TraceId", b"1234")
client_header_trans.set_header(b"B3-ParentSpanId", b"2345")
client_header_trans.set_header(b"B3-SpanId", b"3456")
client_header_trans.set_header(b"B3-Sampled", b"1")
client_header_trans.set_header(b"B3-Flags", b"1")
client = TestService.Client(client_prot)
try:
client.example_simple()
except (TTransportException, EOFError):
pass # we don't have a test response for the client
itrans = TMemoryBuffer(client_memory_trans.getvalue())
iprot = THeaderProtocolFactory().getProtocol(itrans)
self.processor.process(iprot, self.oprot)
self.assertEqual(self.observer.on_server_span_created.call_count, 1)
context, server_span = self.observer.on_server_span_created.call_args[
0]
self.assertEqual(server_span.trace_id, 1234)
self.assertEqual(server_span.parent_id, 2345)
self.assertEqual(server_span.id, 3456)
self.assertTrue(server_span.sampled)
self.assertEqual(server_span.flags, 1)
with self.assertRaises(NoAuthenticationError):
context.request_context.user.id
self.assertEqual(self.server_observer.on_start.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_args[0], (None, ))
def test_expected_exception_not_passed_to_server_span_finish(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans)
client = TestService.Client(client_prot)
try:
client.example_throws(crash=False)
except (TTransportException, EOFError):
pass # we don't have a test response for the client
itrans = TMemoryBuffer(client_memory_trans.getvalue())
iprot = THeaderProtocolFactory().getProtocol(itrans)
self.processor.process(iprot, self.oprot)
self.assertEqual(self.server_observer.on_start.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_args[0], (None, ))
def getProtocol(self, thrift_transport):
'''
Given a Thrift transport instance, which is compatible with
asyncio.Transport, this method should create a Thrift protocol object.
'''
return THeaderProtocolFactory(
client_type=self.client_type).getProtocol(thrift_transport, )
def main():
op = optparse.OptionParser(usage='%prog [options]', add_help_option=False)
op.add_option('-p', '--port',
action='store', type='int', dest='port', default=1234,
help='The server port')
op.add_option('-c', '--with-context',
action='store_true', help='Use the generated ContextIface')
op.add_option('-h', '--header',
action='store_true', help='Use the generated ContextIface')
op.add_option('-?', '--help',
action='help',
help='Show this help message and exit')
(options, args) = op.parse_args()
if args:
op.error('trailing arguments: ' + ' '.join(args))
if options.with_context:
handler = AsyncContextLoadHandler()
processor = LoadTest.ContextProcessor(handler)
else:
handler = AsyncLoadHandler()
processor = LoadTest.Processor(handler)
if options.header:
proto_factory = THeaderProtocolFactory(True,
[THeaderTransport.HEADERS_CLIENT_TYPE,
THeaderTransport.FRAMED_DEPRECATED])
else:
proto_factory = TBinaryProtocolAcceleratedFactory()
server_factory = ThriftServerFactory(processor, proto_factory)
reactor.listenTCP(options.port, server_factory)
print 'Serving requests on port %d...' % (options.port,)
reactor.run()
def test_client_proxy_flow(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans)
class Pool(object):
@contextlib.contextmanager
def connection(self):
yield client_prot
client_factory = ThriftContextFactory(Pool(), TestService.Client)
span = mock.MagicMock()
child_span = span.make_child().__enter__()
child_span.trace_id = 1
child_span.parent_id = 1
child_span.id = 1
child_span.sampled = True
child_span.flags = None
edge_context = self.edge_context_factory.from_upstream(
SERIALIZED_EDGECONTEXT_WITH_VALID_AUTH)
edge_context.attach_context(child_span.context)
client = client_factory.make_object_for_context("test", span)
try:
client.example_simple()
except (TTransportException, EOFError):
pass # we don't have a test response for the client
itrans = TMemoryBuffer(client_memory_trans.getvalue())
iprot = THeaderProtocolFactory().getProtocol(itrans)
self.processor.process(iprot, self.oprot)
context, _ = self.observer.on_server_span_created.call_args[0]
try:
self.assertEqual(context.request_context.user.id, "t2_example")
self.assertEqual(context.request_context.user.roles, set())
self.assertEqual(context.request_context.user.is_logged_in, True)
self.assertEqual(context.request_context.user.loid, "t2_deadbeef")
self.assertEqual(context.request_context.user.cookie_created_ms,
100000)
self.assertEqual(context.request_context.oauth_client.id, None)
self.assertFalse(
context.request_context.oauth_client.is_type("third_party"))
self.assertEqual(context.request_context.session.id, "beefdead")
except jwt.exceptions.InvalidAlgorithmError:
raise unittest.SkipTest("cryptography is not installed")
def connection_made(self, transport):
assert self.transport is None, "Transport already instantiated here."
assert self.client is None, "Client already instantiated here."
# asyncio.Transport
self.transport = transport
# Thrift transport
self.thrift_transport = SenderTransport(self.transport, self.loop)
self.client = self._client_class(
self.thrift_transport, WrappedTransportFactory(self),
THeaderProtocolFactory(client_type=self.client_type))
def _create_server(self, port=0):
handler = CalculatorHandler()
transport = TServerSocket(port)
pfactory = THeaderProtocolFactory(client_types=THEADER_CLIENT_TYPES)
return TNonblockingServer(
Calculator.Processor(handler),
transport,
inputProtocolFactory=pfactory,
outputProtocolFactory=pfactory,
threads=1,
)
def __init__(self, processor: TProcessor, *args: Any, **kwargs: Any):
self.processor = processor
self.transport_factory = TBufferedTransportFactory()
self.protocol_factory = THeaderProtocolFactory(
# allow non-headerprotocol clients to talk with us
allowed_client_types=[
THeaderClientType.HEADERS,
THeaderClientType.FRAMED_BINARY,
THeaderClientType.UNFRAMED_BINARY,
])
super().__init__(*args, **kwargs)
def __init__(self, processor, *args, **kwargs):
self.processor = processor
self.transport_factory = TBufferedTransportFactory()
self.protocol_factory = THeaderProtocolFactory(
# allow non-headerprotocol clients to talk with us
client_types=[
THeaderTransport.HEADERS_CLIENT_TYPE,
THeaderTransport.FRAMED_DEPRECATED,
THeaderTransport.UNFRAMED_DEPRECATED,
],
)
super(GeventServer, self).__init__(*args, **kwargs)
def test_edge_request_headers(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans)
client_header_trans = client_prot.trans
client_header_trans.set_header(b"Edge-Request",
SERIALIZED_EDGECONTEXT_WITH_VALID_AUTH)
client_header_trans.set_header(b"Trace", b"1234")
client_header_trans.set_header(b"Parent", b"2345")
client_header_trans.set_header(b"Span", b"3456")
client_header_trans.set_header(b"Sampled", b"1")
client_header_trans.set_header(b"Flags", b"1")
client = TestService.Client(client_prot)
try:
client.example_simple()
except (TTransportException, EOFError):
pass # we don't have a test response for the client
itrans = TMemoryBuffer(client_memory_trans.getvalue())
iprot = THeaderProtocolFactory().getProtocol(itrans)
self.processor.process(iprot, self.oprot)
context, _ = self.observer.on_server_span_created.call_args[0]
try:
self.assertEqual(context.request_context.user.id, "t2_example")
self.assertEqual(context.request_context.user.roles, set())
self.assertEqual(context.request_context.user.is_logged_in, True)
self.assertEqual(context.request_context.user.loid, "t2_deadbeef")
self.assertEqual(context.request_context.user.cookie_created_ms,
100000)
self.assertEqual(context.request_context.oauth_client.id, None)
self.assertFalse(
context.request_context.oauth_client.is_type("third_party"))
self.assertEqual(context.request_context.session.id, "beefdead")
except jwt.exceptions.InvalidAlgorithmError:
raise unittest.SkipTest("cryptography is not installed")
def setUp(self):
self.otrans = TMemoryBuffer()
self.oprot = THeaderProtocolFactory().getProtocol(self.otrans)
self.observer = mock.Mock(spec=BaseplateObserver)
self.server_observer = mock.Mock(spec=ServerSpanObserver)
def _register_mock(context, server_span):
server_span.register(self.server_observer)
self.observer.on_server_span_created.side_effect = _register_mock
self.logger = mock.Mock(spec=logging.Logger)
mock_filewatcher = mock.Mock(spec=FileWatcher)
mock_filewatcher.get_data.return_value = {
"secrets": {
"secret/authentication/public-key": {
"type": "versioned",
"current": AUTH_TOKEN_PUBLIC_KEY,
},
},
"vault": {
"token": "test",
"url": "http://vault.example.com:8200/",
}
}
self.secrets = store.SecretsStore("/secrets")
self.secrets._filewatcher = mock_filewatcher
baseplate = Baseplate()
baseplate.register(self.observer)
self.edge_context_factory = EdgeRequestContextFactory(self.secrets)
handler = TestHandler()
processor = TestService.Processor(handler)
self.processor = baseplateify_processor(processor, self.logger,
baseplate,
self.edge_context_factory)
)
from thrift.transport.TTransport import TTransportBase, TTransportException
if six.PY3:
import asyncio
else:
# pyre-fixme[21]: Could not find module `trollius`.
import trollius as asyncio
# We support the deprecated FRAMED transport for old fb303
# clients that were otherwise failing miserably.
THEADER_CLIENT_TYPES = {
CLIENT_TYPE.HEADER,
CLIENT_TYPE.FRAMED_DEPRECATED,
}
_default_thpfactory = THeaderProtocolFactory(client_types=THEADER_CLIENT_TYPES)
THeaderProtocol = _default_thpfactory.getProtocol
logger = logging.getLogger(__name__)
if six.PY2:
class PermissionError(IOError):
pass
class ResourceWarning(RuntimeWarning):
pass
class TReadOnlyBuffer(TTransportBase):
"""Leaner version of TMemoryBuffer that is resettable."""
def main():
op = optparse.OptionParser(usage='%prog [options]', add_help_option=False)
op.add_option('-p',
'--port',
action='store',
type='int',
dest='port',
default=1234,
help='The server port')
op.add_option('-s',
'--servertype',
action='store',
type='string',
dest='servertype',
default='TNonblockingServer',
help='Type name of server')
op.add_option('-w',
'--num_workers',
action='store',
type='int',
dest='workers',
default=4,
help='Number of worker processes/threads')
op.add_option('-Q',
'--max_queue_size',
action='store',
type='int',
dest='max_queue_size',
default=0,
help='Max queue size, passed to TNonblockingServer')
op.add_option('-h',
'--header',
action='store_true',
help='Use the generated ContextIface')
op.add_option('-?',
'--help',
action='help',
help='Show this help message and exit')
(options, args) = op.parse_args()
if args:
op.error('trailing arguments: ' + ' '.join(args))
handler = LoadHandler()
processor = LoadTest.Processor(handler)
if options.header:
pfactory = THeaderProtocolFactory(True, [
CLIENT_TYPE.HEADER, CLIENT_TYPE.FRAMED_DEPRECATED,
CLIENT_TYPE.UNFRAMED_DEPRECATED, CLIENT_TYPE.HTTP_SERVER
])
if options.servertype == 'TCppServer':
print('C++ ThriftServer, Header transport, backwards compatible '
'with all other types')
elif options.servertype == 'TNonblockingServer':
print('Header transport, backwards compatible with framed')
else:
print('Header transport, backwards compatible with ' +
'unframed, framed, http')
else:
if options.servertype == 'TCppServer':
if not options.header:
op.error('TCppServer cannot be used without header')
elif options.servertype == 'TNonblockingServer':
print('Framed transport')
else:
print('Unframed transport')
pfactory = TBinaryProtocolAcceleratedFactory()
if options.servertype == 'TCppServer':
server = TCppServer.TCppServer(processor)
server.setPort(options.port)
print('Worker threads: ' + str(options.workers))
server.setNumIOWorkerThreads(options.workers)
else:
transport = TSocket.TServerSocket(options.port)
tfactory = TTransport.TBufferedTransportFactory()
if options.servertype == "TNonblockingServer":
server = TNonblockingServer.TNonblockingServer(
processor,
transport,
pfactory,
maxQueueSize=options.max_queue_size)
elif options.servertype == "TProcessPoolServer":
server = TProcessPoolServer.TProcessPoolServer(
processor, transport, tfactory, pfactory)
print('Worker processes: ' + str(options.workers))
server.setNumWorkers(options.workers)
else:
ServerClass = getattr(TServer, options.servertype)
server = ServerClass(processor, transport, tfactory, pfactory)
print('Serving ' + options.servertype + ' requests on port %d...' %
(options.port, ))
server.serve()
def main():
op = optparse.OptionParser(usage='%prog [options]', add_help_option=False)
op.add_option('-p',
'--port',
action='store',
type='int',
dest='port',
default=1234,
help='The server port')
op.add_option('-s',
'--servertype',
action='store',
type='string',
dest='servertype',
default='TGeventServer',
help='Type name of server')
op.add_option('-w',
'--num_workers',
action='store',
type='int',
dest='workers',
default=4,
help='Number of worker processes')
op.add_option('-Q',
'--max_queue_size',
action='store',
type='int',
dest='max_queue_size',
default=0,
help='Max queue size, passed to TNonblockingServer')
op.add_option('-h',
'--header',
action='store_true',
help='Use the generated ContextIface')
op.add_option('-?',
'--help',
action='help',
help='Show this help message and exit')
(options, args) = op.parse_args()
if args:
op.error('trailing arguments: ' + ' '.join(args))
handler = LoadHandler()
if options.servertype == 'TGeventServer':
# only import TGeventServer when necessary. TGeventServer calls
# monkey_patch, which breaks other servers
from apache.thrift.test.sync_load_handler import GeventLoadHandler
from thrift.server import TGeventServer
handler = GeventLoadHandler()
processor = LoadTest.Processor(handler)
if options.header:
pfactory = THeaderProtocolFactory(True, \
[THeaderTransport.HEADERS_CLIENT_TYPE,
THeaderTransport.FRAMED_DEPRECATED,
THeaderTransport.UNFRAMED_DEPRECATED,
THeaderTransport.HTTP_CLIENT_TYPE])
if options.servertype == 'TNonblockingServer':
print 'Header transport, backwards compatible with framed'
else:
print 'Header transport, backwards compatible with ' + \
'unframed, framed, http'
else:
if options.servertype == 'TNonblockingServer':
print 'Framed transport'
else:
print 'Unframed transport'
pfactory = TBinaryProtocolAcceleratedFactory()
transport = TSocket.TServerSocket(options.port)
tfactory = TTransport.TBufferedTransportFactory()
if options.servertype == "TNonblockingServer":
server = TNonblockingServer.TNonblockingServer(
processor,
transport,
pfactory,
maxQueueSize=options.max_queue_size)
elif options.servertype == "TProcessPoolServer":
server = TProcessPoolServer.TProcessPoolServer(processor, transport,
tfactory, pfactory)
print 'Worker processes: ' + str(options.workers)
server.setNumWorkers(options.workers)
elif options.servertype == "TGeventServer":
print 'Worker processes: ' + str(options.workers)
# Gevent makes its own server transport.
server = TGeventServer.TGeventServer(options.port, processor, None,
tfactory, pfactory)
server.setNumWorkers(options.workers)
else:
ServerClass = getattr(TServer, options.servertype)
server = ServerClass(processor, transport, tfactory, pfactory)
print 'Serving ' + options.servertype + \
' requests on port %d...' % (options.port,)
server.serve()
def __init__(self, processor, *args, **kwargs):
self.processor = processor
self.transport_factory = TBufferedTransportFactory()
self.protocol_factory = THeaderProtocolFactory()
super(GeventServer, self).__init__(*args, **kwargs)
def connection_made(self, transport):
self.transport = transport
self.client = self._client_class(SenderTransport(self.transport),
WrappedTransportFactory(),
THeaderProtocolFactory())
def main():
op = optparse.OptionParser(usage="%prog [options]", add_help_option=False)
op.add_option(
"-p",
"--port",
action="store",
type="int",
dest="port",
default=1234,
help="The server port",
)
op.add_option(
"-s",
"--servertype",
action="store",
type="string",
dest="servertype",
default="TCppServer",
help="Type name of server",
)
op.add_option(
"-w",
"--num_workers",
action="store",
type="int",
dest="workers",
default=4,
help="Number of worker processes/threads",
)
op.add_option(
"-h",
"--header",
action="store_true",
default=True,
help="Use the generated ContextIface",
)
op.add_option("-?",
"--help",
action="help",
help="Show this help message and exit")
(options, args) = op.parse_args()
if args:
op.error("trailing arguments: " + " ".join(args))
handler = LoadHandler()
processor = LoadTest.Processor(handler)
if options.header:
pfactory = THeaderProtocolFactory(
True,
[
CLIENT_TYPE.HEADER,
CLIENT_TYPE.FRAMED_DEPRECATED,
CLIENT_TYPE.UNFRAMED_DEPRECATED,
CLIENT_TYPE.HTTP_SERVER,
],
)
if options.servertype == "TCppServer":
print("C++ ThriftServer, Header transport, backwards compatible "
"with all other types")
else:
print("Header transport, backwards compatible with " +
"unframed, framed, http")
else:
if options.servertype == "TCppServer":
if not options.header:
op.error("TCppServer cannot be used without header")
else:
print("Unframed transport")
pfactory = TBinaryProtocolAcceleratedFactory()
if options.servertype == "TCppServer":
server = TCppServer.TCppServer(processor)
server.setPort(options.port)
print("Worker threads: " + str(options.workers))
server.setNumIOWorkerThreads(options.workers)
else:
transport = TSocket.TServerSocket(options.port)
tfactory = TTransport.TBufferedTransportFactory()
ServerClass = getattr(TServer, options.servertype)
server = ServerClass(processor, transport, tfactory, pfactory)
print("Serving " + options.servertype + " requests on port %d..." %
(options.port, ))
server.serve()
