s390/kaslr: generalize and improve random base distribution

Improve the distribution algorithm of random base address to ensure
a uniformity among all suitable addresses. To generate a random value
once, and to build a continuous range in which every value is suitable,
count all the suitable addresses (referred to as positions) that can be
used as a base address. The positions are counted by iterating over the
usable memory ranges. For each range that is big enough to accommodate
the image, count all the suitable addresses where the image can be placed,
while taking reserved memory ranges into consideration.

A new function "iterate_valid_positions()" has dual purpose. Firstly, it
is called to count the positions in a given memory range, and secondly,
to convert a random position back to an address.

"get_random_base()" has been replaced with more generic
"randomize_within_range()" which now could be called for randomizing
base addresses not just for the kernel image.

Acked-by: Alexander Gordeev <agordeev@linux.ibm.com>
Signed-off-by: Vasily Gorbik <gor@linux.ibm.com>

4 files changed, 112 insertions, 11 deletions

diff --git a/arch/s390/boot/boot.h b/arch/s390/boot/boot.h
index 5baa855c21cb..222c6886acf6 100644
--- a/arch/s390/boot/boot.h
+++ b/arch/s390/boot/boot.h
@@ -54,6 +54,7 @@ unsigned long physmem_alloc_top_down(enum reserved_range_type type, unsigned lon
 unsigned long physmem_alloc_range(enum reserved_range_type type, unsigned long size,
 				  unsigned long align, unsigned long min, unsigned long max,
 				  bool die_on_oom);
+unsigned long get_physmem_alloc_pos(void);
 bool ipl_report_certs_intersects(unsigned long addr, unsigned long size,
 				 unsigned long *intersection_start);
 bool is_ipl_block_dump(void);
@@ -66,7 +67,8 @@ void verify_facilities(void);
 void print_missing_facilities(void);
 void sclp_early_setup_buffer(void);
 void print_pgm_check_info(void);
-unsigned long get_random_base(void);
+unsigned long randomize_within_range(unsigned long size, unsigned long align,
+				     unsigned long min, unsigned long max);
 void setup_vmem(unsigned long asce_limit);
 void __printf(1, 2) decompressor_printk(const char *fmt, ...);
 void print_stacktrace(unsigned long sp);
diff --git a/arch/s390/boot/kaslr.c b/arch/s390/boot/kaslr.c
index 71f75f03f800..90602101e2ae 100644
--- a/arch/s390/boot/kaslr.c
+++ b/arch/s390/boot/kaslr.c
@@ -91,16 +91,108 @@ static int get_random(unsigned long limit, unsigned long *value)
 	return 0;
 }
 
