diff --git a/drivers/staging/greybus/core.c b/drivers/staging/greybus/core.c
index b9303c0cd5e4..2fb95744e01c 100644
--- a/drivers/staging/greybus/core.c
+++ b/drivers/staging/greybus/core.c
@@ -10,6 +10,7 @@
 #define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
 
 #define CREATE_TRACE_POINTS
+#include "firmware.h"
 #include "greybus.h"
 #include "greybus_trace.h"
 #include "legacy.h"
@@ -243,9 +244,9 @@ static int __init gb_init(void)
 		goto error_operation;
 	}
 
-	retval = gb_firmware_protocol_init();
+	retval = gb_firmware_init();
 	if (retval) {
-		pr_err("gb_firmware_protocol_init failed\n");
+		pr_err("gb_firmware_init failed\n");
 		goto error_firmware;
 	}
 
@@ -258,7 +259,7 @@ static int __init gb_init(void)
 	return 0;	/* Success */
 
 error_legacy:
-	gb_firmware_protocol_exit();
+	gb_firmware_exit();
 error_firmware:
 	gb_operation_exit();
 error_operation:
@@ -275,7 +276,7 @@ module_init(gb_init);
 static void __exit gb_exit(void)
 {
 	gb_legacy_exit();
-	gb_firmware_protocol_exit();
+	gb_firmware_exit();
 	gb_operation_exit();
 	gb_hd_exit();
 	bus_unregister(&greybus_bus_type);
diff --git a/drivers/staging/greybus/firmware.c b/drivers/staging/greybus/firmware.c
index 37895107e427..17ce573b9162 100644
--- a/drivers/staging/greybus/firmware.c
+++ b/drivers/staging/greybus/firmware.c
@@ -9,11 +9,15 @@
 
 #include <linux/firmware.h>
 
+#include "firmware.h"
 #include "greybus.h"
 
+
 struct gb_firmware {
 	struct gb_connection	*connection;
 	const struct firmware	*fw;
+	u8			protocol_major;
+	u8			protocol_minor;
 };
 
 static void free_firmware(struct gb_firmware *firmware)
@@ -223,8 +227,10 @@ static int gb_firmware_ready_to_boot(struct gb_operation *op)
 	return 0;
 }
 
-static int gb_firmware_request_recv(u8 type, struct gb_operation *op)
+static int gb_firmware_request_handler(struct gb_operation *op)
 {
+	u8 type = op->type;
+
 	switch (type) {
 	case GB_FIRMWARE_TYPE_FIRMWARE_SIZE:
 		return gb_firmware_size_request(op);
@@ -239,60 +245,147 @@ static int gb_firmware_request_recv(u8 type, struct gb_operation *op)
 	}
 }
 
-static int gb_firmware_connection_init(struct gb_connection *connection)
+static int gb_firmware_get_version(struct gb_firmware *firmware)
 {
+	struct gb_bundle *bundle = firmware->connection->bundle;
+	struct gb_firmware_version_request request;
+	struct gb_firmware_version_response response;
+	int ret;
+
+	request.major = GB_FIRMWARE_VERSION_MAJOR;
+	request.minor = GB_FIRMWARE_VERSION_MINOR;
+
+	ret = gb_operation_sync(firmware->connection,
+				GB_FIRMWARE_TYPE_VERSION,
+				&request, sizeof(request), &response,
+				sizeof(response));
+	if (ret) {
+		dev_err(&bundle->dev,
+				"failed to get protocol version: %d\n",
+				ret);
+		return ret;
+	}
+
+	if (response.major > request.major) {
+		dev_err(&bundle->dev,
+				"unsupported major protocol version (%u > %u)\n",
+				response.major, request.major);
+		return -ENOTSUPP;
+	}
+
+	firmware->protocol_major = response.major;
+	firmware->protocol_minor = response.minor;
+
+	dev_dbg(&bundle->dev, "%s - %u.%u\n", __func__, response.major,
+			response.minor);
+
+	return 0;
+}
+
+static int gb_firmware_probe(struct gb_bundle *bundle,
+					const struct greybus_bundle_id *id)
+{
+	struct greybus_descriptor_cport *cport_desc;
+	struct gb_connection *connection;
 	struct gb_firmware *firmware;
 	int ret;
 
+	if (bundle->num_cports != 1)
+		return -ENODEV;
+
+	cport_desc = &bundle->cport_desc[0];
+	if (cport_desc->protocol_id != GREYBUS_PROTOCOL_FIRMWARE)
+		return -ENODEV;
+
 	firmware = kzalloc(sizeof(*firmware), GFP_KERNEL);
 	if (!firmware)
 		return -ENOMEM;
 
-	firmware->connection = connection;
+	connection = gb_connection_create(bundle,
+						le16_to_cpu(cport_desc->id),
+						gb_firmware_request_handler);
+	if (IS_ERR(connection)) {
+		ret = PTR_ERR(connection);
+		goto err_free_firmware;
+	}
+
 	connection->private = firmware;
 
+	firmware->connection = connection;
+
+	greybus_set_drvdata(bundle, firmware);
+
+	ret = gb_connection_enable_tx(connection);
+	if (ret)
+		goto err_connection_destroy;
+
+	ret = gb_firmware_get_version(firmware);
+	if (ret)
+		goto err_connection_disable;
+
 	firmware_es2_fixup_vid_pid(firmware);
 
+	ret = gb_connection_enable(connection);
+	if (ret)
+		goto err_connection_disable;
+
 	/* Tell bootrom we're ready. */
 	ret = gb_operation_sync(connection, GB_FIRMWARE_TYPE_AP_READY, NULL, 0,
 				NULL, 0);
 	if (ret) {
 		dev_err(&connection->bundle->dev,
 				"failed to send AP READY: %d\n", ret);
-		goto err_free_firmware;
+		goto err_connection_disable;
 	}
 
-	dev_dbg(&connection->bundle->dev, "%s: AP_READY sent\n", __func__);
+	dev_dbg(&bundle->dev, "AP_READY sent\n");
 
 	return 0;
 
+err_connection_disable:
+	gb_connection_disable(connection);
+err_connection_destroy:
+	gb_connection_destroy(connection);
 err_free_firmware:
 	kfree(firmware);
 
 	return ret;
 }
 
