# Copyright (c) Meta Platforms, Inc. and affiliates. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # pyre-unsafe from __future__ import absolute_import from __future__ import division from __future__ import print_function from __future__ import unicode_literals import http.server as http_server import zlib from io import BytesIO as StringIO from struct import pack, unpack from thrift.protocol.TBinaryProtocol import TBinaryProtocol from thrift.protocol.TCompactProtocol import getVarint, readVarint, TCompactProtocol from thrift.Thrift import TApplicationException from thrift.transport.TTransport import ( TTransportException, TTransportBase, CReadableTransport, ) xrange = range # Import the snappy module if it is available try: import snappy except ImportError: # If snappy is not available, don't fail immediately. # Only raise an error if we actually ever need to perform snappy # compression. class DummySnappy(object): def compress(self, buf): raise TTransportException( TTransportException.INVALID_TRANSFORM, "snappy module not available" ) def decompress(self, buf): raise TTransportException( TTransportException.INVALID_TRANSFORM, "snappy module not available" ) snappy = DummySnappy() # Import the zstd module if it is available try: import zstd # @manual except ImportError: # If zstd is not available, don't fail immediately. # Only raise an error if we actually ever need to perform zstd # compression. class DummyZstd(object): def ZstdCompressor(self, write_content_size): raise TTransportException( TTransportException.INVALID_TRANSFORM, "zstd module not available" ) def ZstdDecompressor(self): raise TTransportException( TTransportException.INVALID_TRANSFORM, "zstd module not available" ) zstd = DummyZstd() # Definitions from THeader.h class CLIENT_TYPE: HEADER = 0 FRAMED_DEPRECATED = 1 UNFRAMED_DEPRECATED = 2 HTTP_SERVER = 3 HTTP_CLIENT = 4 FRAMED_COMPACT = 5 HTTP_GET = 7 UNKNOWN = 8 UNFRAMED_COMPACT_DEPRECATED = 9 class HEADER_FLAG: SUPPORT_OUT_OF_ORDER = 0x01 DUPLEX_REVERSE = 0x08 class TRANSFORM: NONE = 0x00 ZLIB = 0x01 HMAC = 0x02 SNAPPY = 0x03 QLZ = 0x04 ZSTD = 0x05 class INFO: NORMAL = 1 PERSISTENT = 2 T_BINARY_PROTOCOL = 0 T_COMPACT_PROTOCOL = 2 HEADER_MAGIC = 0x0FFF0000 PACKED_HEADER_MAGIC: bytes = pack(b"!H", HEADER_MAGIC >> 16) HEADER_MASK = 0xFFFF0000 FLAGS_MASK = 0x0000FFFF HTTP_SERVER_MAGIC = 0x504F5354 # POST HTTP_CLIENT_MAGIC = 0x48545450 # HTTP HTTP_GET_CLIENT_MAGIC = 0x47455420 # GET HTTP_HEAD_CLIENT_MAGIC = 0x48454144 # HEAD BIG_FRAME_MAGIC = 0x42494746 # BIGF MAX_FRAME_SIZE = 0x3FFFFFFF MAX_BIG_FRAME_SIZE = 2 ** 61 - 1 class THeaderTransport(TTransportBase, CReadableTransport): """Transport that sends headers. Also understands framed/unframed/HTTP transports and will do the right thing""" __max_frame_size = MAX_FRAME_SIZE # Defaults to current user, but there is also a setter below. __identity = None __first_request = True IDENTITY_HEADER = "identity" ID_VERSION_HEADER = "id_version" ID_VERSION = "1" CLIENT_METADATA_HEADER = "client_metadata" def __init__(self, trans, client_types=None, client_type=None): self.__trans = trans self.__rbuf = StringIO() self.__rbuf_frame = False self.__wbuf = StringIO() self.seq_id = 0 self.__flags = 0 self.__read_transforms = [] self.__write_transforms = [] self.__supported_client_types = set(client_types or (CLIENT_TYPE.HEADER,)) self.__proto_id = T_COMPACT_PROTOCOL # default to compact like c++ self.__client_type = client_type or CLIENT_TYPE.HEADER self.__read_headers = {} self.__read_persistent_headers = {} self.__write_headers = {} self.__write_persistent_headers = {} if self.__client_type in ( CLIENT_TYPE.UNFRAMED_DEPRECATED, CLIENT_TYPE.UNFRAMED_COMPACT_DEPRECATED, CLIENT_TYPE.FRAMED_DEPRECATED, CLIENT_TYPE.FRAMED_COMPACT, ): self.__client_type = CLIENT_TYPE.HEADER self.__supported_client_types.add(self.__client_type) # If we support unframed binary / framed binary also support compact if CLIENT_TYPE.UNFRAMED_DEPRECATED in self.__supported_client_types: self.__supported_client_types.add(CLIENT_TYPE.UNFRAMED_COMPACT_DEPRECATED) if CLIENT_TYPE.FRAMED_DEPRECATED in self.__supported_client_types: self.__supported_client_types.add(CLIENT_TYPE.FRAMED_COMPACT) def set_header_flag(self, flag): self.__flags |= flag def clear_header_flag(self, flag): self.__flags &= ~flag def header_flags(self): return self.__flags def set_max_frame_size(self, size): if size > MAX_BIG_FRAME_SIZE: raise TTransportException( TTransportException.INVALID_FRAME_SIZE, "Cannot set max frame size > %s" % MAX_BIG_FRAME_SIZE, ) if size > MAX_FRAME_SIZE and self.__client_type != CLIENT_TYPE.HEADER: raise TTransportException( TTransportException.INVALID_FRAME_SIZE, "Cannot set max frame size > %s for clients other than HEADER" % MAX_FRAME_SIZE, ) self.__max_frame_size = size def get_peer_identity(self): if self.IDENTITY_HEADER in self.__read_headers: if self.__read_headers[self.ID_VERSION_HEADER] == self.ID_VERSION: return self.__read_headers[self.IDENTITY_HEADER] return None def set_identity(self, identity): self.__identity = identity def get_protocol_id(self): return self.__proto_id def set_protocol_id(self, proto_id): self.__proto_id = proto_id def set_header(self, str_key, str_value): self.__write_headers[str_key] = str_value def get_write_headers(self): return self.__write_headers def get_headers(self): return self.__read_headers def clear_headers(self): self.__write_headers.clear() def set_persistent_header(self, str_key, str_value): self.__write_persistent_headers[str_key] = str_value def get_write_persistent_headers(self): return self.__write_persistent_headers def clear_persistent_headers(self): self.__write_persistent_headers.clear() def add_transform(self, trans_id): self.__write_transforms.append(trans_id) def _reset_protocol(self): # HTTP calls that are one way need to flush here. if self.__client_type == CLIENT_TYPE.HTTP_SERVER: self.flush() # set to anything except unframed self.__client_type = CLIENT_TYPE.UNKNOWN # Read header bytes to check which protocol to decode self.readFrame(0) def getTransport(self): return self.__trans def isOpen(self): return self.getTransport().isOpen() def open(self): return self.getTransport().open() def close(self): return self.getTransport().close() def read(self, sz): ret = self.__rbuf.read(sz) if len(ret) == sz: return ret if self.__client_type in ( CLIENT_TYPE.UNFRAMED_DEPRECATED, CLIENT_TYPE.UNFRAMED_COMPACT_DEPRECATED, ): return ret + self.getTransport().readAll(sz - len(ret)) self.readFrame(sz - len(ret)) return ret + self.__rbuf.read(sz - len(ret)) readAll = read # TTransportBase.readAll does a needless copy here. def readFrame(self, req_sz): self.__rbuf_frame = True word1 = self.getTransport().readAll(4) sz = unpack("!I", word1)[0] proto_id = word1[0] if proto_id == TBinaryProtocol.PROTOCOL_ID: # unframed self.__client_type = CLIENT_TYPE.UNFRAMED_DEPRECATED self.__proto_id = T_BINARY_PROTOCOL if req_sz <= 4: # check for reads < 0. self.__rbuf = StringIO(word1) else: self.__rbuf = StringIO(word1 + self.getTransport().read(req_sz - 4)) elif proto_id == TCompactProtocol.PROTOCOL_ID: self.__client_type = CLIENT_TYPE.UNFRAMED_COMPACT_DEPRECATED self.__proto_id = T_COMPACT_PROTOCOL if req_sz <= 4: # check for reads < 0. self.__rbuf = StringIO(word1) else: self.__rbuf = StringIO(word1 + self.getTransport().read(req_sz - 4)) elif sz == HTTP_SERVER_MAGIC: self.__client_type = CLIENT_TYPE.HTTP_SERVER mf = self.getTransport().handle.makefile("rb", -1) self.handler = RequestHandler(mf, "client_address:port", "") self.header = self.handler.wfile self.__rbuf = StringIO(self.handler.data) else: if sz == BIG_FRAME_MAGIC: sz = unpack("!Q", self.getTransport().readAll(8))[0] # could be header format or framed. Check next two bytes. magic = self.getTransport().readAll(2) proto_id = magic[0] if proto_id == TCompactProtocol.PROTOCOL_ID: self.__client_type = CLIENT_TYPE.FRAMED_COMPACT self.__proto_id = T_COMPACT_PROTOCOL _frame_size_check(sz, self.__max_frame_size, header=False) self.__rbuf = StringIO(magic + self.getTransport().readAll(sz - 2)) elif proto_id == TBinaryProtocol.PROTOCOL_ID: self.__client_type = CLIENT_TYPE.FRAMED_DEPRECATED self.__proto_id = T_BINARY_PROTOCOL _frame_size_check(sz, self.__max_frame_size, header=False) self.__rbuf = StringIO(magic + self.getTransport().readAll(sz - 2)) elif magic == PACKED_HEADER_MAGIC: self.__client_type = CLIENT_TYPE.HEADER _frame_size_check(sz, self.__max_frame_size) # flags(2), seq_id(4), header_size(2) n_header_meta = self.getTransport().readAll(8) self.__flags, self.seq_id, header_size = unpack("!HIH", n_header_meta) data = StringIO() data.write(magic) data.write(n_header_meta) data.write(self.getTransport().readAll(sz - 10)) data.seek(10) self.read_header_format(sz - 10, header_size, data) else: self.__client_type = CLIENT_TYPE.UNKNOWN raise TTransportException( TTransportException.INVALID_CLIENT_TYPE, "Could not detect client transport type", ) if self.__client_type not in self.__supported_client_types: raise TTransportException( TTransportException.INVALID_CLIENT_TYPE, "Client type {} not supported on server".format(self.__client_type), ) def read_header_format(self, sz, header_size, data): # clear out any previous transforms self.__read_transforms = [] header_size = header_size * 4 if header_size > sz: raise TTransportException( TTransportException.INVALID_FRAME_SIZE, "Header size is larger than frame", ) end_header = header_size + data.tell() self.__proto_id = readVarint(data) num_headers = readVarint(data) if self.__proto_id == 1 and self.__client_type != CLIENT_TYPE.HTTP_SERVER: raise TTransportException( TTransportException.INVALID_CLIENT_TYPE, "Trying to recv JSON encoding over binary", ) # Read the headers. Data for each header varies. for _ in range(0, num_headers): trans_id = readVarint(data) if trans_id in (TRANSFORM.ZLIB, TRANSFORM.SNAPPY, TRANSFORM.ZSTD): self.__read_transforms.insert(0, trans_id) elif trans_id == TRANSFORM.HMAC: raise TApplicationException( TApplicationException.INVALID_TRANSFORM, "Hmac transform is no longer supported: %i" % trans_id, ) else: # TApplicationException will be sent back to client raise TApplicationException( TApplicationException.INVALID_TRANSFORM, "Unknown transform in client request: %i" % trans_id, ) # Clear out previous info headers. self.__read_headers.clear() # Read the info headers. while data.tell() < end_header: info_id = readVarint(data) if info_id == INFO.NORMAL: _read_info_headers(data, end_header, self.__read_headers) elif info_id == INFO.PERSISTENT: _read_info_headers(data, end_header, self.__read_persistent_headers) else: break # Unknown header. Stop info processing. if self.__read_persistent_headers: self.__read_headers.update(self.__read_persistent_headers) # Skip the rest of the header data.seek(end_header) payload = data.read(sz - header_size) # Read the data section. self.__rbuf = StringIO(self.untransform(payload)) def write(self, buf): self.__wbuf.write(buf) def transform(self, buf): for trans_id in self.__write_transforms: if trans_id == TRANSFORM.ZLIB: buf = zlib.compress(buf) elif trans_id == TRANSFORM.SNAPPY: buf = snappy.compress(buf) elif trans_id == TRANSFORM.ZSTD: buf = zstd.ZstdCompressor(write_content_size=True).compress(buf) else: raise TTransportException( TTransportException.INVALID_TRANSFORM, "Unknown transform during send", ) return buf def untransform(self, buf): for trans_id in self.__read_transforms: if trans_id == TRANSFORM.ZLIB: buf = zlib.decompress(buf) elif trans_id == TRANSFORM.SNAPPY: buf = snappy.decompress(buf) elif trans_id == TRANSFORM.ZSTD: buf = zstd.ZstdDecompressor().decompress(buf) if trans_id not in self.__write_transforms: self.__write_transforms.append(trans_id) return buf def disable_client_metadata(self): self.__first_request = False def flush(self): self.flushImpl(False) def onewayFlush(self): self.flushImpl(True) def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for CLIENT_TYPE.HEADER @param buf(StringIO): Buffer to write message to @param wout(str): Payload @param wsz(int): Payload length """ transform_data = StringIO() # For now, all transforms don't require data. num_transforms = len(self.__write_transforms) for trans_id in self.__write_transforms: transform_data.write(getVarint(trans_id)) # Add in special flags. if self.__identity: self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION self.__write_headers[self.IDENTITY_HEADER] = self.__identity if self.__first_request: self.__first_request = False self.__write_headers[ self.CLIENT_METADATA_HEADER ] = '{"agent":"THeaderTransport.py"}' info_data = StringIO() # Write persistent kv-headers _flush_info_headers( info_data, self.get_write_persistent_headers(), INFO.PERSISTENT ) # Write non-persistent kv-headers _flush_info_headers(info_data, self.__write_headers, INFO.NORMAL) header_data = StringIO() header_data.write(getVarint(self.__proto_id)) header_data.write(getVarint(num_transforms)) header_size = transform_data.tell() + header_data.tell() + info_data.tell() padding_size = 4 - (header_size % 4) header_size = header_size + padding_size # MAGIC(2) | FLAGS(2) + SEQ_ID(4) + HEADER_SIZE(2) wsz += header_size + 10 if wsz > MAX_FRAME_SIZE: buf.write(pack("!I", BIG_FRAME_MAGIC)) buf.write(pack("!Q", wsz)) else: buf.write(pack("!I", wsz)) buf.write(pack("!HH", HEADER_MAGIC >> 16, self.__flags)) buf.write(pack("!I", self.seq_id)) buf.write(pack("!H", header_size // 4)) buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) buf.write(info_data.getvalue()) # Pad out the header with 0x00 for _ in range(0, padding_size, 1): buf.write(pack("!c", b"\0")) # Send data section buf.write(wout) def flushImpl(self, oneway): wout = self.__wbuf.getvalue() wout = self.transform(wout) wsz = len(wout) # reset wbuf before write/flush to preserve state on underlying failure self.__wbuf.seek(0) self.__wbuf.truncate() if self.__proto_id == 1 and self.__client_type != CLIENT_TYPE.HTTP_SERVER: raise TTransportException( TTransportException.INVALID_CLIENT_TYPE, "Trying to send JSON encoding over binary", ) buf = StringIO() if self.__client_type == CLIENT_TYPE.HEADER: self._flushHeaderMessage(buf, wout, wsz) elif self.__client_type in ( CLIENT_TYPE.FRAMED_DEPRECATED, CLIENT_TYPE.FRAMED_COMPACT, ): buf.write(pack("!i", wsz)) buf.write(wout) elif self.__client_type in ( CLIENT_TYPE.UNFRAMED_DEPRECATED, CLIENT_TYPE.UNFRAMED_COMPACT_DEPRECATED, ): buf.write(wout) elif self.__client_type == CLIENT_TYPE.HTTP_SERVER: # Reset the client type if we sent something - # oneway calls via HTTP expect a status response otherwise buf.write(self.header.getvalue()) buf.write(wout) self.__client_type == CLIENT_TYPE.HEADER elif self.__client_type == CLIENT_TYPE.UNKNOWN: raise TTransportException( TTransportException.INVALID_CLIENT_TYPE, "Unknown client type" ) # We don't include the framing bytes as part of the frame size check frame_size = buf.tell() - (4 if wsz < MAX_FRAME_SIZE else 12) _frame_size_check( frame_size, self.__max_frame_size, header=self.__client_type == CLIENT_TYPE.HEADER, ) self.getTransport().write(buf.getvalue()) if oneway: self.getTransport().onewayFlush() else: self.getTransport().flush() # Implement the CReadableTransport interface. @property def cstringio_buf(self): if not self.__rbuf_frame: self.readFrame(0) return self.__rbuf def cstringio_refill(self, prefix, reqlen): # self.__rbuf will already be empty here because fastproto doesn't # ask for a refill until the previous buffer is empty. Therefore, # we can start reading new frames immediately. # On unframed clients, there is a chance there is something left # in rbuf, and the read pointer is not advanced by fastproto # so seek to the end to be safe self.__rbuf.seek(0, 2) while len(prefix) < reqlen: prefix += self.read(reqlen) self.__rbuf = StringIO(prefix) return self.__rbuf def _serialize_string(str_): if not isinstance(str_, bytes): str_ = str_.encode() return getVarint(len(str_)) + str_ def _flush_info_headers(info_data, write_headers, type) -> None: if len(write_headers) > 0: info_data.write(getVarint(type)) info_data.write(getVarint(len(write_headers))) write_headers_iter = write_headers.items() for str_key, str_value in write_headers_iter: info_data.write(_serialize_string(str_key)) info_data.write(_serialize_string(str_value)) write_headers.clear() def _read_string(bufio, buflimit): str_sz = readVarint(bufio) if str_sz + bufio.tell() > buflimit: raise TTransportException( TTransportException.INVALID_FRAME_SIZE, "String read too big" ) return bufio.read(str_sz) def _read_info_headers(data, end_header, read_headers) -> None: num_keys = readVarint(data) for _ in xrange(num_keys): str_key = _read_string(data, end_header) str_value = _read_string(data, end_header) read_headers[str_key] = str_value def _frame_size_check(sz, set_max_size, header: bool = True) -> None: if sz > set_max_size or (not header and sz > MAX_FRAME_SIZE): raise TTransportException( TTransportException.INVALID_FRAME_SIZE, "%s transport frame was too large" % "Header" if header else "Framed", ) class RequestHandler(http_server.BaseHTTPRequestHandler): # Same as superclass function, but append 'POST' because we # stripped it in the calling function. Would be nice if # we had an ungetch instead def handle_one_request(self): self.raw_requestline = self.rfile.readline() if not self.raw_requestline: self.close_connection = 1 return self.raw_requestline = "POST" + self.raw_requestline if not self.parse_request(): # An error code has been sent, just exit return mname = "do_" + self.command if not hasattr(self, mname): self.send_error(501, "Unsupported method (%r)" % self.command) return method = getattr(self, mname) method() def setup(self): self.rfile = self.request self.wfile = StringIO() # New output buffer def finish(self): if not self.rfile.closed: self.rfile.close() # leave wfile open for reading. def do_POST(self): if int(self.headers["Content-Length"]) > 0: self.data = self.rfile.read(int(self.headers["Content-Length"])) else: self.data = "" # Prepare a response header, to be sent later. self.send_response(200) self.send_header("content-type", "application/x-thrift") self.end_headers()