6 files changed, 52 insertions, 7 deletions

diff --git a/payloads/libpayload/drivers/usb/uhci.c b/payloads/libpayload/drivers/usb/uhci.c
index 198b6c0846..def6b45084 100644
--- a/payloads/libpayload/drivers/usb/uhci.c
+++ b/payloads/libpayload/drivers/usb/uhci.c
@@ -130,7 +130,13 @@ uhci_init (pcidev_t addr)
 	int i;
 	hci_t *controller = new_controller ();
 
+	if (!controller)
+		usb_fatal("Could not create USB controller instance.\n");
+
 	controller->instance = malloc (sizeof (uhci_t));
+	if(!controller->instance)
+		usb_fatal("Not enough memory creating USB controller instance.\n");
+
 	controller->start = uhci_start;
 	controller->stop = uhci_stop;
 	controller->reset = uhci_reset;
@@ -157,6 +163,9 @@ uhci_init (pcidev_t addr)
 	pci_write_config32 (controller->bus_address, 0xc0, 0x8f00);
 
 	UHCI_INST (controller)->framelistptr = memalign (0x1000, 1024 * sizeof (flistp_t *));	/* 4kb aligned to 4kb */
+	if (! UHCI_INST (controller)->framelistptr)
+		usb_fatal("Not enough memory for USB frame list pointer.\n");
+
 	memset (UHCI_INST (controller)->framelistptr, 0,
 		1024 * sizeof (flistp_t));
 
@@ -168,6 +177,8 @@ uhci_init (pcidev_t addr)
 	          for some reason. Not a problem now.
 	   */
 	td_t *antiberserk = memalign(16, sizeof(td_t));
+	if (!antiberserk)
+		usb_fatal("Not enough memory for chipset workaround.\n");
 	memset(antiberserk, 0, sizeof(td_t));
 
 	UHCI_INST (controller)->qh_prei = memalign (16, sizeof (qh_t));
@@ -175,6 +186,12 @@ uhci_init (pcidev_t addr)
 	UHCI_INST (controller)->qh_data = memalign (16, sizeof (qh_t));
 	UHCI_INST (controller)->qh_last = memalign (16, sizeof (qh_t));
 
+	if (! UHCI_INST (controller)->qh_prei ||
+	    ! UHCI_INST (controller)->qh_intr ||
+	    ! UHCI_INST (controller)->qh_data ||
+	    ! UHCI_INST (controller)->qh_last)
+		usb_fatal ("Not enough memory for USB controller queues.\n");
+
 	UHCI_INST (controller)->qh_prei->headlinkptr.ptr =
 		virt_to_phys (UHCI_INST (controller)->qh_intr);
 	UHCI_INST (controller)->qh_prei->headlinkptr.queue_head = 1;
