1# Copyright 2014 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import errno
6import json
7import logging
8import socket
9import time
10
11from telemetry.core import exceptions
12from telemetry.internal.backends.chrome_inspector import websocket
13
14class WebSocketDisconnected(exceptions.Error):
15  """An attempt was made to use a web socket after it had been disconnected."""
16  pass
17
18
19class InspectorWebsocket(object):
20
21  # See http://www.jsonrpc.org/specification#error_object.
22  METHOD_NOT_FOUND_CODE = -32601
23
24  def __init__(self):
25    """Create a websocket handler for communicating with Inspectors."""
26    self._socket = None
27    self._cur_socket_timeout = 0
28    self._next_request_id = 0
29    self._domain_handlers = {}
30    self._pending_callbacks = dict()
31
32  def RegisterDomain(self, domain_name, notification_handler):
33    """Registers a given domain for handling notification methods.
34
35    For example, given inspector_backend:
36       def OnConsoleNotification(msg):
37          if msg['method'] == 'Console.messageAdded':
38             print msg['params']['message']
39       inspector_backend.RegisterDomain('Console', OnConsoleNotification)
40
41    Args:
42      domain_name: The devtools domain name. E.g., 'Tracing', 'Memory', 'Page'.
43      notification_handler: Handler for devtools notification. Will be
44          called if a devtools notification with matching domain is received
45          via DispatchNotifications. The handler accepts a single paramater:
46          the JSON object representing the notification.
47    """
48    assert domain_name not in self._domain_handlers
49    self._domain_handlers[domain_name] = notification_handler
50
51  def UnregisterDomain(self, domain_name):
52    """Unregisters a previously registered domain."""
53    assert domain_name in self._domain_handlers
54    del self._domain_handlers[domain_name]
55
56  def Connect(self, url, timeout=10):
57    """Connects the websocket.
58
59    Raises:
60      websocket.WebSocketException
61      socket.error
62    """
63    assert not self._socket
64    self._socket = websocket.create_connection(url, timeout=timeout)
65    self._cur_socket_timeout = 0
66    self._next_request_id = 0
67
68  def Disconnect(self):
69    """Disconnects the inspector websocket.
70
71    Raises:
72      websocket.WebSocketException
73      socket.error
74    """
75    if self._socket:
76      self._socket.close()
77      self._socket = None
78
79  def SendAndIgnoreResponse(self, req):
80    """Sends a request without waiting for a response.
81
82    Raises:
83      websocket.WebSocketException: Error from websocket library.
84      socket.error: Error from websocket library.
85      exceptions.WebSocketDisconnected: The socket was disconnected.
86    """
87    self._SendRequest(req)
88
89  def _SendRequest(self, req):
90    if not self._socket:
91      raise WebSocketDisconnected()
92    req['id'] = self._next_request_id
93    self._next_request_id += 1
94    data = json.dumps(req)
95    self._socket.send(data)
96    if logging.getLogger().isEnabledFor(logging.DEBUG):
97      logging.debug('sent [%s]', json.dumps(req, indent=2, sort_keys=True))
98
99  def SyncRequest(self, req, timeout=10):
100    """Sends a request and waits for a response.
101
102    Raises:
103      websocket.WebSocketException: Error from websocket library.
104      socket.error: Error from websocket library.
105      exceptions.WebSocketDisconnected: The socket was disconnected.
106    """
107    self._SendRequest(req)
108
109    while True:
110      res = self._Receive(timeout)
111      if 'id' in res and res['id'] == req['id']:
112        return res
113
114  def AsyncRequest(self, req, callback):
115    """Sends an async request and returns immediately.
116
117    Response will be handled in the |callback| later when DispatchNotifications
118    is invoked.
119
120    Args:
121      callback: a function that takes inspector's response as the argument.
122    """
123    self._SendRequest(req)
124    self._pending_callbacks[req['id']] = callback
125
126  def DispatchNotifications(self, timeout=10):
127    """Waits for responses from the websocket, dispatching them as necessary.
128
129    Raises:
130      websocket.WebSocketException: Error from websocket library.
131      socket.error: Error from websocket library.
132      exceptions.WebSocketDisconnected: The socket was disconnected.
133    """
134    self._Receive(timeout)
135
136  def _SetTimeout(self, timeout):
137    if self._cur_socket_timeout != timeout:
138      self._socket.settimeout(timeout)
139      self._cur_socket_timeout = timeout
140
141  def _Receive(self, timeout=10):
142    if not self._socket:
143      raise WebSocketDisconnected()
144
145    self._SetTimeout(timeout)
146
147    while True:
148      try:
149        data = self._socket.recv()
150      except socket.error, e:
151        if e.errno == errno.EAGAIN:
152          # Resource is temporarily unavailable. Try again.
153          # See https://code.google.com/p/chromium/issues/detail?id=545853#c3
154          # for more details.
155          time.sleep(0.1)
156        else:
157          raise
158      else:
159        break
160
161    result = json.loads(data)
162    if logging.getLogger().isEnabledFor(logging.DEBUG):
163      logging.debug(
164          'got [%s]', json.dumps(result, indent=2, sort_keys=True))
165    if 'method' in result:
166      self._HandleNotification(result)
167    elif 'id' in result:
168      self._HandleAsyncResponse(result)
169    return result
170
171  def _HandleNotification(self, result):
172    mname = result['method']
173    dot_pos = mname.find('.')
174    domain_name = mname[:dot_pos]
175    if not domain_name in self._domain_handlers:
176      logging.warn('Unhandled inspector message: %s', result)
177      return
178
179    self._domain_handlers[domain_name](result)
180
181  def _HandleAsyncResponse(self, result):
182    callback = self._pending_callbacks.pop(result['id'], None)
183    if callback:
184      callback(result)
185