greybus: protocol: add a module owner to a protocol

Now that protocols can be in a module, we need to reference count them
to lock them into memory so they can't be removed while in use.  So add
a module owner structure, and have it automatically be assigned when
registering the protocol.
Signed-off-by: Greg Kroah-Hartman <greg@kroah.com>
Reviewed-by: Alex Elder <elder@linaro.org>

drivers/staging/greybus/protocol.c

@@ -39,13 +39,15 @@ static struct gb_protocol *_gb_protocol_find(u8 id, u8 major, u8 minor)
 	return NULL;
 }
 
-int gb_protocol_register(struct gb_protocol *protocol)
+int __gb_protocol_register(struct gb_protocol *protocol, struct module *module)
 {
 	struct gb_protocol *existing;
 	u8 id = protocol->id;
 	u8 major = protocol->major;
 	u8 minor = protocol->minor;
 
+	protocol->owner = module;
+
 	/*
 	 * The protocols list is sorted first by protocol id (low to
 	 * high), then by major version (high to low), and finally
@@ -92,7 +94,7 @@ int gb_protocol_register(struct gb_protocol *protocol)
 
 	return 0;
 }
-EXPORT_SYMBOL_GPL(gb_protocol_register);
+EXPORT_SYMBOL_GPL(__gb_protocol_register);
 
 /*
  * De-register a previously registered protocol.
@@ -135,9 +137,13 @@ struct gb_protocol *gb_protocol_get(u8 id, u8 major, u8 minor)
 	spin_lock_irq(&gb_protocols_lock);
 	protocol = _gb_protocol_find(id, major, minor);
 	if (protocol) {
-		protocol_count = protocol->count;
-		if (protocol_count != U8_MAX)
-			protocol->count++;
+		if (!try_module_get(protocol->owner)) {
+			protocol = NULL;
+		} else {
+			protocol_count = protocol->count;
+			if (protocol_count != U8_MAX)
+				protocol->count++;
+		}
 	}
 	spin_unlock_irq(&gb_protocols_lock);
 
@@ -163,6 +169,7 @@ void gb_protocol_put(struct gb_protocol *protocol)
 		protocol_count = protocol->count;
 		if (protocol_count)
 			protocol->count--;
+		module_put(protocol->owner);
 	}
 	spin_unlock_irq(&gb_protocols_lock);
 	if (protocol)

drivers/staging/greybus/protocol.h

@@ -34,11 +34,15 @@ struct gb_protocol {
 	gb_connection_init_t	connection_init;
 	gb_connection_exit_t	connection_exit;
 	gb_request_recv_t	request_recv;
+	struct module		*owner;
 };
 
-int gb_protocol_register(struct gb_protocol *protocol);
+int __gb_protocol_register(struct gb_protocol *protocol, struct module *module);
 int gb_protocol_deregister(struct gb_protocol *protocol);
 
+#define gb_protocol_register(protocol) \
+	__gb_protocol_register(protocol, THIS_MODULE)
+
 struct gb_protocol *gb_protocol_get(u8 id, u8 major, u8 minor);
 void gb_protocol_put(struct gb_protocol *protocol);