Commit 179ff86e authored by zhifan huang's avatar zhifan huang

update http handler

parent 45b09b5a
...@@ -176,12 +176,16 @@ class RegistryServer(object): ...@@ -176,12 +176,16 @@ class RegistryServer(object):
def recv(self, code): def recv(self, code):
try: try:
prefix, msg = self.sock.recv(1<<16).split('\0', 1) prefix, msg = self.sock.recv(1<<16).split(b'\0', 1)
int(prefix, 2) int(prefix, 2)
except ValueError: except ValueError:
pass pass
else: else:
# let recv return to str
prefix = prefix.decode()
msg = msg.decode()
if msg and ord(msg[0]) == code: if msg and ord(msg[0]) == code:
return prefix, msg[1:] return prefix, msg[1:]
return None, None return None, None
...@@ -253,7 +257,8 @@ class RegistryServer(object): ...@@ -253,7 +257,8 @@ class RegistryServer(object):
# (IOW, do the contrary of newPrefix) # (IOW, do the contrary of newPrefix)
self.timeout = not_after and not_after + GRACE_PERIOD self.timeout = not_after and not_after + GRACE_PERIOD
def handle_request(self, request, method, kw): def handle_request(self, request:BaseHTTPRequestHandler,
method:str, kw: dict[str:str]):
m = getattr(self, method) m = getattr(self, method)
if hasattr(m, '_private'): if hasattr(m, '_private'):
authorized_origin = self.config.authorized_origin authorized_origin = self.config.authorized_origin
...@@ -267,7 +272,7 @@ class RegistryServer(object): ...@@ -267,7 +272,7 @@ class RegistryServer(object):
with self.lock: with self.lock:
session = self.sessions[key] session = self.sessions[key]
for key, protocol in session: for key, protocol in session:
if h == hmac.HMAC(key, request.path, hashlib.sha1).digest(): if h == hmac.HMAC(key, request.path.encode(), hashlib.sha1).digest():
break break
else: else:
raise Exception("Wrong HMAC") raise Exception("Wrong HMAC")
...@@ -284,6 +289,9 @@ class RegistryServer(object): ...@@ -284,6 +289,9 @@ class RegistryServer(object):
request.headers.get("user-agent")) request.headers.get("user-agent"))
try: try:
result = m(**kw) result = m(**kw)
# request.wfile.write need bytes
if isinstance(result, str):
result = result.encode()
except HTTPError as e: except HTTPError as e:
return request.send_error(*e.args) return request.send_error(*e.args)
except: except:
......
...@@ -93,11 +93,11 @@ class TestRegistryServer(unittest.TestCase): ...@@ -93,11 +93,11 @@ class TestRegistryServer(unittest.TestCase):
back_sock = self.server.sock back_sock = self.server.sock
sock = self.server.sock= Mock() sock = self.server.sock= Mock()
sock.recv.side_effect = [ sock.recv.side_effect = [
"0001001001001a_msg", b"0001001001001a_msg",
"0001001001002\0001dqdq", b"0001001001002\0001dqdq",
"0001001001001\000a_msg", b"0001001001001\000a_msg",
"0001001001001\000\4a_msg", b"0001001001001\000\4a_msg",
"0000000000000\0" # ERROR, IndexError: msg is null b"0000000000000\0" # ERROR, IndexError: msg is null
] ]
res1 = self.server.recv(4) res1 = self.server.recv(4)
...@@ -158,9 +158,9 @@ class TestRegistryServer(unittest.TestCase): ...@@ -158,9 +158,9 @@ class TestRegistryServer(unittest.TestCase):
key = b"this_is_a_key" key = b"this_is_a_key"
self.server.sessions[prefix] = [(key, protocol)] self.server.sessions[prefix] = [(key, protocol)]
request = Mock() request = Mock()
request.path = b"/func?a=1&b=2&cn=0000000011111111" request.path = "/func?a=1&b=2&cn=0000000011111111"
request.headers = {registry.HMAC_HEADER: base64.b64encode( request.headers = {registry.HMAC_HEADER: base64.b64encode(
hmac.HMAC(key, request.path, hashlib.sha1).digest())} hmac.HMAC(key, request.path.encode(), hashlib.sha1).digest())}
self.server.handle_request(request, method, params) self.server.handle_request(request, method, params)
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment