def test_id_generator():
min_id = 10
max_id = 20
id_gen = IdGenerator(min_id, max_id)
test_id = 5
assert not id_gen.reserve_id(test_id)
test_id = 25
assert not id_gen.reserve_id(test_id)
test_id = 15
assert not id_gen.is_id_used(test_id)
assert id_gen.reserve_id(test_id)
assert id_gen.is_id_used(test_id)
assert id_gen.release_id(test_id)
assert not id_gen.release_id(test_id)
assert not id_gen.is_id_used(test_id)
for idx in range(min_id, max_id + 1):
new_id = id_gen.get_new_id()
assert idx == new_id
print 'ID generator where all IDs were assigned:'
print str(id_gen)
assert id_gen.get_new_id() is None
assert not id_gen.reserve_id(test_id)
for idx in range(min_id, max_id + 1):
assert id_gen.release_id(idx)
print 'ID generator where all IDs were released:'
print str(id_gen)
assert id_gen.reserve_id(min_id)
assert id_gen.release_id(min_id)
def test_id_alloc_resv():
id_gen = IdGenerator()
max_id = 100
alloc_id_set = set()
for exp_id in range(1, max_id + 1):
alloc_id = id_gen.get_new_id()
assert exp_id == alloc_id
assert alloc_id not in alloc_id_set
assert id_gen.is_id_used(alloc_id)
assert id_gen.max_resved_id == alloc_id
alloc_id_set.add(alloc_id)
for exp_id in range(max_id + 1, max_id * 2 + 1):
assert id_gen.reserve_id(exp_id)
assert id_gen.is_id_used(exp_id)
# reserve used ID will fail
assert not id_gen.reserve_id(exp_id)
assert id_gen.max_resved_id == exp_id
id_list = list(alloc_id_set)
rel_id_set = set()
for idx in range(max_id / 2):
test_id = get_rand_list_item(id_list, True)
assert id_gen.release_id(test_id)
# release non-existent ID will fail
assert not id_gen.release_id(test_id)
assert not id_gen.is_id_used(test_id)
rel_id_set.add(test_id)
sorted_list = sorted(rel_id_set)
for exp_id in sorted_list:
alloc_id = id_gen.get_new_id()
# test if ID allocation is always from the lowest available one
assert exp_id == alloc_id
assert id_gen.max_resved_id == max_id * 2
for test_id in range(1, max_id * 2 + 1):
assert id_gen.release_id(test_id)
assert len(id_gen.resved_ids) == 0
assert id_gen.avail_id == 1
assert id_gen.max_resved_id is None
def test_rule_list():
id_gen = IdGenerator()
prefix_len_list = [8, 16, 24, 32]
action_list = [cfg.VrfIncomingSvcsRule.RULE_ACTION_ALLOW, cfg.VrfIncomingSvcsRule.RULE_ACTION_DENY]
src_ip_list = []
for idx in range(1, 250):
src_ip_list.append('2.2.2.' + str(idx))
rule_list = cfg.VrfIncomingSvcsRuleList()
af = socket.AF_INET
rule1 = cfg.VrfIncomingSvcsRule(cfg.VrfIncomingSvcsRule.RULE_TYPE_ACL, 'test_vrf',
cfg.VrfIncomingSvcsRule.RULE_ACTION_ALLOW, af,
src_ip = socket.inet_pton(af, '1.1.1.0'), prefix_len = 24,
protocol = cfg.VrfIncomingSvcsRule.RULE_PROTO_TCP,
dst_port = 8080, seq_num = 100)
assert rule_list.insert(rule1) is None
rule1.rule_id = 5
assert rule_list.insert(rule1) == 0
assert len(rule_list) == 1
new_rule = copy.deepcopy(rule1)
# duplicate insert is not allowed
assert rule_list.insert(new_rule) is None
new_rule.rule_id = None
# delete rule by rule attribute match
assert rule_list.remove(new_rule) is not None
assert len(rule_list) == 0
rule1.rule_id = 10
assert rule_list.insert(rule1) == 0
assert rule_list.remove_by_id(10) is not None
assert len(rule_list) == 0
assert len(rule_list.seq_num_list) == 0
assert len(rule_list.rule_id_map) == 0
test_rule_num = 20
rule_id_list = []
while test_rule_num > 0:
src_ip = get_rand_list_item(src_ip_list, True)
if src_ip is None:
break
prefix_len = get_rand_list_item(prefix_len_list)
action = get_rand_list_item(action_list)
rule = cfg.VrfIncomingSvcsRule(cfg.VrfIncomingSvcsRule.RULE_TYPE_ACL, "test_vrf", action, af,
src_ip = socket.inet_pton(socket.AF_INET, src_ip),
prefix_len = prefix_len, seq_num = random.randint(1, 1000),
rule_id = id_gen.get_new_id())
print 'Add rule to list: %s' % rule
idx = rule_list.insert(rule)
assert idx is not None
check_rule_list(rule_list)
rule_id_list.append(rule.rule_id)
test_rule_num -= 1
assert len(rule_list) == 20
assert len(rule_id_list) == 20
print 'Rule IDs: %s' % rule_id_list
# clone rule list
new_list = copy.deepcopy(rule_list)
assert len(new_list) == 20
# delete all rules
while len(rule_id_list) > 0:
rule_id = get_rand_list_item(rule_id_list, True)
print 'Delete rule ID %d' % rule_id
assert rule_list.remove_by_id(rule_id) is not None
check_rule_list(rule_list)
assert len(rule_list) == 0
assert len(rule_list.seq_num_list) == 0
assert len(rule_list.rule_id_map) == 0
print 'All rules were deleted'
new_list.clear()
assert len(new_list) == 0
assert len(new_list.seq_num_list) == 0
assert len(new_list.rule_id_map) == 0
class VrfIncomingSvcsRuleCache:
# map: af, vrf_name => rule_list
acl_rules = {socket.AF_INET: {}, socket.AF_INET6: {}}
ip_rules = {socket.AF_INET: {}, socket.AF_INET6: {}}
mutex = threading.RLock()
id_generator = IdGenerator()
@classmethod
def insert_rule(cls, rule):
log_info('Handling add rule: %s' % rule)
cls.mutex.acquire()
rule_list = cls.ip_rules[
rule.af] if rule.rule_type == rule.RULE_TYPE_IP else cls.acl_rules[
rule.af]
if rule.vrf_name not in rule_list:
rule_list[rule.vrf_name] = VrfIncomingSvcsRuleList()
if rule.rule_id is None:
rule.rule_id = cls.id_generator.get_new_id()
if rule.rule_id is None:
log_err('Could not generate new rule ID')
log_info(str(cls.id_generator))
cls.mutex.release()
return False
else:
if cls.id_generator.is_id_used(rule.rule_id):
log_err('Given rule ID %d is used' % rule.rule_id)
log_info(str(cls.id_generator))
cls.mutex.release()
return False
if not cls.id_generator.reserve_id(rule.rule_id):
log_err('Failed to reserve rule ID %d' % rule.rule_id)
log_info(str(cls.id_generator))
cls.mutex.release()
return False
ret_val = True
idx = rule_list[rule.vrf_name].insert(rule)
if idx is not None:
if not IptablesHandler.proc_rule('insert', rule, idx):
log_err('Failed to call iptables to insert ACL rule')
# rollback
rule_list[rule.vrf_name].remove(rule)
ret_val = False
else:
log_err('Failed to insert rule to cache')
cls.id_generator.release_id(rule.rule_id)
ret_val = False
cls.mutex.release()
if ret_val:
log_info('Rule added, ID=%d' % rule.rule_id)
return ret_val
@classmethod
def delete_rule(cls, rule):
log_info('Handling delete rule: %s' % rule)
cls.mutex.acquire()
rule_list = cls.ip_rules[
rule.af] if rule.rule_type == rule.RULE_TYPE_IP else cls.acl_rules[
rule.af]
if rule.vrf_name not in rule_list:
log_err('VRF name %s not found in cache' % rule.vrf_name)
cls.mutex.release()
return False
ret_val = rule_list[rule.vrf_name].remove(rule)
if ret_val is None:
log_err('Failed to delete rule from cache')
cls.mutex.release()
return False
del_rule, idx = ret_val
ret_val = True
if not IptablesHandler.proc_rule('delete', del_rule, idx):
log_err('Failed to call iptables to delete ACL rule')
# rollback
rule_list[del_rule.vrf_name].insert(del_rule)
ret_val = False
if len(rule_list[rule.vrf_name]) == 0:
del rule_list[rule.vrf_name]
cls.mutex.release()
if ret_val:
if not cls.id_generator.release_id(del_rule.rule_id):
log_err('Failed to release rule ID %d' % del_rule.rule_id)
log_info(str(cls.id_generator))
log_info('Rule deleted')
return ret_val
@classmethod
def find_rule_by_id(cls, rule_id):
ret_val = None
cls.mutex.acquire()
for af in [socket.AF_INET, socket.AF_INET6]:
for vrf_name, rule_list in cls.ip_rules[af].items():
if rule_id in rule_list.rule_id_map:
idx = rule_list.rule_id_map[rule_id]
ret_val = (rule_list[idx], idx)
break
if ret_val is not None:
break
for vrf_name, rule_list in cls.acl_rules[af].items():
if rule_id in rule_list.rule_id_map:
idx = rule_list.rule_id_map[rule_id]
ret_val = (rule_list[idx], idx)
break
if ret_val is not None:
break
cls.mutex.release()
return ret_val
@classmethod
def find_rule_by_match(cls, rule):
cls.mutex.acquire()
rule_list = cls.ip_rules[
rule.af] if rule.rule_type == rule.RULE_TYPE_IP else cls.acl_rules[
rule.af]
if rule.vrf_name not in rule_list:
cls.mutex.release()
return None
try:
idx = rule_list[rule.vrf_name].index(rule)
except ValueError:
cls.mutex.release()
return None
return rule_list[rule.vrf_name][idx]
@classmethod
def delete_rule_by_id(cls, rule_id):
log_info('Handling delete rule by ID: %d' % rule_id)
cls.mutex.acquire()
ret_val = True
found_rule = cls.find_rule_by_id(rule_id)
if found_rule is not None:
del_rule, idx = found_rule
rule_list = cls.ip_rules[del_rule.af] \
if del_rule.rule_type == del_rule.RULE_TYPE_IP else cls.acl_rules[del_rule.af]
if rule_list[del_rule.vrf_name].remove_by_id(rule_id) is None:
log_err('Failed to remove rule with ID %d' % rule_id)
cls.mutex.release()
return False
if not IptablesHandler.proc_rule('delete', del_rule, idx):
log_err('Failed to call iptables to delete ACL rule')
# rollback
rule_list = cls.ip_rules[del_rule.af] \
if del_rule.rule_type == del_rule.RULE_TYPE_IP else cls.acl_rules[del_rule.af]
rule_list[del_rule.vrf_name].insert(del_rule)
ret_val = False
if len(rule_list[del_rule.vrf_name]) == 0:
del rule_list[del_rule.vrf_name]
else:
log_err('Rule ID %d not found for delete' % rule_id)
ret_val = False
cls.mutex.release()
if ret_val:
if not cls.id_generator.release_id(rule_id):
log_err('Failed to release rule ID %d' % rule_id)
log_info(str(cls.id_generator))
log_info('Rule deleted')
return ret_val
@classmethod
def replace_rule(cls, rule_id, new_rule):
log_info('Handling replace rule with ID %d' % rule_id)
cls.mutex.acquire()
ret_val = True
found_rule = cls.find_rule_by_id(rule_id)
if found_rule is not None:
old_rule, idx = found_rule
rule_list = cls.ip_rules[old_rule.af] \
if old_rule.rule_type == old_rule.RULE_TYPE_IP else cls.acl_rules[old_rule.af]
if not IptablesHandler.proc_rule('replace', new_rule, idx):
log_err('Failed to call iptables to replace ACL rule')
ret_val = False
else:
rule_list[old_rule.vrf_name][idx] = new_rule
else:
log_err('Rule ID %d not found for replace' % rule_id)
ret_val = False
cls.mutex.release()
if ret_val:
log_info('Rule replaced')
return ret_val
@classmethod
def update_rule(cls, rule_id, **params):
log_info('Handling update rule: ID %d param %s' % (rule_id, params))
cls.mutex.acquire()
found_rule = cls.find_rule_by_id(rule_id)
if found_rule is None:
log_err('Rule ID %d not found for update' % rule_id)
cls.mutex.release()
return False
old_rule, _ = found_rule
upd_rule = copy.deepcopy(old_rule)
changed_attrs = set()
def check_and_set(rule, key, val):
try:
orig_val = getattr(rule, key)
except AttributeError:
orig_val = None
if val != orig_val:
if val is not None:
setattr(rule, key, val)
else:
delattr(rule, key)
changed_attrs.add(key)
for key, val in params.items():
if val is not None:
check_and_set(upd_rule, key, val)
if len(changed_attrs) == 0:
log_info('There is no change to be updated, just return')
cls.mutex.release()
return True
log_info('old_rule: %s' % old_rule)
log_info('new_rule: %s' % upd_rule)
match_rule = cls.find_rule_by_match(upd_rule)
if match_rule is not None and match_rule.rule_id != old_rule.rule_id:
log_err('The updating rule already exists')
cls.mutex.release()
return False
repl_attrs = {
'src_ip', 'prefix_len', 'protocol', 'dst_port', 'action', 'in_intf'
}
no_repl_attrs = changed_attrs.difference(repl_attrs)
if len(no_repl_attrs) > 0:
log_info(
'Non-replacable attributes %s changed, delete and add rule' %
no_repl_attrs)
if not cls.delete_rule_by_id(rule_id):
log_err('Failed to delete existing rule by ID')
cls.mutex.release()
return False
ret_val = cls.insert_rule(upd_rule)
if not ret_val:
log_err('Failed to insert updated rule')
cls.insert_rule(old_rule)
else:
log_info('Only replacable attributes changed, just replace rule')
ret_val = cls.replace_rule(rule_id, upd_rule)
if not ret_val:
log_err('Failed to replace rule')
cls.mutex.release()
if ret_val:
log_info('Rule updated')
return ret_val
@classmethod
def clear_all_rules(cls, flt_vrf=None, flt_af=None):
log_info('Handling clear all rules for VRF %s and AF %s' %
(flt_vrf if flt_vrf is not None else '-',
str(flt_af) if flt_af is not None else '-'))
cls.mutex.acquire()
for af in [socket.AF_INET, socket.AF_INET6]:
if flt_af is not None and flt_af != af:
continue
for vrf_name, rule_list in cls.ip_rules[af].items():
if flt_vrf is not None and flt_vrf != vrf_name:
continue
for rule in rule_list:
IptablesHandler.proc_rule('delete', rule)
cls.id_generator.release_id(rule.rule_id)
rule_list.clear()
del cls.ip_rules[af][vrf_name]
for vrf_name, rule_list in cls.acl_rules[af].items():
if flt_vrf is not None and flt_vrf != vrf_name:
continue
for rule in rule_list:
IptablesHandler.proc_rule('delete', rule)
cls.id_generator.release_id(rule.rule_id)
rule_list.clear()
del cls.acl_rules[af][vrf_name]
cls.mutex.release()
@classmethod
def dump_rules(cls):
cls.mutex.acquire()
for af in [socket.AF_INET, socket.AF_INET6]:
for vrf_name, rule_list in cls.ip_rules[af].items():
log_info('-------------------------------------')
log_info(
' %s IP Rules of VRF %s' %
('IPv4' if af == socket.AF_INET else 'IPv6', vrf_name))
log_info('-------------------------------------')
log_info('\n%s\n' % rule_list)
for vrf_name, rule_list in cls.acl_rules[af].items():
log_info('-------------------------------------')
log_info(
' %s ACL Rules of VRF %s' %
('IPv4' if af == socket.AF_INET else 'IPv6', vrf_name))
log_info('-------------------------------------')
log_info('\n%s\n' % rule_list)
cls.mutex.release()
@classmethod
def get_all_rules(cls, rule_type=None, vrf_name=None, **params):
log_info('Handling get rules for vrf %s' %
(vrf_name if vrf_name is not None else 'ALL'))
cls.mutex.acquire()
ret_list = []
for af in [socket.AF_INET, socket.AF_INET6]:
if rule_type is None or rule_type == VrfIncomingSvcsRule.RULE_TYPE_IP:
for vrf, rule_list in cls.ip_rules[af].items():
if vrf_name is not None and vrf != vrf_name:
continue
for rule in rule_list:
if not rule.match(**params):
continue
ret_list.append(copy.deepcopy(rule))
if rule_type is None or rule_type == VrfIncomingSvcsRule.RULE_TYPE_ACL:
for vrf, rule_list in cls.acl_rules[af].items():
if vrf_name is not None and vrf != vrf_name:
continue
for rule in rule_list:
if not rule.match(**params):
continue
ret_list.append(copy.deepcopy(rule))
cls.mutex.release()
return ret_list