<commit_before>from ethereum.abi import ContractTranslator
import serpent
import requests
import json

class EthJsonRpc:

    def __init__(self, host, port, contract_code=None, contract_address=None):
        self.host = host
        self.port = port
        self.contract_code = None
        self.signature = None
        self.translation = None
        self.contract_address = contract_address
        self.update_code(contract_code)

    def update_code(self, contract_code):
        if contract_code:
            self.contract_code = contract_code
            self.signature = serpent.mk_full_signature(contract_code)
            self.translation = ContractTranslator(self.signature)

    def _call(self, method, params=None, _id=0):
        if params is None:
            params = []
        data = json.dumps({
            'jsonrpc': '2.0',
            'method': method,
            'params': params,
            'id': _id
        })
        response = requests.post("http://{}:{}".format(self.host, self.port), data=data).json()
        if 'result' in response:
            return response['result']
        else:
            raise RuntimeError('Error from RPC call. Returned payload: {0}'.format(response))

    def create_contract(self, contract_code, value=0, from_address=None, gas=0, gas_price=0):
        self.update_code(contract_code)
        byte_code = serpent.compile(contract_code)
        self.contract_address = self.eth_sendTransaction(data=byte_code, value=value, from_address=from_address, gas=gas, gas_price=gas_price)
        return self.contract_address

    def eth_sendTransaction(self, to_address=None, function_name=None, data=None, value=0, from_address=None, gas=0, gas_price=0, code=None):
        """
        Creates new message call transaction or a contract creation, if the data field contains code.
        """
        if code:
            self.update_code(code)
        else:
            self.update_code(self.eth_getCode(to_address))
        if function_name:
            if data is None:
                data = []
            data = self.translation.encode(function_name, data)
        params = {
            'from': from_address,
            'to': to_address,
            'gas': hex(gas) if gas else None,
            'gasPrice': hex(gas_price) if gas_price else None,
            'value': hex(value) if value else None,
            'data': '0x{}'.format(data.encode('hex')) if data else None
        }
        return self._call('eth_sendTransaction', [params])

    def eth_call(self, to_address, function_name, data=None, code=None, default_block="latest"):
        """
        Executes a new message call immediately without creating a transaction on the block chain.
        """
        if code:
            self.update_code(code)
        else:
            self.update_code(self.eth_getCode(to_address))
        if data is None:
            data = []
        data = self.translation.encode(function_name, data)
        params = [
            {
                'to': to_address,
                'data': '0x{}'.format(data.encode('hex'))
            },
            default_block
        ]
        response = self._call('eth_call', params)
        if function_name:
            response = self.translation.decode(function_name, response[2:].decode('hex'))
        return response

    def web3_clientVersion(self):
        """
        Returns the current client version.
        """
        return self._call('web3_clientVersion')

    def web3_sha3(self, data):
        """
        Returns SHA3 of the given data.
        """
        data = str(data).encode('hex')
        return self._call('web3_sha3', [data])

    def net_version(self):
        """
        Returns the current network protocol version.
        """
        return self._call('net_version')

    def net_listening(self):
        """
        Returns true if client is actively listening for network connections.
        """
        return self._call('net_listening')

    def net_peerCount(self):
        """
        Returns number of peers currenly connected to the client.
        """
        return self._call('net_peerCount')

    def eth_version(self):
        """
        Returns the current ethereum protocol version.
        """
        return self._call('eth_version')

    def eth_coinbase(self):
        """
        Returns the client coinbase address.
        """
        return self._call('eth_coinbase')

    def eth_mining(self):
        """
        Returns true if client is actively mining new blocks.
        """
        return self._call('eth_mining')

    def eth_gasPrice(self):
        """
        Returns the current price per gas in wei.
        """
        return self._call('eth_gasPrice')

    def eth_accounts(self):
        """
        Returns a list of addresses owned by client.
        """
        return self._call('eth_accounts')

    def eth_blockNumber(self):
        """
        Returns the number of most recent block.
        """
        return self._call('eth_blockNumber')

    def eth_getBalance(self, address, default_block="latest"):
        """
        Returns the balance of the account of given address.
        """
        return self._call('eth_getBalance', [address, default_block])

    def eth_getStorageAt(self, address, position, default_block="latest"):
        """
        Returns the value from a storage position at a given address.
        """
        return self._call('eth_getStorageAt', [address, hex(position), default_block])

    def eth_getTransactionCount(self, address, default_block="latest"):
        """
        Returns the number of transactions send from a address.
        """
        return self._call('eth_getTransactionCount', [address, default_block])

    def eth_getBlockTransactionCountByHash(self, block_hash):
        """
        Returns the number of transactions in a block from a block matching the given block hash.
        """
        return self._call('eth_getTransactionCount', [block_hash])

    def eth_getBlockTransactionCountByNumber(self, block_number):
        """
        Returns the number of transactions in a block from a block matching the given block number.
        """
        return self._call('eth_getBlockTransactionCountByNumber', [hex(block_number)])

    def eth_getUncleCountByblockHash(self, block_hash):
        """
        Returns the number of uncles in a block from a block matching the given block hash.
        """
        return self._call('eth_getUncleCountByblockHash', [block_hash])

    def eth_getUncleCountByblockNumber(self, block_number):
        """
        Returns the number of uncles in a block from a block matching the given block number.
        """
        return self._call('eth_getUncleCountByblockNumber', [hex(block_number)])

    def eth_getCode(self, address, default_block="latest"):
        """
        Returns code at a given address.
        """
        return self._call('eth_getCode', [address, default_block])

    def eth_getBlockByHash(self, block_hash, transaction_objects=True):
        """
        Returns information about a block by hash.
        """
        return self._call('eth_getBlockByHash', [block_hash, transaction_objects])

    def eth_flush(self):
        """
        """
        return self._call('eth_flush')

    def eth_getBlockByNumber(self, block_number, transaction_objects=True):
        """
        Returns information about a block by hash.
        """
        return self._call('eth_getBlockByNumber', [block_number, transaction_objects])

    def eth_getTransactionByHash(self, transaction_hash):
        """
        Returns the information about a transaction requested by transaction hash.
        """
        return self._call('eth_getTransactionByHash', [transaction_hash])

    def eth_getTransactionByblockHashAndIndex(self, block_hash, index):
        """
        Returns information about a transaction by block hash and transaction index position.
        """
        return self._call('eth_getTransactionByblock_hashAndIndex', [block_hash, hex(index)])

    def eth_getTransactionByblockNumberAndIndex(self, block_number, index):
        """
        Returns information about a transaction by block number and transaction index position.
        """
        return self._call('eth_getTransactionByblock_numberAndIndex', [block_number, hex(index)])

    def eth_getUncleByblockHashAndIndex(self, block_hash, index, transaction_objects=True):
        """
        Returns information about a uncle of a block by hash and uncle index position.
        """
        return self._call('eth_getUncleByblock_hashAndIndex', [block_hash, hex(index), transaction_objects])

    def eth_getUncleByblockNumberAndIndex(self, block_number, index, transaction_objects=True):
        """
        Returns information about a uncle of a block by number and uncle index position.
        """
        return self._call('eth_getUncleByblock_numberAndIndex', [block_number, hex(index), transaction_objects])

    def eth_getCompilers(self):
        """
        Returns a list of available compilers in the client.
        """
        return self._call('eth_getCompilers')

    def eth_compileSolidity(self, code):
        """
        Returns compiled solidity code.
        """
        return self._call('eth_compileSolidity', [code])

    def eth_compileLLL(self, code):
        """
        Returns compiled LLL code.
        """
        return self._call('eth_compileLLL', [code])

    def eth_compileSerpent(self, code):
        """
        Returns compiled serpent code.
        """
        return self._call('eth_compileSerpent', [code])

    def eth_newFilter(self, from_block="latest", to_block="latest", address=None, topics=None):
        """
        Creates a filter object, based on filter options, to notify when the state changes (logs).
        To check if the state has changed, call eth_getFilterChanges.
        """
        _filter = {
            'fromBlock': from_block,
            'toBlock': to_block,
            'address': address,
            'topics': topics
        }
        return self._call('eth_newFilter', [_filter])

    def eth_newBlockFilter(self, default_block="latest"):
        """
        Creates a filter object, based on an option string, to notify when state changes (logs). To check if the state has changed, call eth_getFilterChanges.
        """
        return self._call('eth_newBlockFilter', [default_block])

    def eth_uninstallFilter(self, filter_id):
        """
        Uninstalls a filter with given id. Should always be called when watch is no longer needed. Additonally Filters timeout when they aren't requested with eth_getFilterChanges for a period of time.
        """
        return self._call('eth_uninstallFilter', [hex(filter_id)])

    def eth_getFilterChanges(self, filter_id):
        """
        Polling method for a filter, which returns an array of logs which occurred since last poll.
        """
        return self._call('eth_getFilterChanges', [hex(filter_id)])

    def eth_getFilterLogs(self, filter_id):
        """
        Returns an array of all logs matching filter with given id.
        """
        return self._call('eth_getFilterLogs', [hex(filter_id)])

    def eth_getLogs(self, filter_object):
        """
        Returns an array of all logs matching a given filter object.
        """
        return self._call('eth_getLogs', [filter_object])

    def eth_getWork(self):
        """
        Returns the hash of the current block, the seedHash, and the difficulty to be met.
        """
        return self._call('eth_getWork')

    def eth_submitWork(self, nonce, header, mix_digest):
        """
        Used for submitting a proof-of-work solution.
        """
        return self._call('eth_submitWork', [nonce, header, mix_digest])

    def db_putString(self, database_name, key_name, string):
        """
        Stores a string in the local database.
        """
        return self._call('db_putString', [database_name, key_name, string])

    def db_getString(self, database_name, key_name):
        """
        Stores a string in the local database.
        """
        return self._call('db_getString', [database_name, key_name])

    def db_putHex(self, database_name, key_name, string):
        """
        Stores binary data in the local database.
        """
        return self._call('db_putHex', [database_name, key_name, string.encode('hex')])

    def db_getHex(self, database_name, key_name):
        """
        Returns binary data from the local database.
        """
        return self._call('db_getString', [database_name, key_name]).decode('hex')

    def shh_version(self):
        """
        Returns the current whisper protocol version.
        """
        return self._call('shh_version')

    def shh_post(self, topics, payload, priority, ttl, _from=None, to=None):
        """
        Sends a whisper message.
        """
        whisper_object = {
            'from': _from,
            'to': to,
            'topics': topics,
            'payload': payload,
            'priority': priority,
            'ttl': ttl
        }
        return self._call('shh_post', [whisper_object])

    def shh_newIdentinty(self):
        """
        Creates new whisper identity in the client.
        """
        return self._call('shh_newIdentinty')

    def shh_hasIdentity(self, address):
        """
        Checks if the client hold the private keys for a given identity.
        """
        return self._call('shh_hasIdentity', [address])

    def shh_newGroup(self):
        """
        """
        return self._call('shh_hasIdentity')

    def shh_addToGroup(self):
        """
        """
        return self._call('shh_addToGroup')

    def shh_newFilter(self, to, topics):
        """
        Creates filter to notify, when client receives whisper message matching the filter options.
        """
        _filter = {
            'to': to,
            'topics': topics
        }
        return self._call('shh_newFilter', [_filter])

    def shh_uninstallFilter(self, filter_id):
        """
        Uninstalls a filter with given id. Should always be called when watch is no longer needed.
        Additonally Filters timeout when they aren't requested with shh_getFilterChanges for a period of time.
        """
        return self._call('shh_uninstallFilter', [hex(filter_id)])

    def shh_getFilterChanges(self, filter_id):
        """
        Polling method for whisper filters.
        """
        return self._call('shh_getFilterChanges', [hex(filter_id)])

    def shh_getMessages(self, filter_id):
        """
        Get all messages matching a filter, which are still existing in the node.
        """
        return self._call('shh_getMessages', [hex(filter_id)])
