npdmtool: don't error with missing debug flags, error with multiple set (#50)

Version field was incorrectly labeled "process_category". This is now supported (and defaults to 0 if not present on npdm, 1 if not present on kip). process_category is alias.

Support was added for mesosphere large-address map extension (these bits are reserved in official OS).

Support was added for specifying the signature key generation, which determines modulus used to verify ACID.

Support was added for system call capabilities in range [0x80, 0xBF], which kernel allows since 11.0.0.

"title_id" (and min/max) were renamed to program_id. If program_id not present, title_id used as alias.

configure.ac

@@ -2,7 +2,7 @@
 # Process this file with autoconf to produce a configure script.
 
 AC_PREREQ(2.61)
-AC_INIT([switch-tools],[1.4.1],[https://github.com/switchbrew/switch-tools/issues])
+AC_INIT([switch-tools],[1.13.1],[https://github.com/switchbrew/switch-tools/issues])
 AC_CONFIG_SRCDIR([src/build_pfs0.c])
 
 AM_INIT_AUTOMAKE([subdir-objects])
@@ -31,7 +31,7 @@ NET_LIBS=""
 case "$host" in
    *-*-mingw*)
       NET_LIBS="-lws2_32"
-      CFLAGS="$CFLAGS -D__USE_MINGW_ANSI_STDIO"
+      CFLAGS="$CFLAGS -D__USE_MINGW_ANSI_STDIO -D_WIN32_WINNT=0x0600"
       ;;
 esac
 

src/elf2kip.c

@@ -9,11 +9,6 @@
 #include "blz.h"
 #include "elf64.h"
 
-typedef uint64_t u64;
-typedef uint32_t u32;
-typedef uint16_t u16;
-typedef uint8_t u8;
-
 typedef struct {
     u32 DstOff;
     u32 DecompSz;
@@ -24,8 +19,8 @@ typedef struct {
 typedef struct {
     u8  Magic[4];
     u8  Name[0xC];
-    u64 TitleId;
+    u64 ProgramId;
-    u32 ProcessCategory;
+    u32 Version;
     u8  MainThreadPriority;
     u8  DefaultCpuId;
     u8  Unk;
@@ -60,6 +55,17 @@ uint8_t* ReadEntireFile(const char* fn, size_t* len_out) {
     return buf;
 }
 
+int cJSON_GetString(const cJSON *obj, const char *field, const char **out) {
+  const cJSON *config = cJSON_GetObjectItemCaseSensitive(obj, field);
+  if (cJSON_IsString(config)) {
+    *out = config->valuestring;
+    return 1;
+  } else {
+    fprintf(stderr, "Failed to get %s (field not present).\n", field);
+    return 0;
+  }
+}
+
 int cJSON_GetU8(const cJSON *obj, const char *field, u8 *out) {
     const cJSON *config = cJSON_GetObjectItemCaseSensitive(obj, field);
     if (cJSON_IsNumber(config)) {
@@ -121,10 +127,11 @@ int cJSON_GetBooleanOptional(const cJSON *obj, const char *field, int *out) {
             fprintf(stderr, "Unknown boolean value in %s.\n", field);
             return 0;
         }
+        return 1;
     } else {
         *out = 0;
+        return 0;
     }
-    return 1;
 }
 
 int cJSON_GetU64(const cJSON *obj, const char *field, u64 *out) {
@@ -178,6 +185,32 @@ int cJSON_GetU64FromObjectValue(const cJSON *config, u64 *out) {
     }
 }
 
+int cJSON_GetU32(const cJSON *obj, const char *field, u32 *out) {
+    const cJSON *config = cJSON_GetObjectItemCaseSensitive(obj, field);
+    if (cJSON_IsString(config) && (config->valuestring != NULL)) {
+        char *endptr = NULL;
+        *out = strtoul(config->valuestring, &endptr, 16);
+        if (config->valuestring == endptr) {
+            fprintf(stderr, "Failed to get %s (empty string)\n", field);
+            return 0;
+        } else if (errno == ERANGE) {
+            fprintf(stderr, "Failed to get %s (value out of range)\n", field);
+            return 0;
+        } else if (errno == EINVAL) {
+            fprintf(stderr, "Failed to get %s (not base16 string)\n", field);
+            return 0;
+        } else if (errno) {
+            fprintf(stderr, "Failed to get %s (unknown error)\n", field);
+            return 0;
+        } else {
+            return 1;
+        }
+    } else {
+        fprintf(stderr, "Failed to get %s (field not present).\n", field);
+        return 0;
+    }
+}
+
 int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
     const cJSON *capability = NULL;
     const cJSON *capabilities = NULL;
@@ -202,12 +235,34 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
         goto PARSE_CAPS_END;
     }
 
-    /* Parse title_id. */
+    /* Parse program_id (or deprecated title_id). */
-    if (!cJSON_GetU64(npdm_json, "title_id", &kip_hdr->TitleId)) {
+    if (!cJSON_GetU64(npdm_json, "program_id", &kip_hdr->ProgramId) && !cJSON_GetU64(npdm_json, "title_id", &kip_hdr->ProgramId)) {
         status = 0;
         goto PARSE_CAPS_END;
     }
 
+    /* Parse use secure memory. */
+    /* This field is optional, and defaults to true (set before this function is called). */
+    int use_secure_memory = 1;
+    if (cJSON_GetBooleanOptional(npdm_json, "use_secure_memory", &use_secure_memory)) {
+        if (use_secure_memory) {
+            kip_hdr->Flags |= 0x20;
+        } else {
+            kip_hdr->Flags &= ~0x20;
+        }
+    }
+
+    /* Parse immortality. */
+    /* This field is optional, and defaults to true (set before this function is called). */
+    int immortal = 1;
+    if (cJSON_GetBooleanOptional(npdm_json, "immortal", &immortal)) {
+        if (immortal) {
+            kip_hdr->Flags |= 0x40;
+        } else {
+            kip_hdr->Flags &= ~0x40;
+        }
+    }
+
     /* Parse main_thread_stack_size. */
     u64 stack_size = 0;
     if (!cJSON_GetU64(npdm_json, "main_thread_stack_size", &stack_size)) {
@@ -230,26 +285,43 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
         status = 0;
         goto PARSE_CAPS_END;
     }
-    if (!cJSON_GetU8(npdm_json, "process_category", (u8 *)&kip_hdr->ProcessCategory)) {
+    if (!cJSON_GetU32(npdm_json, "version", &kip_hdr->Version) && !cJSON_GetU32(npdm_json, "process_category", &kip_hdr->Version)) { // optional
-        status = 0;
+        kip_hdr->Version = 1;
-        goto PARSE_CAPS_END;
     }
 
     /* Parse capabilities. */
     capabilities = cJSON_GetObjectItemCaseSensitive(npdm_json, "kernel_capabilities");
-    if (!cJSON_IsObject(capabilities)) {
+    if (!(cJSON_IsArray(capabilities) || cJSON_IsObject(capabilities))) {
-        fprintf(stderr, "Kernel Capabilities must be an object!\n");
+        fprintf(stderr, "Kernel Capabilities must be an array!\n");
         status = 0;
         goto PARSE_CAPS_END;
     }
 
+    int kac_obj = 0;
+    if (cJSON_IsObject(capabilities)) {
+        kac_obj = 1;
+        fprintf(stderr, "Using deprecated kernel_capabilities format. Please turn it into an array.\n");
+    }
+
     u32 cur_cap = 0;
     u32 desc;
     cJSON_ArrayForEach(capability, capabilities) {
         desc = 0;
-        const char *type_str = capability->string;
+        const char *type_str;
+        const cJSON *value;
+
+        if (kac_obj) {
+            type_str = capability->string;
+            value = capability;
+        } else {
+            if (!cJSON_GetString(capability, "type", &type_str)) {
+                status = 0;
+                goto PARSE_CAPS_END;
+            }
+            value = cJSON_GetObjectItemCaseSensitive(capability, "value");
+        }
+
 
-        const cJSON *value = capability;
         if (!strcmp(type_str, "kernel_flags")) {
             if (cur_cap + 1 > 0x20) {
                 fprintf(stderr, "Error: Too many capabilities!\n");
@@ -269,13 +341,17 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
                 status = 0;
                 goto PARSE_CAPS_END;
             }
+
+            u8 real_highest_prio = (lowest_prio < highest_prio) ? lowest_prio : highest_prio;
+            u8 real_lowest_prio  = (lowest_prio > highest_prio) ? lowest_prio : highest_prio;
+
             desc = highest_cpu;
             desc <<= 8;
             desc |= lowest_cpu;
             desc <<= 6;
-            desc |= (lowest_prio & 0x3F);
+            desc |= (real_highest_prio & 0x3F);
             desc <<= 6;
-            desc |= (highest_prio & 0x3F);
+            desc |= (real_lowest_prio & 0x3F);
             kip_hdr->Capabilities[cur_cap++] = (u32)((desc << 4) | (0x0007));
         } else if (!strcmp(type_str, "syscalls")) {
             if (!cJSON_IsObject(value)) {
@@ -284,7 +360,7 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
                 goto PARSE_CAPS_END;
             }
             u32 num_descriptors;
-            u32 descriptors[6] = {0}; /* alignup(0x80/0x18); */
+            u32 descriptors[8] = {0}; /* alignup(0xC0/0x18); */
             char field_name[8] = {0};
             const cJSON *cur_syscall = NULL;
             u64 syscall_value = 0;
@@ -297,14 +373,14 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
                     goto PARSE_CAPS_END;
                 }
 
-                if (syscall_value >= 0x80) {
+                if (syscall_value >= 0xC0) {
-                    fprintf(stderr, "Error: All syscall entries must be numbers in [0, 0x7F]\n");
+                    fprintf(stderr, "Error: All syscall entries must be numbers in [0, 0xBF]\n");
                     status = 0;
                     goto PARSE_CAPS_END;
                 }
                 descriptors[syscall_value / 0x18] |= (1UL << (syscall_value % 0x18));
             }
-            for (unsigned int i = 0; i < 6; i++) {
+            for (unsigned int i = 0; i < 8; i++) {
                 if (descriptors[i]) {
                     if (cur_cap + 1 > 0x20) {
                         fprintf(stderr, "Error: Too many capabilities!\n");
@@ -341,7 +417,8 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
             desc |= is_ro << 24;
             kip_hdr->Capabilities[cur_cap++] = (u32)((desc << 7) | (0x003F));
 
-            desc = (u32)((map_size >> 12) & 0x00FFFFFFULL);
+            desc = (u32)((map_size >> 12) & 0x000FFFFFULL);
+            desc |= (u32)(((map_address >> 36) & 0xFULL) << 20);
             is_io ^= 1;
             desc |= is_io << 24;
             kip_hdr->Capabilities[cur_cap++] = (u32)((desc << 7) | (0x003F));
@@ -358,6 +435,47 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
             }
             desc = (u32)((page_address >> 12) & 0x00FFFFFFULL);
             kip_hdr->Capabilities[cur_cap++] = (u32)((desc << 8) | (0x007F));
+        } else if (!strcmp(type_str, "map_region")) {
+            if (cur_cap + 1 > 0x20) {
+                fprintf(stderr, "Error: Too many capabilities!\n");
+                status = 0;
+                goto PARSE_CAPS_END;
+            }
+            if (!cJSON_IsArray(value)) {
+                fprintf(stderr, "Map Region capability value must be array!\n");
+                status = 0;
+                goto PARSE_CAPS_END;
+            }
+            u8 regions[3] = {0};
+            int is_ro[3] = {0};
+            const cJSON *cur_region = NULL;
+            int index = 0;
+            cJSON_ArrayForEach(cur_region, value) {
+                if (index >= 3) {
+                    fprintf(stderr, "Too many region descriptors!\n");
+                    status = 0;
+                    goto PARSE_CAPS_END;
+                }
+                if (!cJSON_IsObject(cur_region)) {
+                    fprintf(stderr, "Region descriptor value must be object!\n");
+                    status = 0;
+                    goto PARSE_CAPS_END;
+                }
+
+                if (!cJSON_GetU8(cur_region, "region_type", &regions[index]) ||
+                    !cJSON_GetBoolean(cur_region, "is_ro", &is_ro[index])) {
+                    status = 0;
+                    goto PARSE_CAPS_END;
+                }
+
+                index++;
+            }
+
+            u32 capability = 0x3FF;
+            for (int i = 0; i < 3; ++i) {
+                capability |= ((regions[i] & 0x3F) | ((is_ro[i] & 1) << 6)) << (11 + 7 * i);
+            }
+            kip_hdr->Capabilities[cur_cap++] = capability;
         } else if (!strcmp(type_str, "irq_pair")) {
             if (cur_cap + 1 > 0x20) {
                 fprintf(stderr, "Error: Too many capabilities!\n");
@@ -370,17 +488,18 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
                 goto PARSE_CAPS_END;
             }
             const cJSON *irq = NULL;
+            int desc_idx = 0;
             cJSON_ArrayForEach(irq, value) {
-                desc <<= 10;
                 if (cJSON_IsNull(irq)) {
-                    desc |= 0x3FF;
+                    desc |= 0x3FF << desc_idx;
                 } else if (cJSON_IsNumber(irq)) {
-                    desc |= ((u16)(irq->valueint)) & 0x3FF;
+                    desc |= (((u16)(irq->valueint)) & 0x3FF) << desc_idx;
                 } else {
                     fprintf(stderr, "Failed to parse IRQ value.\n");
                     status = 0;
                     goto PARSE_CAPS_END;
                 }
+                desc_idx += 10;
             }
             kip_hdr->Capabilities[cur_cap++] = (u32)((desc << 12) | (0x07FF));
         } else if (!strcmp(type_str, "application_type")) {
@@ -435,6 +554,7 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
             }
             int allow_debug = 0;
             int force_debug = 0;
+            int force_debug_prod = 0;
             if (!cJSON_GetBoolean(value, "allow_debug", &allow_debug)) {
                 status = 0;
                 goto PARSE_CAPS_END;
@@ -443,8 +563,16 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
                 status = 0;
                 goto PARSE_CAPS_END;
             }
-            desc = (allow_debug & 1) | ((force_debug & 1) << 1);
+            if (!cJSON_GetBoolean(value, "force_debug_prod", &force_debug_prod)) {
+                status = 0;
+                goto PARSE_CAPS_END;
+            }
+            desc = (allow_debug & 1) | ((force_debug_prod & 1) << 1) | ((force_debug & 1) << 2);
             kip_hdr->Capabilities[cur_cap++] = (u32)((desc << 17) | (0xFFFF));
+        } else {
+            fprintf(stderr, "Error: unknown capability %s\n", type_str);
+            status = 0;
+            goto PARSE_CAPS_END;
         }
     }
 
@@ -452,6 +580,7 @@ int ParseKipConfiguration(const char *json, KipHeader *kip_hdr) {
         kip_hdr->Capabilities[i] = 0xFFFFFFFF;
     }
 
+    status = 1;
     PARSE_CAPS_END:
     cJSON_Delete(npdm_json);
     return status;
@@ -465,7 +594,7 @@ int main(int argc, char* argv[]) {
 
     KipHeader kip_hdr = {0};
     memcpy(kip_hdr.Magic, "KIP1", 4);
-    kip_hdr.Flags = 0x3F;
+    kip_hdr.Flags = 0x7F;
 
     if (sizeof(KipHeader) != 0x100) {
         fprintf(stderr, "Bad compile environment!\n");
@@ -479,7 +608,10 @@ int main(int argc, char* argv[]) {
         return EXIT_FAILURE;
     }
 
-    ParseKipConfiguration(json, &kip_hdr);
+    if (!ParseKipConfiguration(json, &kip_hdr)) {
+        fprintf(stderr, "Failed to parse kip configuration!\n");
+        return EXIT_FAILURE;
+    }
 
     size_t elf_len;
     uint8_t* elf = ReadEntireFile(argv[1], &elf_len);

src/elf2nro.c

@@ -8,10 +8,6 @@
 #include "elf64.h"
 #include "romfs.h"
 
-typedef uint64_t u64;
-typedef uint32_t u32;
-typedef uint8_t u8;
-
 typedef struct {
     u32 FileOff;
     u32 Size;
@@ -25,9 +21,9 @@ typedef struct {
 
 typedef struct {
     u8  Magic[4];
-    u32 Unk1;
+    u32 version;
     u32 size;
-    u32 Unk2;
+    u32 flags;
     NsoSegment Segments[3];
     u32 bssSize;
     u32 Unk3;
@@ -82,6 +78,7 @@ int main(int argc, char* argv[]) {
         fprintf(stderr, "--nacp=<control.nacp> Embeds control.nacp into the output file.\n");
         fprintf(stderr, "--romfs=<image> Embeds RomFS into the output file.\n");
         fprintf(stderr, "--romfsdir=<directory> Builds and embeds RomFS into the output file.\n");
+        fprintf(stderr, "--alignedheader Sets the \"AlignedHeader\" flag in the output file.\n");
         return EXIT_FAILURE;
     }
 
@@ -106,11 +103,13 @@ int main(int argc, char* argv[]) {
 
     int argi;
     char* icon_path = NULL, *nacp_path = NULL, *romfs_path = NULL, *romfs_dir_path = NULL;
+    u32 aligned_header = 0;
     for (argi=3; argi<argc; argi++) {
         if (strncmp(argv[argi], "--icon=", 7)==0) icon_path = &argv[argi][7];
         if (strncmp(argv[argi], "--nacp=", 7)==0) nacp_path = &argv[argi][7];
         if (strncmp(argv[argi], "--romfs=", 8)==0) romfs_path = &argv[argi][8];
         if (strncmp(argv[argi], "--romfsdir=", 11)==0) romfs_dir_path = &argv[argi][11];
+        if (strncmp(argv[argi], "--alignedheader", 15)==0) aligned_header = 1;
     }
 
     if (romfs_dir_path != NULL && romfs_path != NULL) {
@@ -185,6 +184,25 @@ int main(int argc, char* argv[]) {
         file_off = (file_off + 0xFFF) & ~0xFFF;
     }
 
+    /* Iterate over sections to find build id. */
+    size_t cur_sect_hdr_ofs = hdr->e_shoff;
+    for (unsigned int i = 0; i < hdr->e_shnum; i++) {
+        Elf64_Shdr *cur_shdr = (Elf64_Shdr *)(elf + cur_sect_hdr_ofs);
+        if (cur_shdr->sh_type == SHT_NOTE) {
+            Elf64_Nhdr *note_hdr = (Elf64_Nhdr *)(elf + cur_shdr->sh_offset);
+            u8 *note_name = (u8 *)((uintptr_t)note_hdr + sizeof(Elf64_Nhdr));
+            u8 *note_desc = note_name + note_hdr->n_namesz;
+            if (note_hdr->n_type == NT_GNU_BUILD_ID && note_hdr->n_namesz == 4 && memcmp(note_name, "GNU\x00", 4) == 0) {
+                size_t build_id_size = note_hdr->n_descsz;
+                if (build_id_size > 0x20) {
+                    build_id_size = 0x20;
+                }
+                memcpy(nro_hdr.BuildId, note_desc, build_id_size);
+            }
+        }
+        cur_sect_hdr_ofs += hdr->e_shentsize;
+    }
+
     FILE* out = fopen(argv[2], "wb");
 
     if (out == NULL) {
@@ -194,6 +212,9 @@ int main(int argc, char* argv[]) {
 
     nro_hdr.size = file_off;
 
+    nro_hdr.version = 0;
+    nro_hdr.flags   = (aligned_header << 0);
+
     // TODO check retvals
 
     for (i=0; i<3; i++)

src/elf2nso.c

@@ -158,6 +158,25 @@ int main(int argc, char* argv[]) {
         file_off += comp_sz[i];
     }
     
+    /* Iterate over sections to find build id. */
+    size_t cur_sect_hdr_ofs = hdr->e_shoff;
+    for (unsigned int i = 0; i < hdr->e_shnum; i++) {
+        Elf64_Shdr *cur_shdr = (Elf64_Shdr *)(elf + cur_sect_hdr_ofs);
+        if (cur_shdr->sh_type == SHT_NOTE) {
+            Elf64_Nhdr *note_hdr = (Elf64_Nhdr *)(elf + cur_shdr->sh_offset);
+            u8 *note_name = (u8 *)((uintptr_t)note_hdr + sizeof(Elf64_Nhdr));
+            u8 *note_desc = note_name + note_hdr->n_namesz;
+            if (note_hdr->n_type == NT_GNU_BUILD_ID && note_hdr->n_namesz == 4 && memcmp(note_name, "GNU\x00", 4) == 0) {
+                size_t build_id_size = note_hdr->n_descsz;
+                if (build_id_size > 0x20) {
+                    build_id_size = 0x20;
+                }
+                memcpy(nso_hdr.BuildId, note_desc, build_id_size);
+            }
+        }
+        cur_sect_hdr_ofs += hdr->e_shentsize;
+    }
+
     FILE* out = fopen(argv[2], "wb");
 
     if (out == NULL) {

src/elf_common.h

@@ -995,5 +995,6 @@ typedef struct {
 #define	R_X86_64_TPOFF32	23	/* Offset in static TLS block */
 #define	R_X86_64_IRELATIVE	37
 
+#define NT_GNU_BUILD_ID     3   /* Note type for .note.gnu.build-id */
 
 #endif /* !_SYS_ELF_COMMON_H_ */

src/nacptool.c

@@ -2,6 +2,7 @@
 #include <stdlib.h>
 #include <stdint.h>
 #include <string.h>
+#include <inttypes.h>
 
 typedef uint64_t u64;
 typedef uint32_t u32;
@@ -79,7 +80,7 @@ int main(int argc, char* argv[]) {
     int argi;
     u64 titleid=0;
     for (argi=6; argi<argc; argi++) {
-        if (strncmp(argv[argi], "--titleid=", 10)==0) sscanf(&argv[argi][10], "%016lX", &titleid);
+        if (strncmp(argv[argi], "--titleid=", 10)==0) sscanf(&argv[argi][10], "%016" SCNx64, &titleid);
     }
 
     for (i=0; i<12; i++) {//These are UTF-8.

src/npdmtool.c

@@ -19,9 +19,15 @@ typedef uint8_t u8;
 /* FAC, FAH need to be tightly packed. */
 #pragma pack(push, 1)
 typedef struct {
-    u32 Version;
+    u8 Version;
+    u8 CoiCount;
+    u8 SdoiCount;
+    u8 pad;
     u64 Perms;
-    u8 _0xC[0x20];
+    u64 CoiMin;
+    u64 CoiMax;
+    u64 SdoiMin;
+    u64 SdoiMax;
 } FilesystemAccessControl;
 #pragma pack(pop)
 
@@ -29,17 +35,17 @@ typedef struct {
 typedef struct {
     u32 Version;
     u64 Perms;
-    u32 _0xC;
+    u32 CoiOffset;
-    u32 _0x10;
+    u32 CoiSize;
-    u32 _0x14;
+    u32 SdoiOffset;
-    u32 _0x18;
+    u32 SdoiSize;
 } FilesystemAccessHeader;
 #pragma pack(pop)
 
 typedef struct {
     u32 Magic;
     u8 _0x4[0xC];
-    u64 TitleId;
+    u64 ProgramId;
     u64 _0x18;
     u32 FahOffset;
     u32 FahSize;
@@ -57,8 +63,8 @@ typedef struct {
     u32 Size;
     u32 _0x208;
     u32 Flags;
-    u64 TitleIdRangeMin;
+    u64 ProgramIdRangeMin;
-    u64 TitleIdRangeMax;
+    u64 ProgramIdRangeMax;
     u32 FacOffset;
     u32 FacSize;
     u32 SacOffset;
@@ -70,16 +76,19 @@ typedef struct {
 
 typedef struct {
     u32 Magic;
-    u32 _0x4;
+    u32 SignatureKeyGeneration;
     u32 _0x8;
     u8 MmuFlags;
     u8 _0xD;
     u8 MainThreadPriority;
     u8 DefaultCpuId;
-    u64 _0x10;
+    u32 _0x10;
-    u32 ProcessCategory;
+    u32 SystemResourceSize;
+    u32 Version;
     u32 MainThreadStackSize;
-    char Name[0x50];
+    char Name[0x10];
+    char ProductCode[0x10];
+    u8 _0x40[0x30];
     u32 Aci0Offset;
     u32 Aci0Size;
     u32 AcidOffset;
@@ -113,6 +122,17 @@ uint8_t* ReadEntireFile(const char* fn, size_t* len_out) {
     return buf;
 }
 
+int cJSON_GetString(const cJSON *obj, const char *field, const char **out) {
+  const cJSON *config = cJSON_GetObjectItemCaseSensitive(obj, field);
+  if (cJSON_IsString(config)) {
+    *out = config->valuestring;
+    return 1;
+  } else {
+    fprintf(stderr, "Failed to get %s (field not present).\n", field);
+    return 0;
+  }
+}
+
 int cJSON_GetU8(const cJSON *obj, const char *field, u8 *out) {
     const cJSON *config = cJSON_GetObjectItemCaseSensitive(obj, field);
     if (cJSON_IsNumber(config)) {
@@ -206,6 +226,32 @@ int cJSON_GetU64(const cJSON *obj, const char *field, u64 *out) {
     }
 }
 
+int cJSON_GetU32(const cJSON *obj, const char *field, u32 *out) {
+    const cJSON *config = cJSON_GetObjectItemCaseSensitive(obj, field);
+    if (cJSON_IsString(config) && (config->valuestring != NULL)) {
+        char *endptr = NULL;
+        *out = strtoul(config->valuestring, &endptr, 16);
+        if (config->valuestring == endptr) {
+            fprintf(stderr, "Failed to get %s (empty string)\n", field);
+            return 0;
+        } else if (errno == ERANGE) {
+            fprintf(stderr, "Failed to get %s (value out of range)\n", field);
+            return 0;
+        } else if (errno == EINVAL) {
+            fprintf(stderr, "Failed to get %s (not base16 string)\n", field);
+            return 0;
+        } else if (errno) {
+            fprintf(stderr, "Failed to get %s (unknown error)\n", field);
+            return 0;
+        } else {
+            return 1;
+        }
+    } else {
+        fprintf(stderr, "Failed to get %s (field not present).\n", field);
+        return 0;
+    }
+}
+
 int cJSON_GetU64FromObjectValue(const cJSON *config, u64 *out) {
     if (cJSON_IsString(config) && (config->valuestring != NULL)) {
         char *endptr = NULL;
@@ -244,6 +290,10 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
     const cJSON *service = NULL;
     const cJSON *services = NULL;
     const cJSON *fsaccess = NULL;
+    const cJSON *cois = NULL;
+    const cJSON *coi = NULL;
+    const cJSON *sdois = NULL;
+    const cJSON *sdoi = NULL;
 
     int status = 0;
     cJSON *npdm_json = cJSON_Parse(json);
@@ -292,10 +342,14 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
         status = 0;
         goto NPDM_BUILD_END;
     }
-    if (!cJSON_GetU8(npdm_json, "process_category", (u8 *)&header.ProcessCategory)) {
+
-        status = 0;
+    cJSON_GetU32(npdm_json, "system_resource_size", &header.SystemResourceSize); // optional
-        goto NPDM_BUILD_END;
+
+    /* Get version (deprecated name "process_category"). */
+    if (!cJSON_GetU32(npdm_json, "version", &header.Version) && !cJSON_GetU32(npdm_json, "process_category", &header.Version)) { // optional
+        header.Version = 0;
     }
+
     if (!cJSON_GetU8(npdm_json, "address_space_type", (u8 *)&header.MmuFlags)) {
         status = 0;
         goto NPDM_BUILD_END;
@@ -309,6 +363,33 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
     }
     header.MmuFlags |= is_64_bit;
 
+    int optimize_memory_allocation; // optional
+    if (cJSON_GetBoolean(npdm_json, "optimize_memory_allocation", &optimize_memory_allocation)) {
+        header.MmuFlags |= ((optimize_memory_allocation & 1) << 4);
+    }
+
+    int disable_device_address_space_merge; // optional
+    if (cJSON_GetBoolean(npdm_json, "disable_device_address_space_merge", &disable_device_address_space_merge)) {
+        header.MmuFlags |= ((disable_device_address_space_merge & 1) << 5);
+    }
+
+    int enable_alias_region_extra_size; // optional
+    if (cJSON_GetBoolean(npdm_json, "enable_alias_region_extra_size", &enable_alias_region_extra_size)) {
+        header.MmuFlags |= ((enable_alias_region_extra_size & 1) << 6);
+    }
+
+    int prevent_code_reads; // optional
+    if (cJSON_GetBoolean(npdm_json, "prevent_code_reads", &prevent_code_reads)) {
+        header.MmuFlags |= ((prevent_code_reads & 1) << 7);
+    }
+
+    u8 signature_key_generation; // optional
+    if (cJSON_GetU8(npdm_json, "signature_key_generation", &signature_key_generation)) {
+        header.SignatureKeyGeneration = signature_key_generation;
+    } else {
+        header.SignatureKeyGeneration = 0;
+    }
+
     /* ACID. */
     memset(acid->Signature, 0, sizeof(acid->Signature));
     memset(acid->Modulus, 0, sizeof(acid->Modulus));
@@ -326,19 +407,19 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
     }
     acid->Flags |= (pool_partition & 3) << 2;
 
-    if (!cJSON_GetU64(npdm_json, "title_id_range_min", &acid->TitleIdRangeMin)) {
+    if (!cJSON_GetU64(npdm_json, "program_id_range_min", &acid->ProgramIdRangeMin) && !cJSON_GetU64(npdm_json, "title_id_range_min", &acid->ProgramIdRangeMin)) {
         status = 0;
         goto NPDM_BUILD_END;
     }
-    if (!cJSON_GetU64(npdm_json, "title_id_range_max", &acid->TitleIdRangeMax)) {
+    if (!cJSON_GetU64(npdm_json, "program_id_range_max", &acid->ProgramIdRangeMax) && !cJSON_GetU64(npdm_json, "title_id_range_max", &acid->ProgramIdRangeMax)) {
         status = 0;
         goto NPDM_BUILD_END;
     }
 
     /* ACI0. */
     aci0->Magic = MAGIC_ACI0; /* "ACI0" */
-    /* Parse title_id. */
+    /* Parse program_id (or deprecated title_id). */
-    if (!cJSON_GetU64(npdm_json, "title_id", &aci0->TitleId)) {
+    if (!cJSON_GetU64(npdm_json, "program_id", &aci0->ProgramId) && !cJSON_GetU64(npdm_json, "title_id", &aci0->ProgramId)) {
         status = 0;
         goto NPDM_BUILD_END;
     }
@@ -357,6 +438,14 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
         status = 0;
         goto NPDM_BUILD_END;
     }
+
+    fac->CoiMin    = 0;
+    fac->CoiMax    = 0;
+    fac->SdoiMin   = 0;
+    fac->SdoiMax   = 0;
+    fac->CoiCount  = 0;
+    fac->SdoiCount = 0;
+
     acid->FacOffset = sizeof(NpdmAcid);
     acid->FacSize = sizeof(FilesystemAccessControl);
     acid->SacOffset = (acid->FacOffset + acid->FacSize + 0xF) & ~0xF;
@@ -365,42 +454,155 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
     FilesystemAccessHeader *fah = (FilesystemAccessHeader *)((u8 *)aci0 + sizeof(NpdmAci0));
     fah->Version = 1;
     fah->Perms = fac->Perms;
-    fah->_0xC = 0x1C;
+    fah->CoiOffset = sizeof(FilesystemAccessHeader);
-    fah->_0x14 = 0x1C;
+    fah->CoiSize   = 0;
+
+    cois = cJSON_GetObjectItemCaseSensitive(fsaccess, "content_owner_ids");
+    if (cJSON_IsArray(cois)) {
+        u32 *count = (u32 *)((u8 *)fah + fah->CoiOffset);
+        u64 *id = (u64 *)((u8 *)count + sizeof(u32));
+        cJSON_ArrayForEach(coi, cois) {
+            if (!cJSON_GetU64FromObjectValue(coi, id)) {
+                status = 0;
+                goto NPDM_BUILD_END;
+            }
+            ++id;
+            ++(*count);
+        }
+
+        if (*count > 0) {
+            fah->CoiSize = sizeof(u32) + sizeof(u64) * (*count);
+        }
+    }
+
+    fah->SdoiOffset = fah->CoiOffset + fah->CoiSize;
+    fah->SdoiSize   = 0;
+
+    sdois = cJSON_GetObjectItemCaseSensitive(fsaccess, "save_data_owner_ids");
+    if (cJSON_IsArray(sdois)) {
+        u32 *count = (u32 *)((u8 *)fah + fah->SdoiOffset);
+        cJSON_ArrayForEach(sdoi, sdois) {
+            if (!cJSON_IsObject(sdoi)) {
+                status = 0;
+                goto NPDM_BUILD_END;
+            }
+            ++(*count);
+        }
+
+        u8 *accessibility = (u8 *)count + sizeof(u32);
+        u64 *id = (u64 *)(accessibility + (((*count) + 3ULL) & ~3ULL));
+
+        cJSON_ArrayForEach(sdoi, sdois) {
+            if (!cJSON_GetU8(sdoi, "accessibility", accessibility)) {
+                status = 0;
+                goto NPDM_BUILD_END;
+            }
+            if (!cJSON_GetU64(sdoi, "id", id)) {
+                status = 0;
+                goto NPDM_BUILD_END;
+            }
+
+            ++accessibility;
+            ++id;
+        }
+
+        if (*count > 0) {
+            fah->SdoiSize = sizeof(u32) + sizeof(u8) * ((((*count) + 3ULL) & ~3ULL)) + sizeof(u64) * (*count);
+        }
+    }
+
     aci0->FahOffset = sizeof(NpdmAci0);
-    aci0->FahSize = sizeof(FilesystemAccessHeader);
+    aci0->FahSize = sizeof(FilesystemAccessHeader) + fah->CoiSize + fah->SdoiSize;
     aci0->SacOffset = (aci0->FahOffset + aci0->FahSize + 0xF) & ~0xF;
 
     /* Sac. */
-    services = cJSON_GetObjectItemCaseSensitive(npdm_json, "service_access");
+    u8 *sac = (u8*)aci0 + aci0->SacOffset;
-    if (!cJSON_IsObject(services)) {
+    u32 sac_size = 0;
-        fprintf(stderr, "Service Access must be an object!\n");
+
+    services = cJSON_GetObjectItemCaseSensitive(npdm_json, "service_host");
+    if (services != NULL && !cJSON_IsArray(services)) {
+        fprintf(stderr, "Service Host must be an array!\n");
         status = 0;
         goto NPDM_BUILD_END;
     }
 
-    u8 *sac = (u8*)aci0 + aci0->SacOffset;
-    u32 sac_size = 0;
     cJSON_ArrayForEach(service, services) {
-        if (!cJSON_IsBool(service)) {
+        int is_host = 1;
-            fprintf(stderr, "Services must be of form service_name (str) : is_host (bool)\n");
+        char *service_name;
+
+        if (!cJSON_IsString(service)) {
+            fprintf(stderr, "service_access must be an array of string\n");
             status = 0;
             goto NPDM_BUILD_END;
         }
-        int cur_srv_len = strlen(service->string);
+        service_name = service->valuestring;
+
+        int cur_srv_len = strlen(service_name);
         if (cur_srv_len > 8 || cur_srv_len == 0) {
             fprintf(stderr, "Services must have name length 1 <= len <= 8!\n");
             status = 0;
             goto NPDM_BUILD_END;
         }
         u8 ctrl = (u8)(cur_srv_len - 1);
-        if (cJSON_IsTrue(service)) {
+        if (is_host) {
             ctrl |= 0x80;
         }
         sac[sac_size++] = ctrl;
-        memcpy(sac + sac_size, service->string, cur_srv_len);
+        memcpy(sac + sac_size, service_name, cur_srv_len);
         sac_size += cur_srv_len;
     }
+
+    services = cJSON_GetObjectItemCaseSensitive(npdm_json, "service_access");
+    if (!(services == NULL || cJSON_IsObject(services) || cJSON_IsArray(services))) {
+      fprintf(stderr, "Service Access must be an array!\n");
+      status = 0;
+      goto NPDM_BUILD_END;
+    }
+
+    int sac_obj = 0;
+    if (services != NULL && cJSON_IsObject(services)) {
+      sac_obj = 1;
+      fprintf(stderr, "Using deprecated service_access format. Please turn it into an array.\n");
+    }
+
+    cJSON_ArrayForEach(service, services) {
+        int is_host = 0;
+        char *service_name;
+
+        if (sac_obj) {
+            if (!cJSON_IsBool(service)) {
+                fprintf(stderr, "Services must be of form service_name (str) : is_host (bool)\n");
+                status = 0;
+                goto NPDM_BUILD_END;
+            }
+            is_host = cJSON_IsTrue(service);
+            service_name = service->string;
+        } else {
+            if (!cJSON_IsString(service)) {
+                fprintf(stderr, "service_access must be an array of string\n");
+                status = 0;
+                goto NPDM_BUILD_END;
+            }
+            is_host = 0;
+            service_name = service->valuestring;
+        }
+
+        int cur_srv_len = strlen(service_name);
+        if (cur_srv_len > 8 || cur_srv_len == 0) {
+            fprintf(stderr, "Services must have name length 1 <= len <= 8!\n");
+            status = 0;
+            goto NPDM_BUILD_END;
+        }
+        u8 ctrl = (u8)(cur_srv_len - 1);
+        if (is_host) {
+            ctrl |= 0x80;
+        }
+        sac[sac_size++] = ctrl;
+        memcpy(sac + sac_size, service_name, cur_srv_len);
+        sac_size += cur_srv_len;
+    }
+
+
     memcpy((u8 *)acid + acid->SacOffset, sac, sac_size);
     aci0->SacSize = sac_size;
     acid->SacSize = sac_size;
@@ -409,20 +611,37 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
 
     /* Parse capabilities. */
     capabilities = cJSON_GetObjectItemCaseSensitive(npdm_json, "kernel_capabilities");
-    if (!cJSON_IsObject(capabilities)) {
+    if (!(cJSON_IsArray(capabilities) || cJSON_IsObject(capabilities))) {
-        fprintf(stderr, "Kernel Capabilities must be an object!\n");
+        fprintf(stderr, "Kernel Capabilities must be an array!\n");
         status = 0;
         goto NPDM_BUILD_END;
     }
 
+    int kac_obj = 0;
+    if (cJSON_IsObject(capabilities)) {
+        kac_obj = 1;
+        fprintf(stderr, "Using deprecated kernel_capabilities format. Please turn it into an array.\n");
+    }
+
     u32 *caps = (u32 *)((u8 *)aci0 + aci0->KacOffset);
     u32 cur_cap = 0;
     u32 desc;
     cJSON_ArrayForEach(capability, capabilities) {
         desc = 0;
-        const char *type_str = capability->string;
+        const char *type_str;
+        const cJSON *value;
+
+        if (kac_obj) {
+          type_str = capability->string;
+          value = capability;
+        } else {
+          if (!cJSON_GetString(capability, "type", &type_str)) {
+            status = 0;
+            goto NPDM_BUILD_END;
+          }
+          value = cJSON_GetObjectItemCaseSensitive(capability, "value");
+        }
 
-        const cJSON *value = capability;
         if (!strcmp(type_str, "kernel_flags")) {
             if (!cJSON_IsObject(value)) {
                 fprintf(stderr, "Kernel Flags Capability value must be object!\n");
@@ -437,13 +656,17 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
                 status = 0;
                 goto NPDM_BUILD_END;
             }
+
+            u8 real_highest_prio = (lowest_prio < highest_prio) ? lowest_prio : highest_prio;
+            u8 real_lowest_prio  = (lowest_prio > highest_prio) ? lowest_prio : highest_prio;
+
             desc = highest_cpu;
             desc <<= 8;
             desc |= lowest_cpu;
             desc <<= 6;
-            desc |= (lowest_prio & 0x3F);
+            desc |= (real_highest_prio & 0x3F);
             desc <<= 6;
-            desc |= (highest_prio & 0x3F);
+            desc |= (real_lowest_prio & 0x3F);
             caps[cur_cap++] = (u32)((desc << 4) | (0x0007));
         } else if (!strcmp(type_str, "syscalls")) {
             if (!cJSON_IsObject(value)) {
@@ -452,7 +675,7 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
                 goto NPDM_BUILD_END;
             }
             u32 num_descriptors;
-            u32 descriptors[6] = {0}; /* alignup(0x80/0x18); */
+            u32 descriptors[8] = {0}; /* alignup(0xC0/0x18); */
             char field_name[8] = {0};
             const cJSON *cur_syscall = NULL;
             u64 syscall_value = 0;
@@ -465,14 +688,14 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
                     goto NPDM_BUILD_END;
                 }
 
-                if (syscall_value >= 0x80) {
+                if (syscall_value >= 0xC0) {
-                    fprintf(stderr, "Error: All syscall entries must be numbers in [0, 0x7F]\n");
+                    fprintf(stderr, "Error: All syscall entries must be numbers in [0, 0xBF]\n");
                     status = 0;
                     goto NPDM_BUILD_END;
                 }
                 descriptors[syscall_value / 0x18] |= (1UL << (syscall_value % 0x18));
             }
-            for (unsigned int i = 0; i < 6; i++) {
+            for (unsigned int i = 0; i < 8; i++) {
                 if (descriptors[i]) {
                     desc = descriptors[i] | (i << 24);
                     caps[cur_cap++] = (u32)((desc << 5) | (0x000F));
@@ -499,7 +722,8 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
             desc |= is_ro << 24;
             caps[cur_cap++] = (u32)((desc << 7) | (0x003F));
 
-            desc = (u32)((map_size >> 12) & 0x00FFFFFFULL);
+            desc = (u32)((map_size >> 12) & 0x000FFFFFULL);
+            desc |= (u32)(((map_address >> 36) & 0xFULL) << 20);
             is_io ^= 1;
             desc |= is_io << 24;
             caps[cur_cap++] = (u32)((desc << 7) | (0x003F));
@@ -511,6 +735,47 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
             }
             desc = (u32)((page_address >> 12) & 0x00FFFFFFULL);
             caps[cur_cap++] = (u32)((desc << 8) | (0x007F));
+        } else if (!strcmp(type_str, "map_region")) {
+            if (cur_cap + 1 > 0x20) {
+                fprintf(stderr, "Error: Too many capabilities!\n");
+                status = 0;
+                goto NPDM_BUILD_END;
+            }
+            if (!cJSON_IsArray(value)) {
+                fprintf(stderr, "Map Region capability value must be array!\n");
+                status = 0;
+                goto NPDM_BUILD_END;
+            }
+            u8 regions[3] = {0};
+            int is_ro[3] = {0};
+            const cJSON *cur_region = NULL;
+            int index = 0;
+            cJSON_ArrayForEach(cur_region, value) {
+                if (index >= 3) {
+                    fprintf(stderr, "Too many region descriptors!\n");
+                    status = 0;
+                    goto NPDM_BUILD_END;
+                }
+                if (!cJSON_IsObject(cur_region)) {
+                    fprintf(stderr, "Region descriptor value must be object!\n");
+                    status = 0;
+                    goto NPDM_BUILD_END;
+                }
+
+                if (!cJSON_GetU8(cur_region, "region_type", &regions[index]) ||
+                    !cJSON_GetBoolean(cur_region, "is_ro", &is_ro[index])) {
+                    status = 0;
+                    goto NPDM_BUILD_END;
+                }
+
+                index++;
+            }
+
+            u32 capability = 0x3FF;
+            for (int i = 0; i < 3; ++i) {
+                capability |= ((regions[i] & 0x3F) | ((is_ro[i] & 1) << 6)) << (11 + 7 * i);
+            }
+            caps[cur_cap++] = capability;
         } else if (!strcmp(type_str, "irq_pair")) {
             if (!cJSON_IsArray(value) || cJSON_GetArraySize(value) != 2) {
                 fprintf(stderr, "Error: IRQ Pairs must have size 2 array value.\n");
@@ -518,17 +783,18 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
                 goto NPDM_BUILD_END;
             }
             const cJSON *irq = NULL;
+            int desc_idx = 0;
             cJSON_ArrayForEach(irq, value) {
-                desc <<= 10;
                 if (cJSON_IsNull(irq)) {
-                    desc |= 0x3FF;
+                    desc |= 0x3FF << desc_idx;
                 } else if (cJSON_IsNumber(irq)) {
-                    desc |= ((u16)(irq->valueint)) & 0x3FF;
+                    desc |= (((u16)(irq->valueint)) & 0x3FF) << desc_idx;
                 } else {
                     fprintf(stderr, "Failed to parse IRQ value.\n");
                     status = 0;
                     goto NPDM_BUILD_END;
                 }
+                desc_idx += 10;
             }
             caps[cur_cap++] = (u32)((desc << 12) | (0x07FF));
         } else if (!strcmp(type_str, "application_type")) {
@@ -563,15 +829,16 @@ int CreateNpdm(const char *json, void **dst, u32 *dst_size) {
             }
             int allow_debug = 0;
             int force_debug = 0;
-            if (!cJSON_GetBoolean(value, "allow_debug", &allow_debug)) {
+            int force_debug_prod = 0;
+            cJSON_GetBoolean(value, "allow_debug", &allow_debug);
+            cJSON_GetBoolean(value, "force_debug", &force_debug);
+            cJSON_GetBoolean(value, "force_debug_prod", &force_debug_prod);
+            if ( allow_debug + force_debug + force_debug_prod > 1 ) {
+                fprintf(stderr, "Only one of allow_debug, force_debug, or force_debug_prod can be set!\n");
                 status = 0;
                 goto NPDM_BUILD_END;
             }
-            if (!cJSON_GetBoolean(value, "force_debug", &force_debug)) {
+            desc = (allow_debug & 1) | ((force_debug_prod & 1) << 1) | ((force_debug & 1) << 2);
-                status = 0;
-                goto NPDM_BUILD_END;
-            }
-            desc = (allow_debug & 1) | ((force_debug & 1) << 1);
             caps[cur_cap++] = (u32)((desc << 17) | (0xFFFF));
         }
     }

src/nxlink.c

@@ -14,11 +14,21 @@
 #include <arpa/inet.h>
 #include <netinet/in.h>
 #include <netdb.h>
+#include <poll.h>
+#define closesocket close
 #else
 #include <winsock2.h>
 #include <ws2tcpip.h>
 typedef int socklen_t;
 typedef uint32_t in_addr_t;
+#define SHUT_RD SD_RECEIVE
+#define SHUT_WR SD_SEND
+#define SHUT_RDWR SD_BOTH
+#ifdef EWOULDBLOCK
+#undef EWOULDBLOCK
+#endif
+#define EWOULDBLOCK WSAEWOULDBLOCK
+#define poll WSAPoll
 #endif
 
 #include <zlib.h>
@@ -30,100 +40,95 @@ typedef uint32_t in_addr_t;
 #define NETLOADER_CLIENT_PORT 28771
 
 
-char cmdbuf[3072];
+static char cmdbuf[3072];
-uint32_t cmdlen=0;
+static uint32_t cmdlen=0;
 
 //---------------------------------------------------------------------------------
-void shutdownSocket(int socket) {
+static void shutdownSocket(int socket, int flags) {
 //---------------------------------------------------------------------------------
-#ifdef __WIN32__
+	if (flags)
-	shutdown (socket, SD_SEND);
+		shutdown(socket, flags);
-	closesocket (socket);
+	closesocket(socket);
-#else
-	close(socket);
-#endif
 }
 
 //---------------------------------------------------------------------------------
-static int set_socket_nonblocking(int sock) {
+static int setSocketNonblocking(int sock) {
 //---------------------------------------------------------------------------------
 
 #ifndef __WIN32__
 	int flags = fcntl(sock, F_GETFL);
 
-	if(flags == -1) return -1;
+	if (flags == -1) return -1;
 
 	int rc = fcntl(sock, F_SETFL, flags | O_NONBLOCK);
 
-	if(rc != 0) return -1;
+	if (rc != 0) return -1;
-
 #else
-	u_long opt = 1;
+	u_long iMode = 1; // non-blocking
-	ioctlsocket(sock, FIONBIO, &opt);
+
+	int rc = ioctlsocket(sock, FIONBIO, &iMode);
+
+	if (rc != NO_ERROR) return -1;
 #endif
 
 	return 0;
 }
 
-void socket_error(const char *msg) {
+//---------------------------------------------------------------------------------
+static int socketError(const char *msg) {
+//---------------------------------------------------------------------------------
 #ifndef _WIN32
+	int ret = errno;
+	if (ret == EAGAIN)
+		ret = EWOULDBLOCK;
 	perror(msg);
 #else
+	int ret = WSAGetLastError();
 	wchar_t *s = NULL;
 	FormatMessageW(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
-               NULL, WSAGetLastError(),
+               NULL, ret,
                MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                (LPWSTR)&s, 0, NULL);
 	fprintf(stderr, "%S\n", s);
 	LocalFree(s);
+	if (ret == WSAEWOULDBLOCK)
+		ret = EWOULDBLOCK;
 #endif
-}
 
-/*---------------------------------------------------------------------------------
+	return ret;
-	Subtract the `struct timeval' values Y from X,
-	storing the result in RESULT.
-	Return 1 if the difference is negative, otherwise 0.
-
-	From http://www.gnu.org/software/libtool/manual/libc/Elapsed-Time.html
----------------------------------------------------------------------------------*/
-int timeval_subtract (struct timeval *result, struct timeval *x, struct timeval *y) {
-//---------------------------------------------------------------------------------
-	struct timeval tmp;
-	tmp.tv_sec = y->tv_sec;
-	tmp.tv_usec = y->tv_usec;
-
-	/* Perform the carry for the later subtraction by updating y. */
-	if (x->tv_usec < tmp.tv_usec) {
-		int nsec = (tmp.tv_usec - x->tv_usec) / 1000000 + 1;
-		tmp.tv_usec -= 1000000 * nsec;
-		tmp.tv_sec += nsec;
-	}
-
-	if (x->tv_usec - tmp.tv_usec > 1000000) {
-		int nsec = (x->tv_usec - tmp.tv_usec) / 1000000;
-		tmp.tv_usec += 1000000 * nsec;
-		tmp.tv_sec -= nsec;
-	}
-
-	/*	Compute the time remaining to wait.
-		tv_usec is certainly positive. */
-	result->tv_sec = x->tv_sec - tmp.tv_sec;
-	result->tv_usec = x->tv_usec - tmp.tv_usec;
-
-	/* Return 1 if result is negative. */
-	return x->tv_sec < tmp.tv_sec;
 }
 
 //---------------------------------------------------------------------------------
-void timeval_add (struct timeval *result, struct timeval *x, struct timeval *y) {
+int pollSocket(int fd, int events, int timeout) {
 //---------------------------------------------------------------------------------
-	result->tv_sec = x->tv_sec + y->tv_sec;
+#ifndef __WIN32__
-	result->tv_usec = x->tv_usec + y->tv_usec;
+	struct pollfd pfd;
+#else
+	WSAPOLLFD pfd;
+#endif
 
-	if ( result->tv_usec > 1000000) {
+	pfd.fd = fd;
-		result->tv_sec += result->tv_usec / 1000000;
+	pfd.events = events;
-		result->tv_usec = result->tv_usec % 1000000;
+	pfd.revents = 0;
+
+	int ret = poll(&pfd, 1, timeout);
+	if (ret < 0) {
+		socketError("poll");
+		return -1;
 	}
+
+	if (ret == 0)
+		return -1;
+
+	if (!(pfd.revents & events)) {
+		int err = 0;
+		int len = sizeof(err);
+		getsockopt(fd, SOL_SOCKET, SO_ERROR, (char*)&err, &len);
+		fprintf(stderr, "socket error 0x%x on poll\n", err);
+		return -1;
+	}
+
+	return 0;
 }
 
 //---------------------------------------------------------------------------------
@@ -137,7 +142,7 @@ static struct in_addr findSwitch(int retries) {
 	char mess[] = "nxboot";
 
 	int broadcastSock = socket(PF_INET, SOCK_DGRAM, 0);
-	if(broadcastSock < 0) socket_error("create send socket");
+	if (broadcastSock < 0) socketError("create send socket");
 
 	int optval = 1, len;
 	setsockopt(broadcastSock, SOL_SOCKET, SO_BROADCAST, (char *)&optval, sizeof(optval));
@@ -154,102 +159,89 @@ static struct in_addr findSwitch(int retries) {
 
 	int recvSock = socket(PF_INET, SOCK_DGRAM, 0);
 
-	if (recvSock < 0)  socket_error("create receive socket");
+	if (recvSock < 0)  socketError("create receive socket");
 
-	if(bind(recvSock, (struct sockaddr*) &rs, sizeof(rs)) < 0) socket_error("bind receive socket");
+	if (bind(recvSock, (struct sockaddr*) &rs, sizeof(rs)) < 0) socketError("bind receive socket");
-	set_socket_nonblocking(recvSock);
+	setSocketNonblocking(recvSock);
 
-	struct timeval wanted, now, result;
+	while (retries) {
+		if (sendto(broadcastSock, mess, strlen(mess), 0, (struct sockaddr *)&s, sizeof(s)) < 0)
+			socketError("sendto");
 
-	gettimeofday(&wanted, NULL);
+		if (pollSocket(recvSock, POLLIN, 150) == 0) {
-
+			socklen_t socklen = sizeof(remote);
-	int timeout = retries;
+			len = recvfrom(recvSock, recvbuf, sizeof(recvbuf), 0, (struct sockaddr *)&remote, &socklen);
-	while(timeout) {
+			if (len != -1) {
-		gettimeofday(&now, NULL);
+				if (strncmp("bootnx", recvbuf, strlen("bootnx")) == 0) {
-		if ( timeval_subtract(&result,&wanted,&now)) {
+					break;
-			if(sendto(broadcastSock, mess, strlen(mess), 0, (struct sockaddr *)&s, sizeof(s)) < 0) socket_error("sendto");
+				}
-			result.tv_sec=0;
-			result.tv_usec=150000;
-			timeval_add(&wanted,&now,&result);
-			timeout--;
-		}
-		socklen_t socklen = sizeof(remote);
-		len = recvfrom(recvSock,recvbuf,sizeof(recvbuf),0,(struct sockaddr *)&remote,&socklen);
-		if ( len != -1) {
-			if ( strncmp("bootnx",recvbuf,strlen("bootnx")) == 0) {
-				break;
 			}
 		}
+
+		--retries;
 	}
-	if (timeout == 0) remote.sin_addr.s_addr =  INADDR_NONE;
+
-	shutdownSocket(broadcastSock);
+	if (retries == 0)
-	shutdownSocket(recvSock);
+		remote.sin_addr.s_addr =  INADDR_NONE;
+
+	shutdownSocket(broadcastSock, 0);
+	shutdownSocket(recvSock, SHUT_RD);
+
 	return remote.sin_addr;
 }
 
 //---------------------------------------------------------------------------------
-int sendData(int sock, int sendsize, void *buffer) {
+static int sendData(int sock, int sendsize, void *buffer) {
 //---------------------------------------------------------------------------------
 	char *buf = (char*)buffer;
-	while(sendsize) {
+	while (sendsize) {
+		if (pollSocket(sock, POLLOUT, -1))
+			return 1;
+
 		int len = send(sock, buf, sendsize, 0);
-		if (len == 0) break;
+		if (len == 0)
-		if (len != -1) {
+			return 1;
+
+		if (len == -1) {
+			if (socketError("send") != EWOULDBLOCK)
+				return 1;
+		} else {
 			sendsize -= len;
 			buf += len;
-		} else {
-#ifdef _WIN32
-			int errcode = WSAGetLastError();
-			if (errcode != WSAEWOULDBLOCK) {
-				socket_error("sendData");
-				break;
-			}
-#else
-			if ( errno != EWOULDBLOCK && errno != EAGAIN) {
-				socket_error("sendData");
-				break;
-			}
-#endif
 		}
 	}
+
 	return sendsize != 0;
 }
 
 //---------------------------------------------------------------------------------
-int recvData(int sock, void *buffer, int size, int flags) {
+static int recvData(int sock, void *buffer, int size, int flags) {
 //---------------------------------------------------------------------------------
 	int len, sizeleft = size;
 	char *buf = (char*)buffer;
 	while (sizeleft) {
+		if (pollSocket(sock, POLLIN, -1))
+			return 0;
+
 		len = recv(sock,buf,sizeleft,flags);
-		if (len == 0) {
+		if (len == 0)
-			size = 0;
+			return 0;
-			break;
+
-		}
+		if (len == -1) {
-		if (len != -1) {
+			if (socketError("recv") != EWOULDBLOCK)
+				return 0;
+		} else {
 			sizeleft -=len;
 			buf +=len;
-		} else {
-#ifdef _WIN32
-			int errcode = WSAGetLastError();
-			if (errcode != WSAEWOULDBLOCK) {
-				printf("recvdata error %d\n",errcode);
-				break;
-			}
-#else
-			if ( errno != EWOULDBLOCK && errno != EAGAIN) {
-				socket_error("recvdata");
-				break;
-			}
-#endif
 		}
 	}
+
 	return size;
 }
 
 
 //---------------------------------------------------------------------------------
-int sendInt32LE(int socket, uint32_t size) {
+static int sendInt32LE(int socket, uint32_t size) {
 //---------------------------------------------------------------------------------
 	unsigned char lenbuf[4];
 	lenbuf[0] = size & 0xff;
@@ -257,14 +249,14 @@ int sendInt32LE(int socket, uint32_t size) {
 	lenbuf[2] = (size >> 16) & 0xff;
 	lenbuf[3] = (size >> 24) & 0xff;
 
-	return sendData(socket,4,lenbuf);
+	return sendData(socket, 4, lenbuf);
 }
 
 //---------------------------------------------------------------------------------
-int recvInt32LE(int socket, int32_t *data) {
+static int recvInt32LE(int socket, int32_t *data) {
 //---------------------------------------------------------------------------------
 	unsigned char intbuf[4];
-	int len = recvData(socket,intbuf,4,0);
+	int len = recvData(socket, intbuf, 4, 0);
 
 	if (len == 4) {
 		*data = intbuf[0] & 0xff + (intbuf[1] <<  8) + (intbuf[2] <<  16) + (intbuf[3] <<  24);
@@ -272,14 +264,13 @@ int recvInt32LE(int socket, int32_t *data) {
 	}
 
 	return -1;
-
 }
 
-unsigned char in[ZLIB_CHUNK];
+static unsigned char in[ZLIB_CHUNK];
-unsigned char out[ZLIB_CHUNK];
+static unsigned char out[ZLIB_CHUNK];
 
 //---------------------------------------------------------------------------------
-int sendNROFile(in_addr_t nxaddr, char *name, size_t filesize, FILE *fh) {
+static int sendNROFile(in_addr_t nxaddr, char *name, size_t filesize, FILE *fh) {
 //---------------------------------------------------------------------------------
 
 	int retval = 0;
@@ -297,7 +288,7 @@ int sendNROFile(in_addr_t nxaddr, char *name, size_t filesize, FILE *fh) {
 
 	int sock = socket(AF_INET,SOCK_STREAM,0);
 	if (sock < 0) {
-		socket_error("create connection socket");
+		socketError("create connection socket");
 		return -1;
 	}
 
@@ -307,7 +298,7 @@ int sendNROFile(in_addr_t nxaddr, char *name, size_t filesize, FILE *fh) {
 	s.sin_port = htons(NETLOADER_SERVER_PORT);
 	s.sin_addr.s_addr = nxaddr;
 
-	if (connect(sock,(struct sockaddr *)&s,sizeof(s)) < 0 ) {
+	if (connect(sock,(struct sockaddr *)&s,sizeof(s)) < 0) {
 		struct in_addr address;
 		address.s_addr = nxaddr;
 		fprintf(stderr,"Connection to %s failed\n",inet_ntoa(address));
@@ -336,13 +327,13 @@ int sendNROFile(in_addr_t nxaddr, char *name, size_t filesize, FILE *fh) {
 
 	int response;
 
-	if(recvInt32LE(sock,&response)!=0) {
+	if (recvInt32LE(sock,&response)!=0) {
 		fprintf(stderr,"Invalid response\n");
 		retval = 1;
 		goto error;
 	}
 
-	if(response!=0) {
+	if (response!=0) {
 		switch(response) {
 			case -1:
 				fprintf(stderr,"Failed to create file\n");
@@ -387,7 +378,7 @@ int sendNROFile(in_addr_t nxaddr, char *name, size_t filesize, FILE *fh) {
 					goto error;
 				}
 
-				if(sendData(sock,have,out)) {
+				if (sendData(sock,have,out)) {
 					fprintf(stderr,"Failed sending %s\n", name);
 					retval = 1;
 					(void)deflateEnd(&strm);
@@ -406,14 +397,14 @@ int sendNROFile(in_addr_t nxaddr, char *name, size_t filesize, FILE *fh) {
 
 	printf("%zu sent (%.2f%%), %zd blocks\n",totalsent, (float)(totalsent * 100.0)/ filesize, blocks);
 
-	if(recvInt32LE(sock,&response)!=0) {
+	if (recvInt32LE(sock,&response)!=0) {
 		fprintf(stderr,"Failed sending %s\n",name);
 		retval = 1;
 		goto error;
 	}
 
 
-	if(sendData(sock,cmdlen+4,(unsigned char*)cmdbuf)) {
+	if (sendData(sock,cmdlen+4,(unsigned char*)cmdbuf)) {
 
 		fprintf(stderr,"Failed sending command line\n");
 		retval = 1;
@@ -421,12 +412,12 @@ int sendNROFile(in_addr_t nxaddr, char *name, size_t filesize, FILE *fh) {
 	}
 
 error:
-	shutdownSocket(sock);
+	shutdownSocket(sock, SHUT_WR);
 	return retval;
 }
 
 //---------------------------------------------------------------------------------
-void showHelp() {
+static void showHelp() {
 //---------------------------------------------------------------------------------
 	puts("Usage: nxlink [options] nrofile\n");
 	puts("--help,    -h   Display this information");
@@ -440,7 +431,7 @@ void showHelp() {
 
 
 //---------------------------------------------------------------------------------
-int add_extra_args(int len, char *buf, char *extra_args) {
+static int addExtraArgs(int len, char *buf, char *extra_args) {
 //---------------------------------------------------------------------------------
 
 	if (NULL==extra_args) return len;
@@ -457,7 +448,7 @@ int add_extra_args(int len, char *buf, char *extra_args) {
 		do {
 			c = *src++;
 			extra_len--;
-		} while(c ==' ' && extra_len >= 0);
+		} while (c ==' ' && extra_len >= 0);
 
 		if (c == '\"' || c == '\'') {
 			int quote = c;
@@ -465,7 +456,7 @@ int add_extra_args(int len, char *buf, char *extra_args) {
 				c = *src++;
 				if (c != quote) *dst++ = c;
 				extra_len--;
-			} while(c != quote && extra_len >= 0);
+			} while (c != quote && extra_len >= 0);
 
 			*dst++ = '\0';
 
@@ -475,16 +466,22 @@ int add_extra_args(int len, char *buf, char *extra_args) {
 			*dst++ = c;
 			extra_len--;
 			c = *src++;
-		} while(c != ' ' && extra_len >= 0);
+		} while (c != ' ' && extra_len >= 0);
 
 		*dst++ = '\0';
-	} while(extra_len >= 0);
+	} while (extra_len >= 0);
 
 	return dst - buf;
 }
 
 #define NRO_ARGS	1000
 
+#ifdef __WIN32__
+static void win32_socket_cleanup(void) {
+	WSACleanup();
+}
+#endif
+
 //---------------------------------------------------------------------------------
 int main(int argc, char **argv) {
 //---------------------------------------------------------------------------------
@@ -498,10 +495,10 @@ int main(int argc, char **argv) {
 
 	if (argc < 2) {
 		showHelp();
-		return 1;
+		return EXIT_FAILURE;
 	}
 
-	while(1) {
+	while (1) {
 		static struct option long_options[] = {
 			{"address", required_argument, 0,	'a'},
 			{"retries", required_argument, 0,	'r'},
@@ -519,7 +516,7 @@ int main(int argc, char **argv) {
 
 		/* Detect the end of the options. */
 		if (c == -1)
-		break;
+			break;
 
 		switch(c) {
 
@@ -543,7 +540,7 @@ int main(int argc, char **argv) {
 			break;
 		case 'h':
 			showHelp();
-			break;
+			return EXIT_FAILURE;
 		case NRO_ARGS:
 			extra_args=optarg;
 			break;
@@ -554,7 +551,7 @@ int main(int argc, char **argv) {
 	char *filename = argv[optind++];
 	if (filename== NULL) {
 		showHelp();
-		return 1;
+		return EXIT_FAILURE;
 	}
 
 	memset(cmdbuf, '\0', sizeof(cmdbuf));
@@ -562,7 +559,7 @@ int main(int argc, char **argv) {
 	FILE *fh = fopen(filename,"rb");
 	if (fh == NULL) {
 		fprintf(stderr,"Failed to open %s\n",filename);
-		return -1;
+		return EXIT_FAILURE;
 	}
 
 #ifdef _WIN32
@@ -574,7 +571,7 @@ int main(int argc, char **argv) {
 	fseek(fh,0,SEEK_SET);
 
 	char *basename = NULL;
-	if((basename=strrchr(filename,'/'))!=NULL) {
+	if ((basename=strrchr(filename,'/'))!=NULL) {
 		basename++;
 	} else if ((basename=strrchr(filename,'\\'))!=NULL) {
 		basename++;
@@ -582,7 +579,7 @@ int main(int argc, char **argv) {
 		basename = filename;
 	}
 
-	if(basepath) {
+	if (basepath) {
 		size_t finalpath_len = strlen(basepath);
 		if (basepath[finalpath_len] == '/') {
 			finalpath_len += (strlen(basename) + 1);
@@ -599,12 +596,12 @@ int main(int argc, char **argv) {
 
 	for (int index = optind; index < argc; index++) {
 		int len=strlen(argv[index]);
-		if ( (cmdlen + len + 5 ) >= (sizeof(cmdbuf) - 2) ) break;
+		if ((cmdlen + len + 5) >= (sizeof(cmdbuf) - 2)) break;
 		strcpy(&cmdbuf[cmdlen+4],argv[index]);
 		cmdlen+= len + 1;
 	}
 
-	cmdlen = add_extra_args(cmdlen, &cmdbuf[4], extra_args);
+	cmdlen = addExtraArgs(cmdlen, &cmdbuf[4], extra_args);
 
 	cmdbuf[0] = cmdlen & 0xff;
 	cmdbuf[1] = (cmdlen>>8) & 0xff;
@@ -615,8 +612,9 @@ int main(int argc, char **argv) {
 	WSADATA wsa_data;
 	if (WSAStartup (MAKEWORD(2,2), &wsa_data)) {
 		printf ("WSAStartup failed\n");
-		return 1;
+		return EXIT_FAILURE;
 	}
+	atexit(&win32_socket_cleanup);
 #endif
 
 	struct in_addr nxaddr;
@@ -627,7 +625,7 @@ int main(int argc, char **argv) {
 
 		if (nxaddr.s_addr == INADDR_NONE) {
 			printf("No response from Switch!\n");
-			return 1;
+			return EXIT_FAILURE;
 		}
 
 	} else {
@@ -640,115 +638,97 @@ int main(int argc, char **argv) {
 
 	if (nxaddr.s_addr == INADDR_NONE) {
 		fprintf(stderr,"Invalid address\n");
-		return 1;
+		return EXIT_FAILURE;
 	}
 
 	int res = sendNROFile(nxaddr.s_addr,finalpath,filesize,fh);
 
 	fclose(fh);
 
-	if ( res == 0 && server) {
+	if (res != 0)
-		printf("starting server\n");
+		return EXIT_FAILURE;
 
-		struct sockaddr_in serv_addr;
+	if (!server)
+		return EXIT_SUCCESS;
 
-		memset(&serv_addr, '0', sizeof(serv_addr));
+	printf("starting server\n");
-		serv_addr.sin_family = AF_INET;
-		serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
-		serv_addr.sin_port = htons(NETLOADER_CLIENT_PORT);
 
-		int listenfd = socket(AF_INET, SOCK_STREAM, 0);
+	struct sockaddr_in serv_addr;
-		int rc;
 
-		if(listenfd < 0) {
+	memset(&serv_addr, '0', sizeof(serv_addr));
-			socket_error("socket");
+	serv_addr.sin_family = AF_INET;
-		} else {
+	serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
+	serv_addr.sin_port = htons(NETLOADER_CLIENT_PORT);
 
-			rc = bind(listenfd, (struct sockaddr*)&serv_addr, sizeof(serv_addr));
+	int listenfd = socket(AF_INET, SOCK_STREAM, 0);
+	if (listenfd < 0) {
+		socketError("socket");
+		return EXIT_FAILURE;
+	}
 
-			if(rc != 0) {
+	int rc = bind(listenfd, (struct sockaddr*)&serv_addr, sizeof(serv_addr));
-				socket_error("bind listen socket");
+	if (rc != 0) {
-			} else {
+		socketError("bind listen socket");
-				if (set_socket_nonblocking(listenfd) == -1) {
+		shutdownSocket(listenfd, 0);
-					socket_error("listen fcntl");
+		return EXIT_FAILURE;
+	}
 
-				} else {
+	rc = setSocketNonblocking(listenfd);
-					rc = listen(listenfd, 10);
+	if (rc == -1) {
+		socketError("listen fcntl");
+		shutdownSocket(listenfd, 0);
+		return EXIT_FAILURE;
+	}
 
-					if(rc != 0) {
+	rc = listen(listenfd, 10);
-						socket_error("listen");
+	if (rc != 0) {
-					} else {
+		socketError("listen");
-						printf("server active ...\n");
+		shutdownSocket(listenfd, 0);
+		return EXIT_FAILURE;
+	}
 
-						int datafd = -1;
+	printf("server active ...\n");
 
-						while( listenfd != -1 || datafd != -1) {
+	int datafd = -1;
-							struct sockaddr_in sa_remote;
-							if(listenfd >= 0 && datafd < 0) {
-								socklen_t addrlen = sizeof(sa_remote);
-								datafd = accept(listenfd, (struct sockaddr*)&sa_remote, &addrlen);
 
-								if(datafd < 0) {
+	while (listenfd != -1 || datafd != -1) {
-#ifdef _WIN32
+		struct sockaddr_in sa_remote;
-									int errcode = WSAGetLastError();
-									if (errcode != WSAEWOULDBLOCK) {
-										socket_error("accept");
-										listenfd = -1;
-									}
-#else
-									if ( errno != EWOULDBLOCK && errno != EAGAIN) {
-										socket_error("accept");
-										listenfd = -1;
-									}
-#endif
 
-								} else {
+		if (pollSocket(listenfd >= 0 ? listenfd : datafd, POLLIN, -1))
-									shutdownSocket(listenfd);
+			break;
-									listenfd = -1;
-								}
-							}
 
+		if (listenfd >= 0) {
+			socklen_t addrlen = sizeof(sa_remote);
+			datafd = accept(listenfd, (struct sockaddr*)&sa_remote, &addrlen);
 
-							if(datafd >= 0) {
+			if (datafd < 0 && socketError("accept") != EWOULDBLOCK)
-								char recvbuf[256];
+				break;
-								int len = recv(datafd,recvbuf,256,0);
 
-								if (len > 0 ) {
+			if (datafd >= 0) {
-									recvbuf[len] = 0;
+				shutdownSocket(listenfd, 0);
-									printf("%s", recvbuf);
+				listenfd = -1;
-								} else {
-
-									if (len == -1) {
-#ifdef _WIN32
-										int errcode = WSAGetLastError();
-										if (errcode != WSAEWOULDBLOCK) {
-											socket_error("recvdata");
-											len = 0;
-										}
-#else
-										if ( errno != EWOULDBLOCK && errno != EAGAIN) {
-											perror("recvdata");
-											len = 0;
-										}
-#endif
-									}
-									if (len ==0 ) {
-										shutdownSocket(datafd);
-										datafd = -1;
-									}
-								}
-							}
-						}
-						printf("exiting ... \n");
-					}
-				}
 			}
+		} else {
+			char recvbuf[256];
+			int len = recv(datafd, recvbuf, sizeof(recvbuf), 0);
+
+			if (len == 0 || (len < 0 && socketError("recv") != EWOULDBLOCK)) {
+				shutdownSocket(datafd, 0);
+				datafd = -1;
+				break;
+			}
+
+			if (len > 0)
+				fwrite(recvbuf, 1, len, stdout);
 		}
 	}
 
-#ifdef __WIN32__
+	if (listenfd >= 0)
-	WSACleanup ();
+		shutdownSocket(listenfd, 0);
-#endif
+	if (datafd >= 0)
-	return res;
+		shutdownSocket(datafd, SHUT_RD);
-}
 
+	printf("exiting ... \n");
+
+	return EXIT_SUCCESS;
+}

src/romfs.c

@@ -410,9 +410,7 @@ size_t build_romfs_into_file(filepath_t *in_dirpath, FILE *f_out, off_t base_off
         cur_entry->name_size = name_size;
         memcpy(cur_entry->name, cur_dir->cur_path.char_path + 1, name_size);
 
-        romfs_dirent_ctx_t *temp = cur_dir;
         cur_dir = cur_dir->next;
-        free(temp);
     }
 
     header.header_size = le_dword(sizeof(header));
@@ -474,11 +472,27 @@ size_t build_romfs_into_file(filepath_t *in_dirpath, FILE *f_out, off_t base_off
 
         os_fclose(f_in);
 
+        cur_file = cur_file->next;
+    }
+    free(buffer);
+
+    /* Free all files. */
+    cur_file = romfs_ctx.files;
+    while (cur_file != NULL) {
         romfs_fent_ctx_t *temp = cur_file;
         cur_file = cur_file->next;
         free(temp);
     }
-    free(buffer);
+    romfs_ctx.files = NULL;
+
+    /* Free all directories. */
+    cur_dir = root_ctx;
+    while (cur_dir != NULL) {
+        romfs_dirent_ctx_t *temp = cur_dir;
+        cur_dir = cur_dir->next;
+        free(temp);
+    }
+    root_ctx = NULL;
 
     fseeko64(f_out, base_offset + dir_hash_table_ofs, SEEK_SET);
     if (fwrite(dir_hash_table, 1, romfs_ctx.dir_hash_table_size, f_out) != romfs_ctx.dir_hash_table_size) {