3 files changed, 39 insertions, 13 deletions

diff --git a/security/selinux/ss/policydb.c b/security/selinux/ss/policydb.c
index 674ddfe0ba03..3a29704be8ce 100644
--- a/security/selinux/ss/policydb.c
+++ b/security/selinux/ss/policydb.c
@@ -31,6 +31,7 @@
 #include <linux/string.h>
 #include <linux/errno.h>
 #include <linux/audit.h>
+#include <linux/flex_array.h>
 #include "security.h"
 
 #include "policydb.h"
@@ -739,11 +740,17 @@ void policydb_destroy(struct policydb *p)
         hashtab_map(p->range_tr, range_tr_destroy, NULL);
         hashtab_destroy(p->range_tr);
 
-        if (p->type_attr_map) {
-                for (i = 0; i < p->p_types.nprim; i++)
-                        ebitmap_destroy(&p->type_attr_map[i]);
+        if (p->type_attr_map_array) {
+                for (i = 0; i < p->p_types.nprim; i++) {
+                        struct ebitmap *e;
+
+                        e = flex_array_get(p->type_attr_map_array, i);
+                        if (!e)
+                                continue;
+                        ebitmap_destroy(e);
+                }
+                flex_array_free(p->type_attr_map_array);
         }
-        kfree(p->type_attr_map);
         ebitmap_destroy(&p->policycaps);
         ebitmap_destroy(&p->permissive_map);
 
@@ -2257,19 +2264,33 @@ int policydb_read(struct policydb *p, void *fp)
         if (rc)
                 goto bad;
 
-        p->type_attr_map = kmalloc(p->p_types.nprim * sizeof(struct ebitmap), GFP_KERNEL);
-        if (!p->type_attr_map)
+        rc = -ENOMEM;
+        p->type_attr_map_array = flex_array_alloc(sizeof(struct ebitmap),
+                                                  p->p_types.nprim,
+                                                  GFP_KERNEL | __GFP_ZERO);
+        if (!p->type_attr_map_array)
+                goto bad;
+
+        /* preallocate so we don't have to worry about the put ever failing */
+        rc = flex_array_prealloc(p->type_attr_map_array, 0, p->p_types.nprim - 1,
+                                 GFP_KERNEL | __GFP_ZERO);
+        if (rc)
                 goto bad;
 
         for (i = 0; i < p->p_types.nprim; i++) {
-                ebitmap_init(&p->type_attr_map[i]);
+                struct ebitmap *e = flex_array_get(p->type_attr_map_array, i);
+
+                BUG_ON(!e);
+                ebitmap_init(e);
                 if (p->policyvers >= POLICYDB_VERSION_AVTAB) {
-                        if (ebitmap_read(&p->type_attr_map[i], fp))
+                        rc = ebitmap_read(e, fp);
+                        if (rc)
                                 goto bad;
                 }
                 /* add the type itself as the degenerate case */
-                if (ebitmap_set_bit(&p->type_attr_map[i], i, 1))
-                                goto bad;
+                rc = ebitmap_set_bit(e, i, 1);
+                if (rc)
+                        goto bad;
         }
 
         rc = policydb_bounds_sanity_check(p);
diff --git a/security/selinux/ss/policydb.h b/security/selinux/ss/policydb.h
index 26d9adf8542b..310e94442cb8 100644
--- a/security/selinux/ss/policydb.h
+++ b/security/selinux/ss/policydb.h
@@ -24,6 +24,8 @@
 #ifndef _SS_POLICYDB_H_
 #define _SS_POLICYDB_H_
 
+#include <linux/flex_array.h>
+
 #include "symtab.h"
 #include "avtab.h"
 #include "sidtab.h"
@@ -246,7 +248,7 @@ struct policydb {
         struct hashtab *range_tr;
 
         /* type -> attribute reverse mapping */
-        struct ebitmap *type_attr_map;
+        struct flex_array *type_attr_map_array;
 
         struct ebitmap policycaps;
 
diff --git a/security/selinux/ss/services.c b/security/selinux/ss/services.c
index 1de60ce90d9a..9ea2feca3cd4 100644
--- a/security/selinux/ss/services.c
+++ b/security/selinux/ss/services.c
@@ -50,6 +50,7 @@
 #include <linux/audit.h>
 #include <linux/mutex.h>
 #include <linux/selinux.h>
+#include <linux/flex_array.h>
 #include <net/netlabel.h>
 
 #include "flask.h"
@@ -626,8 +627,10 @@ static void context_struct_compute_av(struct context *scontext,
          */
         avkey.target_class = tclass;
         avkey.specified = AVTAB_AV;
-        sattr = &policydb.type_attr_map[scontext->type - 1];
-        tattr = &policydb.type_attr_map[tcontext->type - 1];
+        sattr = flex_array_get(policydb.type_attr_map_array, scontext->type - 1);
+        BUG_ON(!sattr);
+        tattr = flex_array_get(policydb.type_attr_map_array, tcontext->type - 1);
+        BUG_ON(!tattr);
         ebitmap_for_each_positive_bit(sattr, snode, i) {
                 ebitmap_for_each_positive_bit(tattr, tnode, j) {
                         avkey.source_type = i + 1;