<commit_msg>Removed some weird stuff.<commit_after>from ethereum.abi import ContractTranslator
import serpent
import requests
import json

class EthJsonRpc:

    def __init__(self, host, port, contract_code=None, contract_address=None):
        self.host = host
        self.port = port
        self.contract_code = None
        self.signature = None
        self.translation = None
        self.contract_address = contract_address
        self.update_code(contract_code)

    def update_code(self, contract_code):
        if contract_code:
            self.contract_code = contract_code
            self.signature = serpent.mk_full_signature(contract_code)
            self.translation = ContractTranslator(self.signature)

    def _call(self, method, params=None, _id=0):
        if params is None:
            params = []
        data = json.dumps({
            'jsonrpc': '2.0',
            'method': method,
            'params': params,
            'id': _id
        })
        response = requests.post("http://{}:{}".format(self.host, self.port), data=data).json()
        if 'result' in response:
            return response['result']
        else:
            raise RuntimeError('Error from RPC call. Returned payload: {0}'.format(response))

    def create_contract(self, contract_code, value=0, from_address=None, gas=0, gas_price=0):
        self.update_code(contract_code)
        byte_code = serpent.compile(contract_code)
        if from_address is None:
            from_address = self.eth_accounts()[0]

        self.contract_address = self.eth_sendTransaction(data=byte_code, value=value, from_address=from_address, gas=gas, gas_price=gas_price)
        return self.contract_address

    def eth_sendTransaction(self, to_address=None, function_name=None, data=None, value=0, from_address=None, gas=0, gas_price=0, code=None):
        """
        Creates new message call transaction or a contract creation, if the data field contains code.
        """
        # if code:
        #     self.update_code(code)
        # else:
        #     self.update_code(self.eth_getCode('0x' + to_address.encode('hex')))
        if function_name:
            if data is None:
                data = []
            data = self.translation.encode(function_name, data)
        params = {
            'from': from_address,
            'to': to_address,
            'gas': hex(gas) if gas else None,
            'gasPrice': hex(gas_price) if gas_price else None,
            'value': hex(value) if value else None,
            'data': '0x{}'.format(data.encode('hex')) if data else None
        }
        return self._call('eth_sendTransaction', [params])

    def eth_call(self, to_address, function_name, data=None, code=None, default_block="latest"):
        """
        Executes a new message call immediately without creating a transaction on the block chain.
        """
        #if code:
        #    self.update_code(code)
        #else:
        #    self.update_code(self.eth_getCode('0x' + to_address.encode('hex')))
        if data is None:
            data = []
        data = self.translation.encode(function_name, data)
        params = [
            {
                'to': to_address,
                'data': '0x{}'.format(data.encode('hex'))
            },
            default_block
        ]
        response = self._call('eth_call', params)
        if function_name:
            response = self.translation.decode(function_name, response[2:].decode('hex'))
        return response

    def web3_clientVersion(self):
        """
        Returns the current client version.
        """
        return self._call('web3_clientVersion')

    def web3_sha3(self, data):
        """
        Returns SHA3 of the given data.
        """
        data = str(data).encode('hex')
        return self._call('web3_sha3', [data])

    def net_version(self):
        """
        Returns the current network protocol version.
        """
        return self._call('net_version')

    def net_listening(self):
        """
        Returns true if client is actively listening for network connections.
        """
        return self._call('net_listening')

    def net_peerCount(self):
        """
        Returns number of peers currenly connected to the client.
        """
        return self._call('net_peerCount')

    def eth_version(self):
        """
        Returns the current ethereum protocol version.
        """
        return self._call('eth_version')

    def eth_coinbase(self):
        """
        Returns the client coinbase address.
        """
        return self._call('eth_coinbase')

    def eth_mining(self):
        """
        Returns true if client is actively mining new blocks.
        """
        return self._call('eth_mining')

    def eth_gasPrice(self):
        """
        Returns the current price per gas in wei.
        """
        return self._call('eth_gasPrice')

    def eth_accounts(self):
        """
        Returns a list of addresses owned by client.
        """
        return self._call('eth_accounts')

    def eth_blockNumber(self):
        """
        Returns the number of most recent block.
        """
        return self._call('eth_blockNumber')

    def eth_getBalance(self, address, default_block="latest"):
        """
        Returns the balance of the account of given address.
        """
        return self._call('eth_getBalance', [address, default_block])

    def eth_getStorageAt(self, address, position, default_block="latest"):
        """
        Returns the value from a storage position at a given address.
        """
        return self._call('eth_getStorageAt', [address, hex(position), default_block])

    def eth_getTransactionCount(self, address, default_block="latest"):
        """
        Returns the number of transactions send from a address.
        """
        return self._call('eth_getTransactionCount', [address, default_block])

    def eth_getBlockTransactionCountByHash(self, block_hash):
        """
        Returns the number of transactions in a block from a block matching the given block hash.
        """
        return self._call('eth_getTransactionCount', [block_hash])

    def eth_getBlockTransactionCountByNumber(self, block_number):
        """
        Returns the number of transactions in a block from a block matching the given block number.
        """
        return self._call('eth_getBlockTransactionCountByNumber', [hex(block_number)])

    def eth_getUncleCountByblockHash(self, block_hash):
        """
        Returns the number of uncles in a block from a block matching the given block hash.
        """
        return self._call('eth_getUncleCountByblockHash', [block_hash])

    def eth_getUncleCountByblockNumber(self, block_number):
        """
        Returns the number of uncles in a block from a block matching the given block number.
        """
        return self._call('eth_getUncleCountByblockNumber', [hex(block_number)])

    def eth_getCode(self, address, default_block="latest"):
        """
        Returns code at a given address.
        """
        return self._call('eth_getCode', [address, default_block])

    def eth_getBlockByHash(self, block_hash, transaction_objects=True):
        """
        Returns information about a block by hash.
        """
        return self._call('eth_getBlockByHash', [block_hash, transaction_objects])

    def eth_flush(self):
        """
        """
        return self._call('eth_flush')

    def eth_getBlockByNumber(self, block_number, transaction_objects=True):
        """
        Returns information about a block by hash.
        """
        return self._call('eth_getBlockByNumber', [block_number, transaction_objects])

    def eth_getTransactionByHash(self, transaction_hash):
        """
        Returns the information about a transaction requested by transaction hash.
        """
        return self._call('eth_getTransactionByHash', [transaction_hash])

    def eth_getTransactionByblockHashAndIndex(self, block_hash, index):
        """
        Returns information about a transaction by block hash and transaction index position.
        """
        return self._call('eth_getTransactionByblock_hashAndIndex', [block_hash, hex(index)])

    def eth_getTransactionByblockNumberAndIndex(self, block_number, index):
        """
        Returns information about a transaction by block number and transaction index position.
        """
        return self._call('eth_getTransactionByblock_numberAndIndex', [block_number, hex(index)])

    def eth_getUncleByblockHashAndIndex(self, block_hash, index, transaction_objects=True):
        """
        Returns information about a uncle of a block by hash and uncle index position.
        """
        return self._call('eth_getUncleByblock_hashAndIndex', [block_hash, hex(index), transaction_objects])

    def eth_getUncleByblockNumberAndIndex(self, block_number, index, transaction_objects=True):
        """
        Returns information about a uncle of a block by number and uncle index position.
        """
        return self._call('eth_getUncleByblock_numberAndIndex', [block_number, hex(index), transaction_objects])

    def eth_getCompilers(self):
        """
        Returns a list of available compilers in the client.
        """
        return self._call('eth_getCompilers')

    def eth_compileSolidity(self, code):
        """
        Returns compiled solidity code.
        """
        return self._call('eth_compileSolidity', [code])

    def eth_compileLLL(self, code):
        """
        Returns compiled LLL code.
        """
        return self._call('eth_compileLLL', [code])

    def eth_compileSerpent(self, code):
        """
        Returns compiled serpent code.
        """
        return self._call('eth_compileSerpent', [code])

    def eth_newFilter(self, from_block="latest", to_block="latest", address=None, topics=None):
        """
        Creates a filter object, based on filter options, to notify when the state changes (logs).
        To check if the state has changed, call eth_getFilterChanges.
        """
        _filter = {
            'fromBlock': from_block,
            'toBlock': to_block,
            'address': address,
            'topics': topics
        }
        return self._call('eth_newFilter', [_filter])

    def eth_newBlockFilter(self, default_block="latest"):
        """
        Creates a filter object, based on an option string, to notify when state changes (logs). To check if the state has changed, call eth_getFilterChanges.
        """
        return self._call('eth_newBlockFilter', [default_block])

    def eth_uninstallFilter(self, filter_id):
        """
        Uninstalls a filter with given id. Should always be called when watch is no longer needed. Additonally Filters timeout when they aren't requested with eth_getFilterChanges for a period of time.
        """
        return self._call('eth_uninstallFilter', [hex(filter_id)])

    def eth_getFilterChanges(self, filter_id):
        """
        Polling method for a filter, which returns an array of logs which occurred since last poll.
        """
        return self._call('eth_getFilterChanges', [hex(filter_id)])

    def eth_getFilterLogs(self, filter_id):
        """
        Returns an array of all logs matching filter with given id.
        """
        return self._call('eth_getFilterLogs', [hex(filter_id)])

    def eth_getLogs(self, filter_object):
        """
        Returns an array of all logs matching a given filter object.
        """
        return self._call('eth_getLogs', [filter_object])

    def eth_getWork(self):
        """
        Returns the hash of the current block, the seedHash, and the difficulty to be met.
        """
        return self._call('eth_getWork')

    def eth_submitWork(self, nonce, header, mix_digest):
        """
        Used for submitting a proof-of-work solution.
        """
        return self._call('eth_submitWork', [nonce, header, mix_digest])

    def db_putString(self, database_name, key_name, string):
        """
        Stores a string in the local database.
        """
        return self._call('db_putString', [database_name, key_name, string])

    def db_getString(self, database_name, key_name):
        """
        Stores a string in the local database.
        """
        return self._call('db_getString', [database_name, key_name])

    def db_putHex(self, database_name, key_name, string):
        """
        Stores binary data in the local database.
        """
        return self._call('db_putHex', [database_name, key_name, string.encode('hex')])

    def db_getHex(self, database_name, key_name):
        """
        Returns binary data from the local database.
        """
        return self._call('db_getString', [database_name, key_name]).decode('hex')

    def shh_version(self):
        """
        Returns the current whisper protocol version.
        """
        return self._call('shh_version')

    def shh_post(self, topics, payload, priority, ttl, _from=None, to=None):
        """
        Sends a whisper message.
        """
        whisper_object = {
            'from': _from,
            'to': to,
            'topics': topics,
            'payload': payload,
            'priority': priority,
            'ttl': ttl
        }
        return self._call('shh_post', [whisper_object])

    def shh_newIdentinty(self):
        """
        Creates new whisper identity in the client.
        """
        return self._call('shh_newIdentinty')

    def shh_hasIdentity(self, address):
        """
        Checks if the client hold the private keys for a given identity.
        """
        return self._call('shh_hasIdentity', [address])

    def shh_newGroup(self):
        """
        """
        return self._call('shh_hasIdentity')

    def shh_addToGroup(self):
        """
        """
        return self._call('shh_addToGroup')

    def shh_newFilter(self, to, topics):
        """
        Creates filter to notify, when client receives whisper message matching the filter options.
        """
        _filter = {
            'to': to,
            'topics': topics
        }
        return self._call('shh_newFilter', [_filter])

    def shh_uninstallFilter(self, filter_id):
        """
        Uninstalls a filter with given id. Should always be called when watch is no longer needed.
        Additonally Filters timeout when they aren't requested with shh_getFilterChanges for a period of time.
        """
        return self._call('shh_uninstallFilter', [hex(filter_id)])

    def shh_getFilterChanges(self, filter_id):
        """
        Polling method for whisper filters.
        """
        return self._call('shh_getFilterChanges', [hex(filter_id)])

    def shh_getMessages(self, filter_id):
        """
        Get all messages matching a filter, which are still existing in the node.
        """
        return self._call('shh_getMessages', [hex(filter_id)])
