ws.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166
  1. #!/usr/bin/python
  2. '''
  3. WebSocket server-side load test program. Sends and receives traffic
  4. that has a random payload (length and content) that is checksummed and
  5. given a sequence number. Any errors are reported and counted.
  6. '''
  7. import sys, os, socket, ssl, time, traceback
  8. import random, time
  9. from base64 import b64encode, b64decode
  10. from select import select
  11. sys.path.insert(0,os.path.dirname(__file__) + "/../utils/")
  12. from websocket import *
  13. buffer_size = 65536
  14. max_packet_size = 10000
  15. recv_cnt = send_cnt = 0
  16. def check(buf):
  17. global recv_cnt
  18. try:
  19. data_list = decode(buf)
  20. except:
  21. print "\n<BOF>" + repr(buf) + "<EOF>"
  22. return "Failed to decode"
  23. err = ""
  24. for data in data_list:
  25. if data.count('$') > 1:
  26. raise Exception("Multiple parts within single packet")
  27. if len(data) == 0:
  28. traffic("_")
  29. continue
  30. if data[0] != "^":
  31. err += "buf did not start with '^'\n"
  32. continue
  33. try:
  34. cnt, length, chksum, nums = data[1:-1].split(':')
  35. cnt = int(cnt)
  36. length = int(length)
  37. chksum = int(chksum)
  38. except:
  39. print "\n<BOF>" + repr(data) + "<EOF>"
  40. err += "Invalid data format\n"
  41. continue
  42. if recv_cnt != cnt:
  43. err += "Expected count %d but got %d\n" % (recv_cnt, cnt)
  44. recv_cnt = cnt + 1
  45. continue
  46. recv_cnt += 1
  47. if len(nums) != length:
  48. err += "Expected length %d but got %d\n" % (length, len(nums))
  49. continue
  50. inv = nums.translate(None, "0123456789")
  51. if inv:
  52. err += "Invalid characters found: %s\n" % inv
  53. continue
  54. real_chksum = 0
  55. for num in nums:
  56. real_chksum += int(num)
  57. if real_chksum != chksum:
  58. err += "Expected checksum %d but real chksum is %d\n" % (chksum, real_chksum)
  59. return err
  60. def generate():
  61. global send_cnt, rand_array
  62. length = random.randint(10, max_packet_size)
  63. numlist = rand_array[max_packet_size-length:]
  64. # Error in length
  65. #numlist.append(5)
  66. chksum = sum(numlist)
  67. # Error in checksum
  68. #numlist[0] = 5
  69. nums = "".join( [str(n) for n in numlist] )
  70. data = "^%d:%d:%d:%s$" % (send_cnt, length, chksum, nums)
  71. send_cnt += 1
  72. return encode(data)
  73. def responder(client, delay=10):
  74. global errors
  75. cqueue = []
  76. cpartial = ""
  77. socks = [client]
  78. last_send = time.time() * 1000
  79. while True:
  80. ins, outs, excepts = select(socks, socks, socks, 1)
  81. if excepts: raise Exception("Socket exception")
  82. if client in ins:
  83. buf = client.recv(buffer_size)
  84. if len(buf) == 0: raise Exception("Client closed")
  85. #print "Client recv: %s (%d)" % (repr(buf[1:-1]), len(buf))
  86. if buf[-1] == '\xff':
  87. if cpartial:
  88. err = check(cpartial + buf)
  89. cpartial = ""
  90. else:
  91. err = check(buf)
  92. if err:
  93. traffic("}")
  94. errors = errors + 1
  95. print err
  96. else:
  97. traffic(">")
  98. else:
  99. traffic(".>")
  100. cpartial = cpartial + buf
  101. now = time.time() * 1000
  102. if client in outs and now > (last_send + delay):
  103. last_send = now
  104. #print "Client send: %s" % repr(cqueue[0])
  105. client.send(generate())
  106. traffic("<")
  107. def test_handler(client):
  108. global errors, delay, send_cnt, recv_cnt
  109. send_cnt = 0
  110. recv_cnt = 0
  111. try:
  112. responder(client, delay)
  113. except:
  114. print "accumulated errors:", errors
  115. errors = 0
  116. raise
  117. if __name__ == '__main__':
  118. errors = 0
  119. try:
  120. if len(sys.argv) < 2: raise
  121. listen_port = int(sys.argv[1])
  122. if len(sys.argv) == 3:
  123. delay = int(sys.argv[2])
  124. else:
  125. delay = 10
  126. except:
  127. print "Usage: <listen_port> [delay_ms]"
  128. sys.exit(1)
  129. print "Prepopulating random array"
  130. rand_array = []
  131. for i in range(0, max_packet_size):
  132. rand_array.append(random.randint(0, 9))
  133. settings['listen_port'] = listen_port
  134. settings['daemon'] = False
  135. settings['handler'] = test_handler
  136. start_server()