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