-unsigned long get_random_base(void)
+static void sort_reserved_ranges(struct reserved_range *res, unsigned long size)
 {
-	unsigned long vmlinux_size = vmlinux.image_size + vmlinux.bss_size;
-	unsigned long minimal_pos = vmlinux.default_lma + vmlinux_size;
-	unsigned long random;
+	struct reserved_range tmp;
+	int i, j;
 
-	/* [vmlinux.default_lma + vmlinux.image_size + vmlinux.bss_size : physmem_info.usable] */
-	if (get_random(physmem_info.usable - minimal_pos, &random))
-		return 0;
+	for (i = 1; i < size; i++) {
+		tmp = res[i];
+		for (j = i - 1; j >= 0 && res[j].start > tmp.start; j--)
+			res[j + 1] = res[j];
+		res[j + 1] = tmp;
+	}
+}
+
+static unsigned long iterate_valid_positions(unsigned long size, unsigned long align,
+					     unsigned long _min, unsigned long _max,
+					     struct reserved_range *res, size_t res_count,
+					     bool pos_count, unsigned long find_pos)
+{
+	unsigned long start, end, tmp_end, range_pos, pos = 0;
+	struct reserved_range *res_end = res + res_count;
+	struct reserved_range *skip_res;
+	int i;
+
+	align = max(align, 8UL);
+	_min = round_up(_min, align);
+	for_each_physmem_usable_range(i, &start, &end) {
+		if (_min >= end)
+			continue;
+		start = round_up(start, align);
+		if (start >= _max)
+			break;
+		start = max(_min, start);
+		end = min(_max, end);
+
+		while (start + size <= end) {
+			/* skip reserved ranges below the start */
+			while (res && res->end <= start) {
+				res++;
+				if (res >= res_end)
+					res = NULL;
+			}
+			skip_res = NULL;
+			tmp_end = end;
+			/* has intersecting reserved range */
+			if (res && res->start < end) {
+				skip_res = res;
+				tmp_end = res->start;
+			}
+			if (start + size <= tmp_end) {
+				range_pos = (tmp_end - start - size) / align + 1;
+				if (pos_count) {
+					pos += range_pos;
+				} else {
+					if (range_pos >= find_pos)
+						return start + (find_pos - 1) * align;
+					find_pos -= range_pos;
+				}
+			}
+			if (!skip_res)
+				break;
+			start = round_up(skip_res->end, align);
+		}
+	}
 
-	return physmem_alloc_range(RR_VMLINUX, vmlinux_size, THREAD_SIZE,
-				   vmlinux.default_lma, minimal_pos + random, false);
+	return pos_count ? pos : 0;
+}
+
+/*
+ * Two types of decompressor memory allocations/reserves are considered
+ * differently.
+ *
+ * "Static" or "single" allocations are done via physmem_alloc_range() and
+ * physmem_reserve(), and they are listed in physmem_info.reserved[]. Each
+ * type of "static" allocation can only have one allocation per type and
+ * cannot have chains.
+ *
+ * On the other hand, "dynamic" or "repetitive" allocations are done via
+ * physmem_alloc_top_down(). These allocations are tightly packed together
+ * top down from the end of online memory. physmem_alloc_pos represents
+ * current position where those allocations start.
+ *
+ * Functions randomize_within_range() and iterate_valid_positions()
+ * only consider "dynamic" allocations by never looking above
+ * physmem_alloc_pos. "Static" allocations, however, are explicitly
+ * considered by checking the "res" (reserves) array. The first
+ * reserved_range of a "dynamic" allocation may also be checked along the
+ * way, but it will always be above the maximum value anyway.
+ */
+unsigned long randomize_within_range(unsigned long size, unsigned long align,
+				     unsigned long min, unsigned long max)
+{
+	struct reserved_range res[RR_MAX];
+	unsigned long max_pos, pos;
+
+	memcpy(res, physmem_info.reserved, sizeof(res));
+	sort_reserved_ranges(res, ARRAY_SIZE(res));
+	max = min(max, get_physmem_alloc_pos());
+
+	max_pos = iterate_valid_positions(size, align, min, max, res, ARRAY_SIZE(res), true, 0);
+	if (!max_pos)
+		return 0;
+	if (get_random(max_pos, &pos))
+		return 0;
+	return iterate_valid_positions(size, align, min, max, res, ARRAY_SIZE(res), false, pos + 1);
 }
diff --git a/arch/s390/boot/physmem_info.c b/arch/s390/boot/physmem_info.c
index 4ee9b7381142..0cf79826eef9 100644
--- a/arch/s390/boot/physmem_info.c
+++ b/arch/s390/boot/physmem_info.c
@@ -321,3 +321,8 @@ unsigned long physmem_alloc_top_down(enum reserved_range_type type, unsigned lon
 	physmem_alloc_ranges = ranges_left;
 	return addr;
 }
+
+unsigned long get_physmem_alloc_pos(void)
+{
+	return physmem_alloc_pos;
+}
diff --git a/arch/s390/boot/startup.c b/arch/s390/boot/startup.c
index 45c30c2b1a7a..638a6cf8a075 100644
--- a/arch/s390/boot/startup.c
+++ b/arch/s390/boot/startup.c
@@ -316,7 +316,9 @@ void startup_kernel(void)
 	rescue_initrd(safe_addr, ident_map_size);
 
 	if (kaslr_enabled()) {
-		vmlinux_lma = get_random_base();
+		vmlinux_lma = randomize_within_range(vmlinux.image_size + vmlinux.bss_size,
+						     THREAD_SIZE, vmlinux.default_lma,
+						     ident_map_size);
 		if (vmlinux_lma) {
 			__kaslr_offset = vmlinux_lma - vmlinux.default_lma;
 			offset_vmlinux_info(__kaslr_offset);