def test_remove_after_releasing_some_addresses(self):
""" remove after releasing some addresses shouldn't remove block """
# Assign IP block
self._stub.AddIPBlock(self._block_msg)
# Allocate 2 IPs
alloc_request0 = AllocateIPRequest(
sid=self._sid0,
version=AllocateIPRequest.IPV4,
apn=self._apn0,
)
ip_msg0 = self._stub.AllocateIPAddress(alloc_request0)
alloc_request1 = AllocateIPRequest(
sid=self._sid1,
version=AllocateIPRequest.IPV4,
apn=self._apn0,
)
self._stub.AllocateIPAddress(alloc_request1)
# Test remove without force -- should not remove block
remove_request0 = RemoveIPBlockRequest(
ip_blocks=[self._block_msg],
force=False,
)
resp = self._stub.RemoveIPBlock(remove_request0)
expect = RemoveIPBlockResponse()
self.assertEqual(expect, resp)
# Ensure that block has not been removed
resp = self._stub.ListAddedIPv4Blocks(Void())
expect = ListAddedIPBlocksResponse()
expect.ip_block_list.extend([self._block_msg])
self.assertEqual(expect, resp)
# Release the allocated IPs
release_request0 = ReleaseIPRequest(
sid=self._sid0,
ip=ip_msg0.ip_list[0],
apn=self._apn0,
)
resp = self._stub.ReleaseIPAddress(release_request0)
self.assertEqual(resp, Void())
# Test remove without force -- should not remove block
remove_request1 = RemoveIPBlockRequest(
ip_blocks=[self._block_msg],
force=False,
)
resp = self._stub.RemoveIPBlock(remove_request1)
expect = RemoveIPBlockResponse()
self.assertEqual(expect, resp)
# Ensure that block has not been removed
resp = self._stub.ListAddedIPv4Blocks(Void())
expect = ListAddedIPBlocksResponse()
expect.ip_block_list.extend([self._block_msg])
self.assertEqual(expect, resp)
def test_release_ip_address(self):
""" test ReleaseIPAddress """
self._stub.AddIPBlock(self._block_msg)
alloc_request0 = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4,
apn=self._apn0)
ip_msg0 = self._stub.AllocateIPAddress(alloc_request0)
alloc_request1 = AllocateIPRequest(sid=self._sid1,
version=AllocateIPRequest.IPV4,
apn=self._apn0)
ip_msg1 = self._stub.AllocateIPAddress(alloc_request1)
# release ip_msg0
release_request0 = ReleaseIPRequest(sid=self._sid0,
ip=ip_msg0.ip_list[0],
apn=self._apn0)
resp = self._stub.ReleaseIPAddress(release_request0)
self.assertEqual(resp, Void())
resp = self._stub.ListAllocatedIPs(self._block_msg)
tmp = ListAllocatedIPsResponse()
tmp.ip_list.extend([ip_msg1.ip_list[0]])
self.assertEqual(resp, tmp)
# release ip_msg1
release_request1 = ReleaseIPRequest(sid=self._sid1,
ip=ip_msg1.ip_list[0],
apn=self._apn0)
resp = self._stub.ReleaseIPAddress(release_request1)
resp = self._stub.ListAllocatedIPs(self._block_msg)
self.assertEqual(len(resp.ip_list), 0)
def allocate_ip_handler(client, args):
try:
sid_msg = SIDUtils.to_pb(args.sid)
except ValueError:
print("Invalid SubscriberID format: %s" % args.sid)
return
request = AllocateIPRequest()
if int(args.version) == 4:
request.version = AllocateIPRequest.IPV4
elif int(args.version) == 6:
request.version = AllocateIPRequest.IPV6
else:
print("Error: IP version %d is not supported yet" % args.version)
return
request.sid.CopyFrom(sid_msg)
request.apn = args.apn
response = client.AllocateIPAddress(request)
ip_list_msg = response.ip_list
for ip_msg in ip_list_msg:
if ip_msg.version == IPAddress.IPV4:
ip = ipaddress.IPv4Address(ip_msg.address)
print("IPv4 address allocated: %s" % ip)
elif ip_msg.version == IPAddress.IPV6:
ip = ipaddress.IPv6Address(ip_msg.address)
print("IPv6 address allocated: %s" % ip)
else:
print("Error: unknown IP version")
def test_get_subscriber_ip_table(self):
""" test GetSubscriberIPTable """
self._stub.AddIPBlock(self._block_msg)
resp = self._stub.GetSubscriberIPTable(Void())
self.assertEqual(len(resp.entries), 0)
alloc_request0 = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4)
ip_msg0 = self._stub.AllocateIPAddress(alloc_request0)
entry0 = SubscriberIPTableEntry(sid=self._sid0, ip=ip_msg0)
resp = self._stub.GetSubscriberIPTable(Void())
self.assertTrue(entry0 in resp.entries)
alloc_request1 = AllocateIPRequest(sid=self._sid1,
version=AllocateIPRequest.IPV4)
ip_msg1 = self._stub.AllocateIPAddress(alloc_request1)
entry1 = SubscriberIPTableEntry(sid=self._sid1, ip=ip_msg1)
resp = self._stub.GetSubscriberIPTable(Void())
self.assertTrue(entry0 in resp.entries)
self.assertTrue(entry1 in resp.entries)
# keep in table after in release
release_request0 = ReleaseIPRequest(sid=self._sid0, ip=ip_msg0)
resp = self._stub.ReleaseIPAddress(release_request0)
resp = self._stub.GetSubscriberIPTable(Void())
self.assertTrue(entry0 in resp.entries)
self.assertTrue(entry1 in resp.entries)
def test_run_out_of_ip(self):
""" should raise RESOURCE_EXHAUSTED when running out of IP """
# The subnet is provisioned with 16 addresses
# Inside ip_address_man.py 11 addresses are reserved,
# 2 addresses are not usable (all zeros and all ones)
# Thus, we have a usable pool of 3 IP addresses;
# first three allocations should succeed, while the fourth
# request should raise RESOURCE_EXHAUSTED error
self._stub.AddIPBlock(self._block_msg)
request = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4,
apn=self._apn0)
self._stub.AllocateIPAddress(request)
request.apn = self._apn1
self._stub.AllocateIPAddress(request)
request.sid.CopyFrom(self._sid1)
self._stub.AllocateIPAddress(request)
request.sid.CopyFrom(self._sid2)
with self.assertRaises(grpc.RpcError) as err:
self._stub.AllocateIPAddress(request)
self.assertEqual(err.exception.code(),
grpc.StatusCode.RESOURCE_EXHAUSTED)
def test_allocate_ip_address(self):
""" test AllocateIPAddress and ListAllocatedIPs """
self._stub.AddIPBlock(self._block_msg)
# allocate 1st IP
request = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4,
apn=self._apn0)
ip_msg0 = self._stub.AllocateIPAddress(request)
self.assertEqual(ip_msg0.ip_list[0].version, AllocateIPRequest.IPV4)
ip0 = ipaddress.ip_address(ip_msg0.ip_list[0].address)
self.assertTrue(ip0 in self._block)
# TODO: uncomment the code below when ip_allocator
# actually rejects with DuplicatedIPAllocationError
# with self.assertRaises(grpc.RpcError) as err:
# self._stub.AllocateIPAddress(request)
# self.assertEqual(err.exception.code(),
# grpc.StatusCode.ALREADY_EXISTS)
# allocate 2nd IP
request.sid.CopyFrom(self._sid1)
ip_msg2 = self._stub.AllocateIPAddress(request)
self.assertEqual(ip_msg2.ip_list[0].version, AllocateIPRequest.IPV4)
ip2 = ipaddress.ip_address(ip_msg2.ip_list[0].address)
self.assertTrue(ip2 in self._block)
def test_release_unknown_sid_apn_ip_tuple(self):
""" releasing unknown sid-apn-ip tuple should raise NOT_FOUND """
self._stub.AddIPBlock(self._block_msg)
alloc_request0 = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4,
apn=self._apn0)
ip_msg0 = self._stub.AllocateIPAddress(alloc_request0)
request = ReleaseIPRequest(sid=SIDUtils.to_pb("IMSI12345"),
ip=ip_msg0.ip_list[0],
apn=self._apn0)
with self.assertRaises(grpc.RpcError) as err:
self._stub.ReleaseIPAddress(request)
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
ip_bytes = bytes(map(int, '10.0.0.0'.split('.')))
request.ip.CopyFrom(IPAddress(version=IPAddress.IPV4,
address=ip_bytes))
with self.assertRaises(grpc.RpcError) as err:
self._stub.ReleaseIPAddress(request)
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
request = ReleaseIPRequest(sid=self._sid0,
ip=ip_msg0.ip_list[0],
apn=self._apn1)
with self.assertRaises(grpc.RpcError) as err:
self._stub.ReleaseIPAddress(request)
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
def test_ipv6_unimplemented(self):
""" ipv6 requests should raise UNIMPLEMENTED """
ip = ipaddress.ip_address("fc::")
block = IPBlock(version=IPBlock.IPV6,
net_address=ip.packed,
prefix_len=120)
# AddIPBlock
with self.assertRaises(grpc.RpcError) as err:
self._stub.AddIPBlock(block)
self.assertEqual(err.exception.code(), grpc.StatusCode.UNIMPLEMENTED)
# ListAllocatedIPs
with self.assertRaises(grpc.RpcError) as err:
self._stub.ListAllocatedIPs(block)
self.assertEqual(err.exception.code(), grpc.StatusCode.UNIMPLEMENTED)
# AllocateIPAddress
request = AllocateIPRequest(sid=self._sid1,
version=AllocateIPRequest.IPV6,
apn=self._apn0)
with self.assertRaises(grpc.RpcError) as err:
self._stub.AllocateIPAddress(request)
self.assertEqual(err.exception.code(), grpc.StatusCode.UNIMPLEMENTED)
# ReleaseIPAddress
release_request = ReleaseIPRequest(
sid=self._sid0,
ip=IPAddress(version=IPAddress.IPV6),
apn=self._apn0)
with self.assertRaises(grpc.RpcError) as err:
self._stub.ReleaseIPAddress(release_request)
self.assertEqual(err.exception.code(), grpc.StatusCode.UNIMPLEMENTED)
def test_get_subscriber_id_from_ip(self):
""" test GetSubscriberIDFromIP """
self._stub.AddIPBlock(self._block_msg)
alloc_request0 = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4)
ip_msg0 = self._stub.AllocateIPAddress(alloc_request0)
sid_pb_returned = self._stub.GetSubscriberIDFromIP(ip_msg0)
self.assertEqual(SIDUtils.to_str(self._sid0),
SIDUtils.to_str(sid_pb_returned))
def test_multiple_apn_ipallocation(self):
""" test AllocateIPAddress for multiple APNs """
self._stub.AddIPBlock(self._block_msg)
# allocate 1st IP
request = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4,
apn=self._apn0)
ip_msg0 = self._stub.AllocateIPAddress(request)
self.assertEqual(ip_msg0.ip_list[0].version, AllocateIPRequest.IPV4)
ip0 = ipaddress.ip_address(ip_msg0.ip_list[0].address)
self.assertTrue(ip0 in self._block)
# allocate 2nd IP from another APN to the same user
request.apn = self._apn1
ip_msg1 = self._stub.AllocateIPAddress(request)
self.assertEqual(ip_msg1.ip_list[0].version, AllocateIPRequest.IPV4)
ip1 = ipaddress.ip_address(ip_msg1.ip_list[0].address)
self.assertTrue(ip1 in self._block)
def allocate_ip_handler(client, args):
try:
sid_msg = SIDUtils.to_pb(args.sid)
except ValueError:
print("Invalid SubscriberID format: %s" % args.sid)
return
request = AllocateIPRequest()
request.version = AllocateIPRequest.IPV4
request.sid.CopyFrom(sid_msg)
ip_msg = client.AllocateIPAddress(request)
if ip_msg.version == IPAddress.IPV4:
ip = ipaddress.IPv4Address(ip_msg.address)
print("IPv4 address allocated: %s" % ip)
elif ip_msg.version == IPAddress.IPV6:
ip = ipaddress.IPv6Address(ip_msg.address)
print("IPv6 address allocated: %s" % ip)
else:
print("Error: unknown IP version")
def _build_allocate_ip_data(num_subs: int):
active_sids = _load_subs(num_subs)
allocate_ip_reqs = []
for sid in active_sids:
ip_req = AllocateIPRequest(sid=sid, version=AllocateIPRequest.IPV4,
apn='magma.ipv4') # hardcoding APN
ip_req_dict = json_format.MessageToDict(ip_req)
# Dumping AllocateIP request into json
allocate_ip_reqs.append(ip_req_dict)
with open('allocate_data.json', 'w') as file:
json.dump(allocate_ip_reqs, file, separators=(',', ':'))
def test_get_ip_for_subscriber(self):
""" test GetIPForSubscriber """
self._stub.AddIPBlock(self._block_msg)
alloc_request0 = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4)
ip_msg0 = self._stub.AllocateIPAddress(alloc_request0)
ip0 = ipaddress.ip_address(ip_msg0.address)
ip_msg0_returned = self._stub.GetIPForSubscriber(self._sid0)
ip0_returned = ipaddress.ip_address(ip_msg0_returned.address)
self.assertEqual(ip0, ip0_returned)
def test_run_out_of_ip(self):
""" should raise RESOURCE_EXHAUSTED when running out of IP """
self._stub.AddIPBlock(self._block_msg)
request = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4)
self._stub.AllocateIPAddress(request)
request.sid.CopyFrom(self._sid1)
self._stub.AllocateIPAddress(request)
request.sid.CopyFrom(self._sid2)
with self.assertRaises(grpc.RpcError) as err:
self._stub.AllocateIPAddress(request)
self.assertEqual(err.exception.code(),
grpc.StatusCode.RESOURCE_EXHAUSTED)
def test_ipv6(self):
""" ipv6 requests should work for allocate / release IP requests """
# AllocateIPAddress
request = AllocateIPRequest(sid=self._sid1,
version=AllocateIPRequest.IPV6,
apn=self._apn0)
ip_msg = self._stub.AllocateIPAddress(request)
self.assertTrue(
ipaddress.ip_address(ip_msg.ip_list[0].address) in
self._ipv6_block)
# ReleaseIPAddress
release_request = ReleaseIPRequest(sid=self._sid1,
ip=ip_msg.ip_list[0],
apn=self._apn0)
self._stub.ReleaseIPAddress(release_request)
def test_list_allocated_ips(self):
""" test list allocated IPs from a IP block """
self._stub.AddIPBlock(self._block_msg)
# list empty allocated IPs
resp = self._stub.ListAllocatedIPs(self._block_msg)
self.assertEqual(len(resp.ip_list), 0)
# list after allocating one IP
request = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4)
ip_msg0 = self._stub.AllocateIPAddress(request)
resp = self._stub.ListAllocatedIPs(self._block_msg)
self.assertNotEqual(resp, None)
tmp = ListAllocatedIPsResponse()
tmp.ip_list.extend([ip_msg0])
self.assertEqual(resp, tmp)
def test_remove_allocated_assigned_block_with_force(self):
""" remove should return nothing """
self._stub.AddIPBlock(self._block_msg)
alloc_request0 = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4)
self._stub.AllocateIPAddress(alloc_request0)
remove_request0 = RemoveIPBlockRequest(ip_blocks=[self._block_msg],
force=True)
resp = self._stub.RemoveIPBlock(remove_request0)
expect = RemoveIPBlockResponse()
expect.ip_blocks.extend([self._block_msg])
self.assertEqual(expect, resp)
resp = self._stub.ListAddedIPv4Blocks(Void())
expect = ListAddedIPBlocksResponse()
self.assertEqual(expect, resp)
def test_allocate_ip_address(self):
""" test AllocateIPAddress and ListAllocatedIPs """
self._stub.AddIPBlock(self._block_msg)
# allocate 1st IP
request = AllocateIPRequest(sid=self._sid0,
version=AllocateIPRequest.IPV4)
ip_msg0 = self._stub.AllocateIPAddress(request)
self.assertEqual(ip_msg0.version, AllocateIPRequest.IPV4)
ip0 = ipaddress.ip_address(ip_msg0.address)
self.assertTrue(ip0 in self._block)
# duplicated request should fail
with self.assertRaises(grpc.RpcError) as err:
self._stub.AllocateIPAddress(request)
self.assertEqual(err.exception.code(), grpc.StatusCode.ALREADY_EXISTS)
# allocate 2nd IP
request.sid.CopyFrom(self._sid1)
ip_msg2 = self._stub.AllocateIPAddress(request)
self.assertEqual(ip_msg2.version, AllocateIPRequest.IPV4)
ip2 = ipaddress.ip_address(ip_msg2.address)
self.assertTrue(ip2 in self._block)