<|endoftext|>


<commit_before>import traceback
import bcrypt
import uuid
import threading

from appdaemon.appdaemon import AppDaemon
import appdaemon.utils as utils
from appdaemon.stream.socketio_handler import SocketIOHandler
from appdaemon.stream.ws_handler import WSHandler
from appdaemon.stream.sockjs_handler import SockJSHandler

class ADStream:
    def __init__(self, ad: AppDaemon, app, transport):

        self.AD = ad
        self.logger = ad.logging.get_child("_stream")
        self.access = ad.logging.get_access()
        self.app = app
        self.transport = transport
        self.handlers = {}
        self.handlers_lock = threading.RLock()

        if self.transport == "ws":
            self.stream_handler = WSHandler(self, app, "/stream", self.AD)
        elif self.transport == "socketio":
            self.stream_handler = SocketIOHandler(self, app, "/stream", self.AD)
        elif self.transport == "sockjs":
            self.stream_handler = SockJSHandler(self, app, "/stream", self.AD)
        else:
            self.logger.warning("Unknown stream type: %s", transport)

    def get_handler(self, id):
        with self.handlers_lock:
            for handle in self.handlers:
                if self.handlers[handle].stream.client_id == id:
                    return self.handlers[handle]
        return None

    def get_handle(self, id):
        with self.handlers_lock:
            for handle in self.handlers:
                if self.handlers[handle].stream.client_id == id:
                    return handle
        return None

    async def on_connect(self, request):
        # New connection - create a handler and add it to the list
        handle = uuid.uuid4().hex
        rh = RequestHandler(self.AD, self, handle, request)
        with self.handlers_lock:
            self.handlers[handle] = rh
        await rh.stream.run()

    async def on_disconnect(self, handle):
        with self.handlers_lock:
            del self.handlers[handle]

    async def process_event(self, data):  # noqa: C901
        try:
            with self.handlers_lock:
                if len(self.handlers) > 0:
                    # self.logger.debug("Sending data: %s", data)
                    for handler in self.handlers:
                        await self.handlers[handler]._event(data)

        except Exception:
            self.logger.warning("-" * 60)
            self.logger.warning("Unexpected error during 'process_event()'")
            self.logger.warning("-" * 60)
            self.logger.warning(traceback.format_exc())
            self.logger.warning("-" * 60)