-static void gb_firmware_connection_exit(struct gb_connection *connection)
+static void gb_firmware_disconnect(struct gb_bundle *bundle)
 {
-	struct gb_firmware *firmware = connection->private;
+	struct gb_firmware *firmware = greybus_get_drvdata(bundle);
+
+	dev_dbg(&bundle->dev, "%s\n", __func__);
+
+	gb_connection_disable(firmware->connection);
 
 	/* Release firmware */
 	if (firmware->fw)
 		free_firmware(firmware);
 
-	connection->private = NULL;
+	gb_connection_destroy(firmware->connection);
 	kfree(firmware);
-
-	dev_dbg(&connection->bundle->dev, "%s\n", __func__);
 }
 
-static struct gb_protocol firmware_protocol = {
-	.name			= "firmware",
-	.id			= GREYBUS_PROTOCOL_FIRMWARE,
-	.major			= GB_FIRMWARE_VERSION_MAJOR,
-	.minor			= GB_FIRMWARE_VERSION_MINOR,
-	.connection_init	= gb_firmware_connection_init,
-	.connection_exit	= gb_firmware_connection_exit,
-	.request_recv		= gb_firmware_request_recv,
+static const struct greybus_bundle_id gb_firmware_id_table[] = {
+	{ GREYBUS_DEVICE_CLASS(GREYBUS_CLASS_FIRMWARE) },
+	{ }
 };
-gb_builtin_protocol_driver(firmware_protocol);
+
+static struct greybus_driver gb_firmware_driver = {
+	.name		= "firmware",
+	.probe		= gb_firmware_probe,
+	.disconnect	= gb_firmware_disconnect,
+	.id_table	= gb_firmware_id_table,
+};
+
+int gb_firmware_init(void)
+{
+	return greybus_register(&gb_firmware_driver);
+}
+
+void gb_firmware_exit(void)
+{
+	greybus_deregister(&gb_firmware_driver);
+}
diff --git a/drivers/staging/greybus/firmware.h b/drivers/staging/greybus/firmware.h
index 548d297eec63..f25744c1648e 100644
--- a/drivers/staging/greybus/firmware.h
+++ b/drivers/staging/greybus/firmware.h
@@ -10,7 +10,7 @@
 #ifndef __FIRMWARE_H
 #define __FIRMWARE_H
 
-int gb_firmware_protocol_init(void);
-void gb_firmware_protocol_exit(void);
+int gb_firmware_init(void);
+void gb_firmware_exit(void);
 
 #endif /* __FIRMWARE_H */
diff --git a/drivers/staging/greybus/greybus.h b/drivers/staging/greybus/greybus.h
index 27679468e997..2fc79faad28e 100644
--- a/drivers/staging/greybus/greybus.h
+++ b/drivers/staging/greybus/greybus.h
@@ -27,7 +27,6 @@
 #include "manifest.h"
 #include "hd.h"
 #include "svc.h"
-#include "firmware.h"
 #include "control.h"
 #include "interface.h"
 #include "bundle.h"
diff --git a/drivers/staging/greybus/greybus_protocols.h b/drivers/staging/greybus/greybus_protocols.h
index 3f51fb5d6e8b..89db93282cac 100644
--- a/drivers/staging/greybus/greybus_protocols.h
+++ b/drivers/staging/greybus/greybus_protocols.h
@@ -203,6 +203,7 @@ struct gb_control_interface_version_response {
 #define GB_FIRMWARE_VERSION_MINOR		0x01
 
 /* Greybus firmware request types */
+#define GB_FIRMWARE_TYPE_VERSION		0x01
 #define GB_FIRMWARE_TYPE_FIRMWARE_SIZE		0x02
 #define GB_FIRMWARE_TYPE_GET_FIRMWARE		0x03
 #define GB_FIRMWARE_TYPE_READY_TO_BOOT		0x04
@@ -226,6 +227,16 @@ struct gb_control_interface_version_response {
 /* Max firmware data fetch size in bytes */
 #define GB_FIRMWARE_FETCH_MAX			2000
 
+struct gb_firmware_version_request {
+	__u8	major;
+	__u8	minor;
+} __packed;
+
+struct gb_firmware_version_response {
+	__u8	major;
+	__u8	minor;
+} __packed;
+
 /* Firmware protocol firmware size request/response */
 struct gb_firmware_size_request {
 	__u8			stage;
diff --git a/drivers/staging/greybus/legacy.c b/drivers/staging/greybus/legacy.c
index 900521ab6931..057029dc3ee4 100644
--- a/drivers/staging/greybus/legacy.c
+++ b/drivers/staging/greybus/legacy.c
@@ -247,7 +247,6 @@ static const struct greybus_bundle_id legacy_id_table[] = {
 	{ GREYBUS_DEVICE_CLASS(GREYBUS_CLASS_CAMERA) },
 	{ GREYBUS_DEVICE_CLASS(GREYBUS_CLASS_LIGHTS) },
 	{ GREYBUS_DEVICE_CLASS(GREYBUS_CLASS_LOOPBACK) },
-	{ GREYBUS_DEVICE_CLASS(GREYBUS_CLASS_FIRMWARE) },
 	{ GREYBUS_DEVICE_CLASS(GREYBUS_CLASS_RAW) },
 	{ }
 };