1 files changed, 55 insertions, 12 deletions

diff --git a/tests/sepol_wrap.cpp b/tests/sepol_wrap.cpp
index 8fea2d5b..d537b7e0 100644
--- a/tests/sepol_wrap.cpp
+++ b/tests/sepol_wrap.cpp
@@ -181,7 +181,7 @@ void *load_policy(const char *policy_path)
 
 /* items needed to iterate over the avtab */
 struct avtab_iter {
-    avtab_t avtab;
+    avtab_t *avtab;
     uint32_t i;
     avtab_ptr_t cur;
 };
@@ -198,9 +198,9 @@ static int get_avtab_allow_rule(char *out, size_t max_size, policydb_t *db,
 {
     size_t len;
 
-    for (; avtab_i->i < avtab_i->avtab.nslot; (avtab_i->i)++) {
+    for (; avtab_i->i < avtab_i->avtab->nslot; (avtab_i->i)++) {
         if (avtab_i->cur == NULL) {
-            avtab_i->cur = avtab_i->avtab.htable[avtab_i->i];
+            avtab_i->cur = avtab_i->avtab->htable[avtab_i->i];
         }
         for (; avtab_i->cur; avtab_i->cur = (avtab_i->cur)->next) {
             if (!((avtab_i->cur)->key.specified & AVTAB_ALLOWED)) continue;
@@ -233,6 +233,37 @@ int get_allow_rule(char *out, size_t len, void *policydbp, void *avtab_iterp)
     return get_avtab_allow_rule(out, len, db, avtab_i);
 }
 
+static avtab_iter *init_avtab_common(avtab_t *in)
+{
+    struct avtab_iter *out = (struct avtab_iter *)
+                            calloc(1, sizeof(struct avtab_iter));
+    if (!out) {
+        std::cerr << "Failed to allocate avtab iterator" << std::endl;
+        return NULL;
+    }
+
+    out->avtab = in;
+    return out;
+}
+
+void *init_avtab(void *policydbp)
+{
+    policydb_t *p = static_cast<policydb_t *>(policydbp);
+    return static_cast<void *>(init_avtab_common(&p->te_avtab));
+}
+
+void *init_cond_avtab(void *policydbp)
+{
+    policydb_t *p = static_cast<policydb_t *>(policydbp);
+    return static_cast<void *>(init_avtab_common(&p->te_cond_avtab));
+}
+
+void destroy_avtab(void *avtab_iterp)
+{
+    struct avtab_iter *avtab_i = static_cast<struct avtab_iter *>(avtab_iterp);
+    free(avtab_i);
+}
+
 /*
  * <sepol/policydb/expand.h->conditional.h> uses 'bool' as a variable name
  * inside extern "C" { .. } construct, which clang doesn't like.
@@ -240,45 +271,57 @@ int get_allow_rule(char *out, size_t len, void *policydbp, void *avtab_iterp)
  */
 extern "C" int expand_avtab(policydb_t *p, avtab_t *a, avtab_t *expa);
 
-static avtab_iter *init_avtab_common(avtab_t *in, policydb_t *p)
+static avtab_iter *init_expanded_avtab_common(avtab_t *in, policydb_t *p)
 {
     struct avtab_iter *out = (struct avtab_iter *)
                             calloc(1, sizeof(struct avtab_iter));
     if (!out) {
+        std::cerr << "Failed to allocate avtab iterator" << std::endl;
+        return NULL;
+    }
+
+    avtab_t *avtab = (avtab_t *) calloc(1, sizeof(avtab_t));
+
+    if (!avtab) {
         std::cerr << "Failed to allocate avtab" << std::endl;
+        free(out);
         return NULL;
     }
 
-    if (avtab_init(&out->avtab)) {
+    out->avtab = avtab;
+    if (avtab_init(out->avtab)) {
         std::cerr << "Failed to initialize avtab" << std::endl;
+        free(avtab);
         free(out);
         return NULL;
     }
 
-    if (expand_avtab(p, in, &out->avtab)) {
+    if (expand_avtab(p, in, out->avtab)) {
         std::cerr << "Failed to expand avtab" << std::endl;
+        free(avtab);
         free(out);
         return NULL;
     }
     return out;
 }
 
-void *init_avtab(void *policydbp)
+void *init_expanded_avtab(void *policydbp)
 {
     policydb_t *p = static_cast<policydb_t *>(policydbp);
-    return static_cast<void *>(init_avtab_common(&p->te_avtab, p));
+    return static_cast<void *>(init_expanded_avtab_common(&p->te_avtab, p));
 }
 
-void *init_cond_avtab(void *policydbp)
+void *init_expanded_cond_avtab(void *policydbp)
 {
     policydb_t *p = static_cast<policydb_t *>(policydbp);
-    return static_cast<void *>(init_avtab_common(&p->te_cond_avtab, p));
+    return static_cast<void *>(init_expanded_avtab_common(&p->te_cond_avtab, p));
 }
 
-void destroy_avtab(void *avtab_iterp)
+void destroy_expanded_avtab(void *avtab_iterp)
 {
     struct avtab_iter *avtab_i = static_cast<struct avtab_iter *>(avtab_iterp);
-    avtab_destroy(&avtab_i->avtab);
+    avtab_destroy(avtab_i->avtab);
+    free(avtab_i->avtab);
     free(avtab_i);
 }