## Any method here that doesn't begin with "_" will be exposed to the stream
## directly. Only Create public methods here if you wish to make them
## stream commands.
class RequestHandler:
    def __init__(self, ad: AppDaemon, adstream, handle, request):
        self.AD = ad
        self.handle = handle
        self.adstream = adstream
        self.authed = False
        self.client_name = None
        self.subscriptions = {
            "state": {},
            "event": {},
        }

        self.logger = ad.logging.get_child("_stream")
        self.access = ad.logging.get_access()

        if self.AD.http.password is None:
            self.authed = True

        # Create a stream
        #
        self.stream = self.adstream.stream_handler.makeStream(
            self.AD, request, on_message=self._on_message, on_disconnect=self._on_disconnect
        )
        #

    async def _on_message(self, data):
        await self._request(data)

    async def _on_disconnect(self):
        await self.adstream.on_disconnect(self.handle)
        self.access.info("Client disconnection from %s", self.client_name)
        event_data = {
            "event_type": "stream_disconnected",
            "data": {"client_name": self.client_name},
        }

        await self.AD.events.process_event("admin", event_data)

    async def _event(self, data):
        response = {"data": data}
        if data["event_type"] == "state_changed":
            for handle, sub in self.subscriptions["state"].items():
                if sub["namespace"].endswith("*"):
                    if not data["namespace"].startswith(sub["namespace"][:-1]):
                        continue
                else:
                    if not data["namespace"] == sub["namespace"]:
                        continue

                if sub["entity_id"].endswith("*"):
                    if not data["data"]["entity_id"].startswith(sub["entity_id"][:-1]):
                        continue
                else:
                    if not data["data"]["entity_id"] == sub["entity_id"]:
                        continue

                response["response_id"] = sub["response_id"]
                response["response_type"] = "state_changed"
                await self._respond(response)
                break
        else:
            for handle, sub in self.subscriptions["event"].items():
                if sub["namespace"].endswith("*"):
                    if not data["namespace"].startswith(sub["namespace"][:-1]):
                        continue
                else:
                    if not data["namespace"] == sub["namespace"]:
                        continue

                if sub["event"].endswith("*"):
                    if not data["event_type"].startswith(sub["event"][:-1]):
                        continue
                else:
                    if not data["event_type"] == sub["event"]:
                        continue

                response["response_id"] = sub["response_id"]
                response["response_type"] = "event"
                await self._respond(response)
                break

    async def _respond(self, data):
        self.logger.debug("--> %s", data)
        await self.stream.sendclient(data)

    async def _response_success(self, msg, data=None):
        response = {"response_type": msg["request_type"]}
        if "request_id" in msg:
            response["response_id"] = msg["request_id"]
        response["response_success"] = True
        if data is None:
            response["data"] = {}
        else:
            response["data"] = data

        response["request"] = msg

        await self._respond(response)

    async def _response_error(self, msg, error):
        response = {"response_type": msg["request_type"]}
        if "request_id" in msg:
            response["response_id"] = msg["request_id"]
        response["response_success"] = False
        response["response_error"] = error
        response["request"] = msg

        await self._respond(response)

    async def _request(self, msg):
        self.logger.debug("<-- %s", msg)

        if "request_type" not in msg:
            return await self._response_error(msg, "invalid request")

        if msg["request_type"][0] == "_":
            return await self._response_error(msg, "forbidden request")

        if not hasattr(self, msg["request_type"]):
            return await self._response_error(msg, "unavailable request")

        fn = getattr(self, msg["request_type"])

        if not callable(fn):
            return await self._response_error(msg, "uncallable request")

        request_data = msg.get("data", {})
        request_id = msg.get("request_id", None)

        try:
            data = await fn(request_data, request_id)
            if data is not None or request_id is not None:
                return await self._response_success(msg, data)
        except RequestHandlerException as e:
            return await self._response_error(msg, str(e))
        except Exception as e:
            await self._response_error(msg, "Unknown error occurred, check AppDaemon logs: {}".format(str(e)))
            raise

    async def _check_adcookie(self, cookie):
        return await utils.run_in_executor(self, bcrypt.checkpw, str.encode(self.AD.http.password), str.encode(cookie))

    async def _auth_data(self, data):
        if "password" in data:
            if data["password"] == self.AD.http.password:
                self.authed = True
                return

        if "cookie" in data:
            if await self._check_adcookie(data["cookie"]):
                self.authed = True
                return

    async def hello(self, data, request_id):
        if "client_name" not in data:
            raise RequestHandlerException("client_name required")
        else:
            self.client_name = data["client_name"]

        if self.AD.http.password is None:
            self.authed = True

        if not self.authed:
            await self._auth_data(data)

        if not self.authed:
            raise RequestHandlerException("authorization failed")

        self.access.info("New client %s connected", data["client_name"])
        response_data = {"version": utils.__version__}

        event_data = {
            "event_type": "stream_connected",
            "data": {"client_name": self.client_name},
        }

        await self.AD.events.process_event("admin", event_data)

        return response_data

    async def get_services(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        return self.AD.services.list_services()

    async def fire_event(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "namespace" not in data:
            raise RequestHandlerException("invalid namespace")

        if "event" not in data:
            raise RequestHandlerException("invalid event")

        event_data = data.get("data", {})

        return await self.AD.events.fire_event(data["namespace"], data["event"], **event_data)

    async def call_service(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "namespace" not in data:
            raise RequestHandlerException("invalid namespace")

        if "service" not in data:
            raise RequestHandlerException("invalid service")
        else:
            service = data["service"]

        if "domain" not in data:
            d, s = service.split("/")
            if d and s:
                domain = d
                service = s
            else:
                raise RequestHandlerException("invalid domain")
        else:
            domain = data["domain"]

        if "data" not in data:
            service_data = {}
        else:
            del data["data"]["service"]
            service_data = data["data"]

        return await self.AD.services.call_service(data["namespace"], domain, service, service_data)

    async def get_state(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        namespace = data.get("namespace", None)
        entity_id = data.get("entity_id", None)

        if entity_id is not None and namespace is None:
            raise RequestHandlerException("entity_id cannot be set without namespace")

        return self.AD.state.get_entity(namespace, entity_id, self.client_name)

    async def listen_state(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "namespace" not in data:
            raise RequestHandlerException("invalid namespace")

        if "entity_id" not in data:
            raise RequestHandlerException("invalid entity_id")

        handle = data.get("handle", uuid.uuid4().hex)

        if handle in self.subscriptions["state"]:
            raise RequestHandlerException("handle already exists")

        self.subscriptions["state"][handle] = {
            "response_id": request_id,
            "namespace": data["namespace"],
            "entity_id": data["entity_id"],
        }

        return handle

    async def cancel_listen_state(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "handle" not in data:
            raise RequestHandlerException("invalid handle")

        if data["handle"] not in self.subscriptions["state"]:
            raise RequestHandlerException("invalid handle")

        del self.subscriptions["state"][data["handle"]]

        return True

    async def listen_event(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "namespace" not in data:
            raise RequestHandlerException("invalid namespace")

        if "event" not in data:
            raise RequestHandlerException("invalid event")

        handle = data.get("handle", uuid.uuid4().hex)

        if handle in self.subscriptions["event"]:
            raise RequestHandlerException("handle already exists")

        self.subscriptions["event"][handle] = {
            "response_id": request_id,
            "namespace": data["namespace"],
            "event": data["event"],
        }

        return handle

    async def cancel_listen_event(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "handle" not in data:
            raise RequestHandlerException("invalid handle")

        if data["handle"] not in self.subscriptions["event"]:
            raise RequestHandlerException("invalid handle")

        del self.subscriptions["event"][data["handle"]]

        return True

class RequestHandlerException(Exception):
    pass
<commit_msg>Update adstream.py<commit_after>import traceback
import bcrypt
import uuid
import threading

from appdaemon.appdaemon import AppDaemon
import appdaemon.utils as utils
from appdaemon.stream.socketio_handler import SocketIOHandler
from appdaemon.stream.ws_handler import WSHandler
from appdaemon.stream.sockjs_handler import SockJSHandler

class ADStream:
    def __init__(self, ad: AppDaemon, app, transport):

        self.AD = ad
        self.logger = ad.logging.get_child("_stream")
        self.access = ad.logging.get_access()
        self.app = app
        self.transport = transport
        self.handlers = {}
        self.handlers_lock = threading.RLock()

        if self.transport == "ws":
            self.stream_handler = WSHandler(self, app, "/stream", self.AD)
        elif self.transport == "socketio":
            self.stream_handler = SocketIOHandler(self, app, "/stream", self.AD)
        elif self.transport == "sockjs":
            self.stream_handler = SockJSHandler(self, app, "/stream", self.AD)
        else:
            self.logger.warning("Unknown stream type: %s", transport)

    def get_handler(self, id):
        with self.handlers_lock:
            for handle in self.handlers:
                if self.handlers[handle].stream.client_id == id:
                    return self.handlers[handle]
        return None

    def get_handle(self, id):
        with self.handlers_lock:
            for handle in self.handlers:
                if self.handlers[handle].stream.client_id == id:
                    return handle
        return None

    async def on_connect(self, request):
        # New connection - create a handler and add it to the list
        handle = uuid.uuid4().hex
        rh = RequestHandler(self.AD, self, handle, request)
        with self.handlers_lock:
            self.handlers[handle] = rh
        await rh.stream.run()

    async def on_disconnect(self, handle):
        with self.handlers_lock:
            del self.handlers[handle]

    async def process_event(self, data):  # noqa: C901
        try:
            with self.handlers_lock:
                if len(self.handlers) > 0:
                    # self.logger.debug("Sending data: %s", data)
                    for handler in self.handlers:
                        await self.handlers[handler]._event(data)

        except Exception:
            self.logger.warning("-" * 60)
            self.logger.warning("Unexpected error during 'process_event()'")
            self.logger.warning("-" * 60)
            self.logger.warning(traceback.format_exc())
            self.logger.warning("-" * 60)

## Any method here that doesn't begin with "_" will be exposed to the stream
## directly. Only Create public methods here if you wish to make them
## stream commands.
class RequestHandler:
    def __init__(self, ad: AppDaemon, adstream, handle, request):
        self.AD = ad
        self.handle = handle
        self.adstream = adstream
        self.authed = False
        self.client_name = None
        self.subscriptions = {
            "state": {},
            "event": {},
        }

        self.logger = ad.logging.get_child("_stream")
        self.access = ad.logging.get_access()

        if self.AD.http.password is None:
            self.authed = True

        # Create a stream
        #
        self.stream = self.adstream.stream_handler.makeStream(
            self.AD, request, on_message=self._on_message, on_disconnect=self._on_disconnect
        )
        #

    async def _on_message(self, data):
        await self._request(data)

    async def _on_disconnect(self):
        await self.adstream.on_disconnect(self.handle)
        self.access.info("Client disconnection from %s", self.client_name)
        event_data = {
            "event_type": "stream_disconnected",
            "data": {"client_name": self.client_name},
        }

        await self.AD.events.process_event("admin", event_data)

    async def _event(self, data):
        response = {"data": data}
        if data["event_type"] == "state_changed":
            for handle, sub in self.subscriptions["state"].items():
                if sub["namespace"].endswith("*"):
                    if not data["namespace"].startswith(sub["namespace"][:-1]):
                        continue
                else:
                    if not data["namespace"] == sub["namespace"]:
                        continue

                if sub["entity_id"].endswith("*"):
                    if not data["data"]["entity_id"].startswith(sub["entity_id"][:-1]):
                        continue
                else:
                    if not data["data"]["entity_id"] == sub["entity_id"]:
                        continue

                response["response_id"] = sub["response_id"]
                response["response_type"] = "state_changed"
                await self._respond(response)
                break
        else:
            for handle, sub in self.subscriptions["event"].items():
                if sub["namespace"].endswith("*"):
                    if not data["namespace"].startswith(sub["namespace"][:-1]):
                        continue
                else:
                    if not data["namespace"] == sub["namespace"]:
                        continue

                if sub["event"].endswith("*"):
                    if not data["event_type"].startswith(sub["event"][:-1]):
                        continue
                else:
                    if not data["event_type"] == sub["event"]:
                        continue

                response["response_id"] = sub["response_id"]
                response["response_type"] = "event"
                await self._respond(response)
                break

    async def _respond(self, data):
        self.logger.debug("--> %s", data)
        await self.stream.sendclient(data)

    async def _response_success(self, msg, data=None):
        response = {"response_type": msg["request_type"]}
        if "request_id" in msg:
            response["response_id"] = msg["request_id"]
        response["response_success"] = True
        if data is None:
            response["data"] = {}
        else:
            response["data"] = data

        response["request"] = msg

        await self._respond(response)

    async def _response_error(self, msg, error):
        response = {"response_type": msg["request_type"]}
        if "request_id" in msg:
            response["response_id"] = msg["request_id"]
        response["response_success"] = False
        response["response_error"] = error
        response["request"] = msg

        await self._respond(response)

    async def _request(self, msg):
        self.logger.debug("<-- %s", msg)

        if "request_type" not in msg:
            return await self._response_error(msg, "invalid request")

        if msg["request_type"][0] == "_":
            return await self._response_error(msg, "forbidden request")

        if not hasattr(self, msg["request_type"]):
            return await self._response_error(msg, "unavailable request")

        fn = getattr(self, msg["request_type"])

        if not callable(fn):
            return await self._response_error(msg, "uncallable request")

        request_data = msg.get("data", {})
        request_id = msg.get("request_id", None)

        try:
            data = await fn(request_data, request_id)
            if data is not None or request_id is not None:
                return await self._response_success(msg, data)
        except RequestHandlerException as e:
            return await self._response_error(msg, str(e))
        except Exception as e:
            await self._response_error(msg, "Unknown error occurred, check AppDaemon logs: {}".format(str(e)))
            raise

    async def _check_adcookie(self, cookie):
        return await utils.run_in_executor(self, bcrypt.checkpw, str.encode(self.AD.http.password), str.encode(cookie))

    async def _auth_data(self, data):
        if "password" in data:
            if data["password"] == self.AD.http.password:
                self.authed = True
                return

        if "cookie" in data:
            if await self._check_adcookie(data["cookie"]):
                self.authed = True
                return

    async def hello(self, data, request_id):
        if "client_name" not in data:
            raise RequestHandlerException("client_name required")
        else:
            self.client_name = data["client_name"]

        if self.AD.http.password is None:
            self.authed = True

        if not self.authed:
            await self._auth_data(data)

        if not self.authed:
            raise RequestHandlerException("authorization failed")

        self.access.info("New client %s connected", data["client_name"])
        response_data = {"version": utils.__version__}

        event_data = {
            "event_type": "stream_connected",
            "data": {"client_name": self.client_name},
        }

        await self.AD.events.process_event("admin", event_data)

        return response_data

    async def get_services(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        return self.AD.services.list_services()

    async def fire_event(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "namespace" not in data:
            raise RequestHandlerException("invalid namespace")

        if "event" not in data:
            raise RequestHandlerException("invalid event")

        event_data = data.get("data", {})

        return await self.AD.events.fire_event(data["namespace"], data["event"], **event_data)

    async def call_service(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "namespace" not in data:
            raise RequestHandlerException("invalid namespace")

        if "service" not in data:
            raise RequestHandlerException("invalid service")
        else:
            service = data["service"]

        if "domain" not in data:
            d, s = service.split("/")
            if d and s:
                domain = d
                service = s
            else:
                raise RequestHandlerException("invalid domain")
        else:
            domain = data["domain"]

        if "data" not in data:
            service_data = {}
        else:
            service_data = data["data"]

        return await self.AD.services.call_service(data["namespace"], domain, service, service_data)

    async def get_state(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        namespace = data.get("namespace", None)
        entity_id = data.get("entity_id", None)

        if entity_id is not None and namespace is None:
            raise RequestHandlerException("entity_id cannot be set without namespace")

        return self.AD.state.get_entity(namespace, entity_id, self.client_name)

    async def listen_state(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "namespace" not in data:
            raise RequestHandlerException("invalid namespace")

        if "entity_id" not in data:
            raise RequestHandlerException("invalid entity_id")

        handle = data.get("handle", uuid.uuid4().hex)

        if handle in self.subscriptions["state"]:
            raise RequestHandlerException("handle already exists")

        self.subscriptions["state"][handle] = {
            "response_id": request_id,
            "namespace": data["namespace"],
            "entity_id": data["entity_id"],
        }

        return handle

    async def cancel_listen_state(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "handle" not in data:
            raise RequestHandlerException("invalid handle")

        if data["handle"] not in self.subscriptions["state"]:
            raise RequestHandlerException("invalid handle")

        del self.subscriptions["state"][data["handle"]]

        return True

    async def listen_event(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "namespace" not in data:
            raise RequestHandlerException("invalid namespace")

        if "event" not in data:
            raise RequestHandlerException("invalid event")

        handle = data.get("handle", uuid.uuid4().hex)

        if handle in self.subscriptions["event"]:
            raise RequestHandlerException("handle already exists")

        self.subscriptions["event"][handle] = {
            "response_id": request_id,
            "namespace": data["namespace"],
            "event": data["event"],
        }

        return handle

    async def cancel_listen_event(self, data, request_id):
        if not self.authed:
            raise RequestHandlerException("unauthorized")

        if "handle" not in data:
            raise RequestHandlerException("invalid handle")

        if data["handle"] not in self.subscriptions["event"]:
            raise RequestHandlerException("invalid handle")

        del self.subscriptions["event"][data["handle"]]

        return True

class RequestHandlerException(Exception):
    pass
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# h2olog - A BPF-backed request logging client for the H2O server.
#
# USAGE: $ sudo h2olog -p $(pgrep -o h2o)
#
# Copyright 2019 Fastly, Toru Maesaka

from bcc import BPF, USDT
from collections import OrderedDict
import getopt, json, sys

bpf = """
#define MAX_STR_LEN 128
#define MIN(x, y) (((x) < (y)) ? (x) : (y))

/*
 * "trace_line_t" is a general structure to store the data emitted by
 * USDT probes. This structure is pushed into the BPF ring buffer.
 */
struct trace_line_t {
    u64 conn_id;
    u64 req_id;
    u32 http_version;
    u32 http_status;
    u64 header_name_len;
    u64 header_value_len;
    char header_name[MAX_STR_LEN];
    char header_value[MAX_STR_LEN];
};

BPF_PERF_OUTPUT(rxbuf);
BPF_PERF_OUTPUT(txbuf);

int trace_receive_req(struct pt_regs *ctx) {
    struct trace_line_t line = {};

    bpf_usdt_readarg(1, ctx, &line.conn_id);
    bpf_usdt_readarg(2, ctx, &line.req_id);
    bpf_usdt_readarg(3, ctx, &line.http_version);

    if (rxbuf.perf_submit(ctx, &line, sizeof(line)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_receive_req_header(struct pt_regs *ctx) {
    struct trace_line_t line = {};
    void *pos = NULL;

    bpf_usdt_readarg(1, ctx, &line.conn_id);
    bpf_usdt_readarg(2, ctx, &line.req_id);

    // Extract the Header Name
    bpf_usdt_readarg(3, ctx, &pos);
    bpf_usdt_readarg(4, ctx, &line.header_name_len);
    line.header_name_len = MIN(MAX_STR_LEN, line.header_name_len);
    bpf_probe_read(&line.header_name, line.header_name_len, pos);

    // Extract the Header Value
    bpf_usdt_readarg(5, ctx, &pos);
    bpf_usdt_readarg(6, ctx, &line.header_value_len);
    line.header_value_len = MIN(MAX_STR_LEN, line.header_value_len);
    bpf_probe_read(&line.header_value, line.header_value_len, pos);

    if (rxbuf.perf_submit(ctx, &line, sizeof(line)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_send_resp(struct pt_regs *ctx) {
    struct trace_line_t line = {};

    bpf_usdt_readarg(1, ctx, &line.conn_id);
    bpf_usdt_readarg(2, ctx, &line.req_id);
    bpf_usdt_readarg(3, ctx, &line.http_status);

    if (txbuf.perf_submit(ctx, &line, sizeof(line)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}
"""

quic_bpf = """
#define MAX_STR_LEN 32

struct st_quicly_conn_t {
    u32 dummy[4];
    u32 master_id;
};

struct quic_event_t {
    char type[MAX_STR_LEN];
    char dcid[MAX_STR_LEN];
    u64 at;
    u32 master_conn_id;
    u64 packet_num;
    u64 packet_len;
    u32 ack_only;
    u8 first_octet;
    u32 newly_acked;
};

BPF_PERF_OUTPUT(events);

int trace_accept(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "accept");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &pos);
    bpf_probe_read(&event.dcid, MAX_STR_LEN, pos);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_receive(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "receive");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &pos);
    bpf_probe_read(&event.dcid, MAX_STR_LEN, pos);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_idle_timeout(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "idle_timeout");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_stateless_reset_receive(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "stateless_reset_receive");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_packet_prepare(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "packet_prepare");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.first_octet);
    bpf_usdt_readarg(4, ctx, &pos);
    bpf_probe_read(&event.dcid, MAX_STR_LEN, pos);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_packet_commit(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "packet_commit");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.packet_num);
    bpf_usdt_readarg(4, ctx, &event.packet_len);
    bpf_usdt_readarg(5, ctx, &event.ack_only);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_packet_acked(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "packet_acked");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.packet_num);
    bpf_usdt_readarg(4, ctx, &event.newly_acked);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_packet_lost(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "packet_lost");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.packet_num);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}
"""

def handle_req_line(cpu, data, size):
    line = b["rxbuf"].event(data)
    if line.http_version:
        v = "HTTP/%d.%d" % (line.http_version / 256, line.http_version % 256)
        print("%u %u RxProtocol %s" % (line.conn_id, line.req_id, v))
    else:
        print("%u %u RxHeader   %s %s" % (line.conn_id, line.req_id, line.header_name, line.header_value))

def handle_resp_line(cpu, data, size):
    line = b["txbuf"].event(data)
    print("%u %u TxStatus   %d" % (line.conn_id, line.req_id, line.http_status))

def load_common_fields(hsh, line):
    for k in ['at', 'type', 'master_conn_id']:
        hsh[k] = getattr(line, k)

def handle_quic_line(cpu, data, size):
    line = b["events"].event(data)
    rv = OrderedDict()
    load_common_fields(rv, line)

    if line.type == "accept":
        rv["dcid"] = getattr(line, "dcid")
    elif line.type == "receive":
        rv["dcid"] = getattr(line, "dcid")
    elif line.type == "packet_prepare":
        for k in ["first_octet", "dcid"]:
            rv[k] = getattr(line, k)
    elif line.type == "packet_commit":
        for k in ["packet_num", "packet_len", "ack_only"]:
            rv[k] = getattr(line, k)
    elif line.type == "packet_acked":
        for k in ["packet_num", "newly_acked"]:
            rv[k] = getattr(line, k)
    elif line.type == "packet_lost":
        for k in ["packet_num"]:
            rv[k] = getattr(line, k)

    print(json.dumps(rv))

def usage():
    print ("USAGE: h2olog -p PID")
    print ("       h2olog quic -p PID")
    exit()

def trace_http():
    b["rxbuf"].open_perf_buffer(handle_req_line)
    b["txbuf"].open_perf_buffer(handle_resp_line)

    while 1:
        try:
            b.perf_buffer_poll()
        except KeyboardInterrupt:
            exit()

def trace_quic():
    b["events"].open_perf_buffer(handle_quic_line)
    while 1:
        try:
            b.perf_buffer_poll()
        except KeyboardInterrupt:
            exit()

if len(sys.argv) < 1:
    usage()

tracer_func = trace_http
optidx = 1
if sys.argv[1] == "quic":
    tracer_func = trace_quic
    optidx = 2

try:
    h2o_pid = 0
    opts, args = getopt.getopt(sys.argv[optidx:], 'p:')
    for opt, arg in opts:
        if opt == "-p":
            h2o_pid = arg
except getopt.error as msg:
    print(msg)
    sys.exit(2)

if h2o_pid == 0:
    usage()

u = USDT(pid=int(h2o_pid))
if sys.argv[1] == "quic":
    u.enable_probe(probe="accept", fn_name="trace_accept")
    u.enable_probe(probe="receive", fn_name="trace_receive")
    u.enable_probe(probe="idle_timeout", fn_name="trace_idle_timeout")
    u.enable_probe(probe="stateless_reset_receive", fn_name="trace_stateless_reset_receive")
    u.enable_probe(probe="packet_prepare", fn_name="trace_packet_prepare")
    u.enable_probe(probe="packet_commit", fn_name="trace_packet_commit")
    u.enable_probe(probe="packet_acked", fn_name="trace_packet_acked")
    u.enable_probe(probe="packet_lost", fn_name="trace_packet_lost")
    b = BPF(text=quic_bpf, usdt_contexts=[u])
else:
    u.enable_probe(probe="receive_request", fn_name="trace_receive_req")
    u.enable_probe(probe="receive_request_header", fn_name="trace_receive_req_header")
    u.enable_probe(probe="send_response", fn_name="trace_send_resp")
    b = BPF(text=bpf, usdt_contexts=[u])

tracer_func()
<commit_msg>h2olog/quic: support the version_switch probe<commit_after>#!/usr/bin/env python
#
# h2olog - A BPF-backed request logging client for the H2O server.
#
# USAGE: $ sudo h2olog -p $(pgrep -o h2o)
#
# Copyright 2019 Fastly, Toru Maesaka

from bcc import BPF, USDT
from collections import OrderedDict
import getopt, json, sys

bpf = """
#define MAX_STR_LEN 128
#define MIN(x, y) (((x) < (y)) ? (x) : (y))

/*
 * "trace_line_t" is a general structure to store the data emitted by
 * USDT probes. This structure is pushed into the BPF ring buffer.
 */
struct trace_line_t {
    u64 conn_id;
    u64 req_id;
    u32 http_version;
    u32 http_status;
    u64 header_name_len;
    u64 header_value_len;
    char header_name[MAX_STR_LEN];
    char header_value[MAX_STR_LEN];
};

BPF_PERF_OUTPUT(rxbuf);
BPF_PERF_OUTPUT(txbuf);

int trace_receive_req(struct pt_regs *ctx) {
    struct trace_line_t line = {};

    bpf_usdt_readarg(1, ctx, &line.conn_id);
    bpf_usdt_readarg(2, ctx, &line.req_id);
    bpf_usdt_readarg(3, ctx, &line.http_version);

    if (rxbuf.perf_submit(ctx, &line, sizeof(line)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_receive_req_header(struct pt_regs *ctx) {
    struct trace_line_t line = {};
    void *pos = NULL;

    bpf_usdt_readarg(1, ctx, &line.conn_id);
    bpf_usdt_readarg(2, ctx, &line.req_id);

    // Extract the Header Name
    bpf_usdt_readarg(3, ctx, &pos);
    bpf_usdt_readarg(4, ctx, &line.header_name_len);
    line.header_name_len = MIN(MAX_STR_LEN, line.header_name_len);
    bpf_probe_read(&line.header_name, line.header_name_len, pos);

    // Extract the Header Value
    bpf_usdt_readarg(5, ctx, &pos);
    bpf_usdt_readarg(6, ctx, &line.header_value_len);
    line.header_value_len = MIN(MAX_STR_LEN, line.header_value_len);
    bpf_probe_read(&line.header_value, line.header_value_len, pos);

    if (rxbuf.perf_submit(ctx, &line, sizeof(line)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_send_resp(struct pt_regs *ctx) {
    struct trace_line_t line = {};

    bpf_usdt_readarg(1, ctx, &line.conn_id);
    bpf_usdt_readarg(2, ctx, &line.req_id);
    bpf_usdt_readarg(3, ctx, &line.http_status);

    if (txbuf.perf_submit(ctx, &line, sizeof(line)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}
"""

quic_bpf = """
#define MAX_STR_LEN 32

struct st_quicly_conn_t {
    u32 dummy[4];
    u32 master_id;
};

struct quic_event_t {
    char type[MAX_STR_LEN];
    char dcid[MAX_STR_LEN];
    u64 at;
    u32 master_conn_id;
    u64 packet_num;
    u64 packet_len;
    u32 ack_only;
    u8 first_octet;
    u32 newly_acked;
    u32 new_version;
};

BPF_PERF_OUTPUT(events);

int trace_accept(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "accept");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &pos);
    bpf_probe_read(&event.dcid, MAX_STR_LEN, pos);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_receive(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "receive");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &pos);
    bpf_probe_read(&event.dcid, MAX_STR_LEN, pos);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_version_switch(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "version_switch");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.new_version);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_idle_timeout(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "idle_timeout");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_stateless_reset_receive(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "stateless_reset_receive");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_packet_prepare(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "packet_prepare");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.first_octet);
    bpf_usdt_readarg(4, ctx, &pos);
    bpf_probe_read(&event.dcid, MAX_STR_LEN, pos);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_packet_commit(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "packet_commit");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.packet_num);
    bpf_usdt_readarg(4, ctx, &event.packet_len);
    bpf_usdt_readarg(5, ctx, &event.ack_only);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_packet_acked(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "packet_acked");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.packet_num);
    bpf_usdt_readarg(4, ctx, &event.newly_acked);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}

int trace_packet_lost(struct pt_regs *ctx) {
    void *pos = NULL;
    struct quic_event_t event = {};
    struct st_quicly_conn_t conn = {};
    sprintf(event.type, "packet_lost");

    bpf_usdt_readarg(1, ctx, &pos);
    bpf_probe_read(&conn, sizeof(conn), pos);
    event.master_conn_id = conn.master_id;
    bpf_usdt_readarg(2, ctx, &event.at);
    bpf_usdt_readarg(3, ctx, &event.packet_num);

    if (events.perf_submit(ctx, &event, sizeof(event)) < 0)
        bpf_trace_printk("failed to perf_submit\\n");

    return 0;
}
"""

def handle_req_line(cpu, data, size):
    line = b["rxbuf"].event(data)
    if line.http_version:
        v = "HTTP/%d.%d" % (line.http_version / 256, line.http_version % 256)
        print("%u %u RxProtocol %s" % (line.conn_id, line.req_id, v))
    else:
        print("%u %u RxHeader   %s %s" % (line.conn_id, line.req_id, line.header_name, line.header_value))

def handle_resp_line(cpu, data, size):
    line = b["txbuf"].event(data)
    print("%u %u TxStatus   %d" % (line.conn_id, line.req_id, line.http_status))

def load_common_fields(hsh, line):
    for k in ['at', 'type', 'master_conn_id']:
        hsh[k] = getattr(line, k)

def handle_quic_line(cpu, data, size):
    line = b["events"].event(data)
    rv = OrderedDict()
    load_common_fields(rv, line)

    if line.type == "accept":
        rv["dcid"] = getattr(line, "dcid")
    elif line.type == "receive":
        rv["dcid"] = getattr(line, "dcid")
    elif line.type == "version_switch":
        rv["new_version"] = getattr(line, "new_version")
    elif line.type == "packet_prepare":
        for k in ["first_octet", "dcid"]:
            rv[k] = getattr(line, k)
    elif line.type == "packet_commit":
        for k in ["packet_num", "packet_len", "ack_only"]:
            rv[k] = getattr(line, k)
    elif line.type == "packet_acked":
        for k in ["packet_num", "newly_acked"]:
            rv[k] = getattr(line, k)
    elif line.type == "packet_lost":
        for k in ["packet_num"]:
            rv[k] = getattr(line, k)

    print(json.dumps(rv))

def usage():
    print ("USAGE: h2olog -p PID")
    print ("       h2olog quic -p PID")
    exit()

def trace_http():
    b["rxbuf"].open_perf_buffer(handle_req_line)
    b["txbuf"].open_perf_buffer(handle_resp_line)

    while 1:
        try:
            b.perf_buffer_poll()
        except KeyboardInterrupt:
            exit()

def trace_quic():
    b["events"].open_perf_buffer(handle_quic_line)
    while 1:
        try:
            b.perf_buffer_poll()
        except KeyboardInterrupt:
            exit()

if len(sys.argv) < 1:
    usage()

tracer_func = trace_http
optidx = 1
if sys.argv[1] == "quic":
    tracer_func = trace_quic
    optidx = 2

try:
    h2o_pid = 0
    opts, args = getopt.getopt(sys.argv[optidx:], 'p:')
    for opt, arg in opts:
        if opt == "-p":
            h2o_pid = arg
except getopt.error as msg:
    print(msg)
    sys.exit(2)

if h2o_pid == 0:
    usage()

u = USDT(pid=int(h2o_pid))
if sys.argv[1] == "quic":
    u.enable_probe(probe="accept", fn_name="trace_accept")
    u.enable_probe(probe="receive", fn_name="trace_receive")
    u.enable_probe(probe="version_switch", fn_name="trace_version_switch")
    u.enable_probe(probe="idle_timeout", fn_name="trace_idle_timeout")
    u.enable_probe(probe="stateless_reset_receive", fn_name="trace_stateless_reset_receive")
    u.enable_probe(probe="packet_prepare", fn_name="trace_packet_prepare")
    u.enable_probe(probe="packet_commit", fn_name="trace_packet_commit")
    u.enable_probe(probe="packet_acked", fn_name="trace_packet_acked")
    u.enable_probe(probe="packet_lost", fn_name="trace_packet_lost")
    b = BPF(text=quic_bpf, usdt_contexts=[u])
else:
    u.enable_probe(probe="receive_request", fn_name="trace_receive_req")
    u.enable_probe(probe="receive_request_header", fn_name="trace_receive_req_header")
    u.enable_probe(probe="send_response", fn_name="trace_send_resp")
    b = BPF(text=bpf, usdt_contexts=[u])

tracer_func()
<|endoftext|>