1# Copyright 2014 Google Inc. All rights reserved.
2#
3# Use of this source code is governed by a BSD-style
4# license that can be found in the COPYING file or at
5# https://developers.google.com/open-source/licenses/bsd
6
7
8from mod_pywebsocket import util
9
10
11class XHRBenchmarkHandler(object):
12    def __init__(self, headers, rfile, wfile):
13        self._logger = util.get_class_logger(self)
14
15        self.headers = headers
16        self.rfile = rfile
17        self.wfile = wfile
18
19    def do_send(self):
20        content_length = int(self.headers.getheader('Content-Length'))
21
22        self._logger.debug('Requested to receive %s bytes', content_length)
23
24        RECEIVE_BLOCK_SIZE = 1024 * 1024
25
26        bytes_to_receive = content_length
27        while bytes_to_receive > 0:
28            bytes_to_receive_in_this_loop = bytes_to_receive
29            if bytes_to_receive_in_this_loop > RECEIVE_BLOCK_SIZE:
30                bytes_to_receive_in_this_loop = RECEIVE_BLOCK_SIZE
31            received_data = self.rfile.read(bytes_to_receive_in_this_loop)
32            for c in received_data:
33                if c != 'a':
34                    self._logger.debug('Request body verification failed')
35                    return
36            bytes_to_receive -= len(received_data)
37        if bytes_to_receive < 0:
38            self._logger.debug('Received %d more bytes than expected' %
39                               (-bytes_to_receive))
40            return
41
42        # Return the number of received bytes back to the client.
43        response_body = '%d' % content_length
44        self.wfile.write(
45            'HTTP/1.1 200 OK\r\n'
46            'Content-Type: text/html\r\n'
47            'Content-Length: %d\r\n'
48            '\r\n%s' % (len(response_body), response_body))
49        self.wfile.flush()
50
51    def do_receive(self):
52        content_length = int(self.headers.getheader('Content-Length'))
53        request_body = self.rfile.read(content_length)
54
55        request_array = request_body.split(' ')
56        if len(request_array) < 2:
57            self._logger.debug('Malformed request body: %r', request_body)
58            return
59
60        # Parse the size parameter.
61        bytes_to_send = request_array[0]
62        try:
63            bytes_to_send = int(bytes_to_send)
64        except ValueError, e:
65            self._logger.debug('Malformed size parameter: %r', bytes_to_send)
66            return
67        self._logger.debug('Requested to send %s bytes', bytes_to_send)
68
69        # Parse the transfer encoding parameter.
70        chunked_mode = False
71        mode_parameter = request_array[1]
72        if mode_parameter == 'chunked':
73            self._logger.debug('Requested chunked transfer encoding')
74            chunked_mode = True
75        elif mode_parameter != 'none':
76            self._logger.debug('Invalid mode parameter: %r', mode_parameter)
77            return
78
79        # Write a header
80        response_header = (
81            'HTTP/1.1 200 OK\r\n'
82            'Content-Type: application/octet-stream\r\n')
83        if chunked_mode:
84            response_header += 'Transfer-Encoding: chunked\r\n\r\n'
85        else:
86            response_header += (
87                'Content-Length: %d\r\n\r\n' % bytes_to_send)
88        self.wfile.write(response_header)
89        self.wfile.flush()
90
91        # Write a body
92        SEND_BLOCK_SIZE = 1024 * 1024
93
94        while bytes_to_send > 0:
95            bytes_to_send_in_this_loop = bytes_to_send
96            if bytes_to_send_in_this_loop > SEND_BLOCK_SIZE:
97                bytes_to_send_in_this_loop = SEND_BLOCK_SIZE
98
99            if chunked_mode:
100                self.wfile.write('%x\r\n' % bytes_to_send_in_this_loop)
101            self.wfile.write('a' * bytes_to_send_in_this_loop)
102            if chunked_mode:
103                self.wfile.write('\r\n')
104            self.wfile.flush()
105
106            bytes_to_send -= bytes_to_send_in_this_loop
107
108        if chunked_mode:
109            self.wfile.write('0\r\n\r\n')
110            self.wfile.flush()
111