Exemplo n.º 1
0
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
Exemplo n.º 2
0
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
Exemplo n.º 3
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