printf("First hex is %x\n", *((unsigned int *)addr));
 }
 
-static void write_bytes(char *addr)
+static void write_bytes(char *addr, size_t length)
 {
        unsigned long i;
 
-       for (i = 0; i < LENGTH; i++)
+       for (i = 0; i < length; i++)
                *(addr + i) = (char)i;
 }
 
-static int read_bytes(char *addr)
+static int read_bytes(char *addr, size_t length)
 {
        unsigned long i;
 
        check_bytes(addr);
-       for (i = 0; i < LENGTH; i++)
+       for (i = 0; i < length; i++)
                if (*(addr + i) != (char)i) {
                        printf("Mismatch at %lu\n", i);
                        return 1;
 
        printf("Returned address is %p\n", addr);
        check_bytes(addr);
-       write_bytes(addr);
-       ret = read_bytes(addr);
+       write_bytes(addr, length);
+       ret = read_bytes(addr, length);
 
        /* munmap() length of MAP_HUGETLB memory must be hugepage aligned */
-       if (munmap(addr, LENGTH)) {
+       if (munmap(addr, length)) {
                perror("munmap");
                exit(1);
        }