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")