ubnt-bs: add new utility for Ubiquiti boot select

Signed-off-by: Bjørn Mork <bjorn@mork.no>

2 files changed, 239 insertions, 0 deletions

diff --git a/ubnt-bs/Makefile b/ubnt-bs/Makefile
new file mode 100644
index 0000000..c471e6e
--- /dev/null
+++ b/ubnt-bs/Makefile
@@ -0,0 +1,45 @@
+# 
+# Copyright (C) 2017 Bjørn Mork <bjorn@mork.no>
+#
+# This is free software, licensed under the GNU General Public License v2.
+# See /LICENSE for more information.
+#
+
+include $(TOPDIR)/rules.mk
+
+PKG_NAME:=ubnt-bs
+PKG_VERSION:=0.01
+PKG_RELEASE:=1
+PKG_LICENSE:=GPL-2.0
+PKG_MAINTAINER:=Bjørn Mork <bjorn@mork.no>
+
+include $(INCLUDE_DIR)/package.mk
+
+define Package/ubnt-bs
+  SECTION:=utils
+  CATEGORY:=Utilities
+  TITLE:="Boot select utility for Ubiquiti UniFi AC etc"
+endef
+
+define Package/covery/description
+  Boot select utility for Ubiquiti UniFi AC etc
+endef
+
+define Build/Prepare
+	$(INSTALL_DIR) $(PKG_BUILD_DIR)
+	$(CP) ./src/* $(PKG_BUILD_DIR)/
+endef
+
+define Build/Configure
+endef
+
+define Build/Compile
+	$(TARGET_CC) $(TARGET_CFLAGS) -o $(PKG_BUILD_DIR)/ubnt-bs $(PKG_BUILD_DIR)/ubnt-bs.c
+endef
+
+define Package/covery/install
+	$(INSTALL_DIR) $(1)/usr/sbin
+	$(INSTALL_BIN) $(PKG_BUILD_DIR)/ubnt-bs $(1)/usr/sbin/
+endef
+
+$(eval $(call BuildPackage,ubnt-bs))
diff --git a/ubnt-bs/src/ubnt-bs.c b/ubnt-bs/src/ubnt-bs.c
new file mode 100644
index 0000000..6649045
--- /dev/null
+++ b/ubnt-bs/src/ubnt-bs.c
@@ -0,0 +1,194 @@
+/*
+ * ubnt-bs - LEDE "boot select" utility for Ubiquiti UniFi AC etc
+ *
+ * Copyright (C) 2017 Bjørn Mork <bjorn@mork.no>
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License v2
+ * as published by the Free Software Foundation.
+ */
+
+#include <byteswap.h>
+#include <endian.h>
+#include <unistd.h>
+#include <stdlib.h>
+#include <stdio.h>
+#include <string.h>
+#include <getopt.h>
+#include <sysexits.h>
+#include <fcntl.h>
+#include <sys/ioctl.h>
+#include <sys/param.h>
+#include <mtd/mtd-user.h>
+
+#define VERSION "0.01"
+
+#ifdef TEST
+static const char bsname[] = "\"mtdram test device\"";
+#else
+static const char bsname[] = "\"bs\"";
+#endif
+
+static const struct option long_options[] = {
+  { "show",      0, NULL, 'i' },
+  { "primary",   0, NULL, '0' },
+  { "secondary", 0, NULL, '1' },
+  { "verbose",   0, NULL, 'v' },
+  { "help",      0, NULL, 'h' },
+  { "version",   0, NULL, 'V' },
+  { 0, 0, 0, 0 }
+};
+
+const char *program;
+
+static void usage(int ex)
+{
+	fprintf(stderr,
+		"Usage: %s [options]\n"
+		"       Updates boot select partition\n"
+		"Options:\n"
+		"  -i  --show         Show current status\n"
+		"  -0  --primary      Boot from primary (kernel0/firmware) partition\n"
+		"  -1  --secondary    Boot from secondary (kernel1/ubnt-airos) partition\n"
+		"  -v  --verbose      Verbose output\n"
+		"  -h  --help         Display this text\n"
+		"  -V  --version      Display version\n",
+		program);
+	exit(ex);
+}
+
+/* we define these as big endian */
+#define MAGIC  0xa34de82b
+#define BSMASK 0x80000000
+
+#if __BYTE_ORDER == __BIG_ENDIAN
+#define cpu_to_be32(x)  (x)
+#define be32_to_cpu(x)  (x)
+#elif __BYTE_ORDER == __LITTLE_ENDIAN
+#define cpu_to_be32(x)  bswap_32(x)
+#define be32_to_cpu(x)  bswap_32(x)
+#else
+#error "Unsupported endianness"
+#endif
+
+struct bs_rec {
+	__be32 flags;
+	__be32 magic;
+};
+
+/* use globals to reduce stack abuse */
+struct mtd_info_user info;
+struct erase_info_user li;
+
+static int mtd_open(const char *mtd, struct mtd_info_user *info)
+{
+        FILE *fp = fopen("/proc/mtd", "r");
+        int fd, i = -1;
+	char buf[PATH_MAX];
+
+	if (!fp)
+		return -1;
+	while (fgets(buf, sizeof(buf), fp)) {
+		if (sscanf(buf, "mtd%d:", &i) && strstr(buf, mtd))
+			break;
+		i = -1;
+	}
+	fclose(fp);
+	if (i < 0)
+		return -1;
+
+	snprintf(buf, sizeof(buf),"/dev/mtd%d", i);
+        fd = open(buf, O_RDWR | O_SYNC);
+        if (fd < 0) {
+                fprintf(stderr, "Could not open mtd device: %s\n", mtd);
+                return fd;
+        }
+
+        if (ioctl(fd, MEMGETINFO, info)) {
+                fprintf(stderr, "Could not get mtd info from %s\n", mtd);
+                close(fd);
+                return -1;
+        }
+        return fd;
+}
+
+static int erase_mtd(int fd, struct mtd_info_user *info)
+{
+	struct erase_info_user erase;
+
+        erase.start = 0;
+        erase.length = info->size;
+        ioctl(fd, MEMUNLOCK, &erase);
+        if (ioctl (fd, MEMERASE, &erase) < 0)
+                return -1;
+        return 0;
+}
+
+int main(int argc, char *argv[])
+{
+	int opt, fd, ret = -1, cmd = -1, verbose = 0;
+	struct bs_rec *rec;
+	char *p = NULL;
+
+	program = argv[0];
+	if (argc < 2 )
+		usage(EX_USAGE);
+	optopt = 0;
+	while ((opt = getopt_long(argc, argv, "+icdxavhV?", long_options, NULL)) != EOF ) {
+		switch(opt) {
+		case 'i':
+		case '0':
+		case '1':
+			if (cmd > 0)
+				usage(-1);
+			cmd = opt - '0';
+			break;
+		case 'v':
+			verbose = 1;
+			break;
+		case 'V':
+			fprintf(stderr, "%s (%s)\n", program, VERSION);
+			exit(0);
+		default:
+			/* optopt will be set if this was an unrecognized option, i.e. *not* 'h' or '?' */
+			usage(optopt ? EX_USAGE : 0);
+			break;
+		}
+	}
+	if (cmd < 0)
+		usage(-1);
+
+	fd = mtd_open(bsname, &info);
+	if (fd < 0) {
+		fprintf(stderr, "ERROR: no %s found in /proc/mtd\n", bsname);
+		goto err;
+	}
+	if (verbose)
+		fprintf(stderr, "opened %s with size %u\n", bsname, info.size);
+	p = malloc(info.erasesize);
+	if (!p) {
+		fprintf(stderr, "ERROR: failed to allocate %u byte buffer\n", info.erasesize);
+		goto err;
+	}
+
+	if (read(fd, p, info.erasesize) != info.erasesize) {
+		fprintf(stderr, "ERROR: failed to read %u bytes from %s\n", info.erasesize, bsname);
+		goto err;
+	}
+
+	rec = (struct bs_rec *)p;
+	if (cpu_to_be32(rec->magic) != MAGIC) {
+		fprintf(stderr, "ERROR: magic mismatch: 0x%04x != 0x%04x\n", cpu_to_be32(rec->magic) != MAGIC);
+		goto err;
+	}
+
+	fprintf(stderr, "INFO: current boot selection is \"%s\"\n", cpu_to_be32(rec->flags) & BSMASK ? "secondary" : "primary");
+	ret = 0;
+err:
+	if (fd < 0)
+		return fd;
+	close(fd);
+	if (p)
+		free(p);
+	return ret;
+}