def HandleRegister(self):
"""Registers a data server in the master."""
if not self.MASTER:
self._EmptyResponse(constants.RESPONSE_NOT_MASTER_SERVER)
return
request = rdf_data_server.DataStoreRegistrationRequest(self.post_data)
port = request.port
addr = self.client_address[0]
token = request.token
if not self.NONCE_STORE.ValidateAuthTokenServer(token):
self._EmptyResponse(constants.RESPONSE_SERVER_NOT_AUTHORIZED)
return
newserver = self.MASTER.RegisterServer(addr, port)
if newserver:
self.data_server = newserver
index = newserver.Index()
body = sutils.SIZE_PACKER.pack(index)
# Need to send back the encrypted client credentials.
body += self.NONCE_STORE.EncryptClientCredentials()
self._Response(constants.RESPONSE_OK, body)
else:
# Could not register the Data Server.
logging.warning(
"Could not register server %s:%d. Maybe not allowed?", addr,
port)
self._EmptyResponse(constants.RESPONSE_SERVER_NOT_ALLOWED)
def _DoRegister(self):
try:
username, password = self.handler_cls.NONCE_STORE.GetServerCredentials(
)
# First get a nonce.
res = self.pool.urlopen("POST",
"/server/handshake",
"",
headers={})
if res.status != constants.RESPONSE_OK:
raise errors.DataServerError(
"Could not register data server at "
"data master.")
nonce = res.data
token = self.handler_cls.NONCE_STORE.GenerateServerAuthToken(nonce)
request = rdf_data_server.DataStoreRegistrationRequest(
token=token, port=self.my_port)
body = request.SerializeToString()
headers = {"Content-Length": len(body)}
res = self.pool.urlopen("POST",
"/server/register",
headers=headers,
body=body)
if res.status == constants.RESPONSE_SERVER_NOT_AUTHORIZED:
raise errors.DataServerError("Wrong server password.")
if res.status == constants.RESPONSE_SERVER_NOT_ALLOWED:
raise errors.DataServerError(
"Server not part of this server group.")
if res.status == constants.RESPONSE_NOT_MASTER_SERVER:
raise errors.DataServerError(
"Server %s:%d is not a master server.", self.master_addr,
self.master_port)
if res.status != constants.RESPONSE_OK:
raise errors.DataServerError(
"Could not register data server at data "
"master.")
logging.info("DataServer fully registered.")
id_str = res.data[:sutils.SIZE_PACKER.size]
self.index = sutils.SIZE_PACKER.unpack(id_str)[0]
creds_str = res.data[sutils.SIZE_PACKER.size:]
# Read client credentials so we know who to allow data store access.
creds = auth.ClientCredentials()
creds.InitializeFromEncryption(creds_str, username, password)
self.handler_cls.NONCE_STORE.SetClientCredentials(creds)
return True
except (urllib3.exceptions.HTTPError, urllib3.exceptions.PoolError):
return False
def testRegister(self):
"""Create master and register other servers."""
m = master.DataMaster(self.ports[0], self.mock_service)
self.assertNotEqual(m, None)
self.assertFalse(m.AllRegistered())
servers = [None]
for (i, port) in enumerate(self.ports):
if i == 0:
# Skip master server.
continue
self.assertFalse(m.AllRegistered())
server = m.RegisterServer(self.host, port)
servers.append(server)
self.assertNotEqual(server, None)
self.assertEqual(server.Address(), self.host)
self.assertEqual(server.Port(), port)
self.assertEqual(server.Index(), i)
self.assertTrue(m.AllRegistered())
# Try to register something that does not exist.
self.assertFalse(m.RegisterServer(self.host, 7004))
# Deregister a server.
m.DeregisterServer(servers[1])
self.assertFalse(m.AllRegistered())
# Register again.
m.RegisterServer(servers[1].Address(), servers[1].Port())
self.assertTrue(m.AllRegistered())
for port in self.ports:
for response_sequence in [[
constants.RESPONSE_OK,
constants.RESPONSE_SERVER_NOT_AUTHORIZED
], [constants.RESPONSE_OK, constants.RESPONSE_SERVER_NOT_ALLOWED],
[
constants.RESPONSE_OK,
constants.RESPONSE_NOT_MASTER_SERVER
]]:
response_mocks = []
for response_status in response_sequence:
response_mocks.append(MockResponse(response_status))
pool_class = GetMockHTTPConnectionPoolClass(response_mocks)
with libutils.Stubber(connectionpool, "HTTPConnectionPool",
pool_class):
m = data_server.StandardDataServer(
port, data_server.DataServerHandler)
m.handler_cls.NONCE_STORE = auth.NonceStore()
self.assertRaises(errors.DataServerError, m._DoRegister)
# Ensure two requests have been made.
self.assertEqual(len(pool_class.requests), 2)
# Ensure the register body is non-empty.
self.assertTrue(pool_class.requests[1]["body"])
# Ensure that the register body is a valid rdfvalue.
rdf_data_server.DataStoreRegistrationRequest(
pool_class.requests[1]["body"])
# Ensure the requests are POST requests.
self.assertEqual(pool_class.requests[0]["method"], "POST")
self.assertEqual(pool_class.requests[1]["method"], "POST")
# Ensure the correct URLs are hit according to the API.
self.assertEqual(pool_class.requests[0]["url"],
"/server/handshake")
self.assertEqual(pool_class.requests[1]["url"],
"/server/register")