@@ -508,11 +525,16 @@ uhci_create_intr_queue (endpoint_t *ep, int reqsize, int reqcount, int reqtiming
 	td_t *tds = memalign(16, sizeof(td_t) * reqcount);
 	qh_t *qh = memalign(16, sizeof(qh_t));
 
+	if (!data || !tds || !qh)
+		usb_fatal ("Not enough memory to create USB intr queue prerequisites.\n");
+
 	qh->elementlinkptr.ptr = virt_to_phys(tds);
 	qh->elementlinkptr.queue_head = 0;
 	qh->elementlinkptr.terminate = 0;
 
 	intr_q *q = malloc(sizeof(intr_q));
+	if (!q)
+		usb_fatal ("Not enough memory to create USB intr queue.\n");
 	q->qh = qh;
 	q->tds = tds;
 	q->data = data;
diff --git a/payloads/libpayload/drivers/usb/uhci_rh.c b/payloads/libpayload/drivers/usb/uhci_rh.c
index cc3c600289..2f4c7d839c 100644
--- a/payloads/libpayload/drivers/usb/uhci_rh.c
+++ b/payloads/libpayload/drivers/usb/uhci_rh.c
@@ -157,6 +157,9 @@ uhci_rh_init (usbdev_t *dev)
 	uhci_rh_enable_port (dev, 1);
 	uhci_rh_enable_port (dev, 2);
 	dev->data = malloc (sizeof (rh_inst_t));
+	if (!dev->data)
+		usb_fatal ("Not enough memory for UHCI RH.\n");
+
 	RH_INST (dev)->port[0] = -1;
 	RH_INST (dev)->port[1] = -1;
 
diff --git a/payloads/libpayload/drivers/usb/usb.c b/payloads/libpayload/drivers/usb/usb.c
index d536d31d86..25e80065d7 100644
--- a/payloads/libpayload/drivers/usb/usb.c
+++ b/payloads/libpayload/drivers/usb/usb.c
@@ -33,14 +33,16 @@
 hci_t *usb_hcs = 0;
 
 hci_t *
-new_controller ()
+new_controller (void)
 {
 	hci_t *controller = malloc (sizeof (hci_t));
 
-	/* atomic */
-	controller->next = usb_hcs;
-	usb_hcs = controller;
-	/* atomic end */
+	if (controller) {
+		/* atomic */
+		controller->next = usb_hcs;
+		usb_hcs = controller;
+		/* atomic end */
+	}
 
 	return controller;
 }
@@ -48,13 +50,13 @@ new_controller ()
 void
 detach_controller (hci_t *controller)
 {
-	if (controller == 0)
+	if (controller == NULL)
 		return;
 	if (usb_hcs == controller) {
 		usb_hcs = controller->next;
 	} else {
 		hci_t *it = usb_hcs;
-		while (it != 0) {
+		while (it != NULL) {
 			if (it->next == controller) {
 				it->next = controller->next;
 				return;
@@ -386,3 +388,10 @@ usb_attach_device(hci_t *controller, int hubaddress, int port, int lowspeed)
 	newdev_t->init (newdev_t);
 	return newdev;
 }
+
+void
+usb_fatal (const char *message)
+{
+	printf(message);
+	for (;;) ;
+}
diff --git a/payloads/libpayload/drivers/usb/usbhid.c b/payloads/libpayload/drivers/usb/usbhid.c
index 44be9ce8d9..ee217c36d8 100644
--- a/payloads/libpayload/drivers/usb/usbhid.c
+++ b/payloads/libpayload/drivers/usb/usbhid.c
@@ -169,6 +169,8 @@ usb_hid_init (usbdev_t *dev)
 			boot_protos[interface->bInterfaceProtocol]);
 		if (interface->bInterfaceProtocol == hid_boot_proto_keyboard) {
 			dev->data = malloc (sizeof (usbhid_inst_t));
+			if (!dev->data)
+				usb_fatal("Not enough memory for USB HID device.\n");
 			printf ("  configuring...\n");
 			usb_hid_set_protocol(dev, interface, hid_proto_boot);
 			usb_hid_set_idle(dev, interface, 0);
diff --git a/payloads/libpayload/drivers/usb/usbhub.c b/payloads/libpayload/drivers/usb/usbhub.c
index 4625246e22..0a5d0ebcb5 100644
--- a/payloads/libpayload/drivers/usb/usbhub.c
+++ b/payloads/libpayload/drivers/usb/usbhub.c
@@ -128,6 +128,9 @@ usb_hub_init (usbdev_t *dev)
 
 	dev->data = malloc (sizeof (usbhub_inst_t));
 
+	if (!dev->data)
+		usb_fatal("Not enough memory for USB hub.\n");
+
 	HUB_INST (dev)->descriptor =
 		(hub_descriptor_t *) get_descriptor (dev,
 						     gen_bmRequestType
@@ -137,6 +140,9 @@ usb_hub_init (usbdev_t *dev)
 	HUB_INST (dev)->num_ports = HUB_INST (dev)->descriptor->bNbrPorts;
 	HUB_INST (dev)->ports =
 		malloc (sizeof (int) * (HUB_INST (dev)->num_ports + 1));
+	if (! HUB_INST (dev)->ports)
+		usb_fatal("Not enough memory for USB hub ports.\n");
+
 	for (i = 1; i <= HUB_INST (dev)->num_ports; i++)
 		HUB_INST (dev)->ports[i] = -1;
 	for (i = 1; i <= HUB_INST (dev)->num_ports; i++)
diff --git a/payloads/libpayload/drivers/usb/usbmsc.c b/payloads/libpayload/drivers/usb/usbmsc.c
index ad4a10c146..f24bd6d857 100644
--- a/payloads/libpayload/drivers/usb/usbmsc.c
+++ b/payloads/libpayload/drivers/usb/usbmsc.c
@@ -346,6 +346,9 @@ usb_msc_init (usbdev_t *dev)
 	}
 
 	dev->data = malloc (sizeof (usbmsc_inst_t));
+	if (!dev->data)
+		usb_fatal("Not enough memory for USB MSC device.\n");
+
 	MSC_INST (dev)->bulk_in = 0;
 	MSC_INST (dev)->bulk_out = 0;