return rc;
 }
 
+/**
+ * _ctl_diag_get_bufftype - return diag buffer type
+ *              either TRACE, SNAPSHOT, or EXTENDED
+ * @ioc: per adapter object
+ * @unique_id: specifies the unique_id for the buffer
+ *
+ * returns MPT3_DIAG_UID_NOT_FOUND if the id not found
+ */
+static u8
+_ctl_diag_get_bufftype(struct MPT3SAS_ADAPTER *ioc, u32 unique_id)
+{
+       u8  index;
+
+       for (index = 0; index < MPI2_DIAG_BUF_TYPE_COUNT; index++) {
+               if (ioc->unique_id[index] == unique_id)
+                       return index;
+       }
+
+       return MPT3_DIAG_UID_NOT_FOUND;
+}
 
 /**
  * _ctl_diag_register_2 - wrapper for registering diag buffer support
                return -EPERM;
        }
 
+       if (diag_register->unique_id == 0) {
+               ioc_err(ioc,
+                   "%s: Invalid UID(0x%08x), buffer_type(0x%02x)\n", __func__,
+                   diag_register->unique_id, buffer_type);
+               return -EINVAL;
+       }
+
        if (ioc->diag_buffer_status[buffer_type] &
            MPT3_DIAG_BUFFER_IS_REGISTERED) {
-               ioc_err(ioc, "%s: already has a registered buffer for buffer_type(0x%02x)\n",
-                       __func__, buffer_type);
-               return -EINVAL;
+               /*
+                * If driver posts buffer initially, then an application wants
+                * to Register that buffer (own it) without Releasing first,
+                * the application Register command MUST have the same buffer
+                * type and size in the Register command (obtained from the
+                * Query command). Otherwise that Register command will be
+                * failed. If the application has released the buffer but wants
+                * to re-register it, it should be allowed as long as the
+                * Unique-Id/Size match.
+                */
+
+               if (ioc->unique_id[buffer_type] == MPT3DIAGBUFFUNIQUEID &&
+                   ioc->diag_buffer_sz[buffer_type] ==
+                   diag_register->requested_buffer_size) {
+
+                       if (!(ioc->diag_buffer_status[buffer_type] &
+                            MPT3_DIAG_BUFFER_IS_RELEASED)) {
+                               dctlprintk(ioc, ioc_info(ioc,
+                                   "%s: diag_buffer (%d) ownership changed. old-ID(0x%08x), new-ID(0x%08x)\n",
+                                   __func__, buffer_type,
+                                   ioc->unique_id[buffer_type],
+                                   diag_register->unique_id));
+
+                               /*
+                                * Application wants to own the buffer with
+                                * the same size.
+                                */
+                               ioc->unique_id[buffer_type] =
+                                   diag_register->unique_id;
+                               rc = 0; /* success */
+                               goto out;
+                       }
+               } else if (ioc->unique_id[buffer_type] !=
+                   MPT3DIAGBUFFUNIQUEID) {
+                       if (ioc->unique_id[buffer_type] !=
+                           diag_register->unique_id ||
+                           ioc->diag_buffer_sz[buffer_type] !=
+                           diag_register->requested_buffer_size ||
+                           !(ioc->diag_buffer_status[buffer_type] &
+                           MPT3_DIAG_BUFFER_IS_RELEASED)) {
+                               ioc_err(ioc,
+                                   "%s: already has a registered buffer for buffer_type(0x%02x)\n",
+                                   __func__, buffer_type);
+                               return -EINVAL;
+                       }
+               } else {
+                       ioc_err(ioc, "%s: already has a registered buffer for buffer_type(0x%02x)\n",
+                           __func__, buffer_type);
+                       return -EINVAL;
+               }
        }
 
        if (diag_register->requested_buffer_size % 4)  {
                ioc->diag_trigger_master.MasterData =
                    (MASTER_TRIGGER_FW_FAULT + MASTER_TRIGGER_ADAPTER_RESET);
                diag_register.buffer_type = MPI2_DIAG_BUF_TYPE_TRACE;
-               diag_register.unique_id = 0x7075900;
+               diag_register.unique_id =
+                   (ioc->hba_mpi_version_belonged == MPI2_VERSION) ?
+                   (MPT2DIAGBUFFUNIQUEID):(MPT3DIAGBUFFUNIQUEID);
 
                if (trace_buff_size != 0) {
                        diag_register.requested_buffer_size = trace_buff_size;
        dctlprintk(ioc, ioc_info(ioc, "%s\n",
                                 __func__));
 
-       buffer_type = karg.unique_id & 0x000000ff;
+       buffer_type = _ctl_diag_get_bufftype(ioc, karg.unique_id);
+       if (buffer_type == MPT3_DIAG_UID_NOT_FOUND) {
+               ioc_err(ioc, "%s: buffer with unique_id(0x%08x) not found\n",
+                   __func__, karg.unique_id);
+               return -EINVAL;
+       }
+
        if (!_ctl_diag_capability(ioc, buffer_type)) {
                ioc_err(ioc, "%s: doesn't have capability for buffer_type(0x%02x)\n",
                        __func__, buffer_type);
                return -EINVAL;
        }
 
-       if (karg.unique_id & 0xffffff00) {
+       if (karg.unique_id) {
                if (karg.unique_id != ioc->unique_id[buffer_type]) {
                        ioc_err(ioc, "%s: unique_id(0x%08x) is not registered\n",
                                __func__, karg.unique_id);
        dctlprintk(ioc, ioc_info(ioc, "%s\n",
                                 __func__));
 
-       buffer_type = karg.unique_id & 0x000000ff;
+       buffer_type = _ctl_diag_get_bufftype(ioc, karg.unique_id);
+       if (buffer_type == MPT3_DIAG_UID_NOT_FOUND) {
+               ioc_err(ioc, "%s: buffer with unique_id(0x%08x) not found\n",
+                   __func__, karg.unique_id);
+               return -EINVAL;
+       }
+
        if (!_ctl_diag_capability(ioc, buffer_type)) {
                ioc_err(ioc, "%s: doesn't have capability for buffer_type(0x%02x)\n",
                        __func__, buffer_type);
        dctlprintk(ioc, ioc_info(ioc, "%s\n",
                                 __func__));
 
-       buffer_type = karg.unique_id & 0x000000ff;
+       buffer_type = _ctl_diag_get_bufftype(ioc, karg.unique_id);
+       if (buffer_type == MPT3_DIAG_UID_NOT_FOUND) {
+               ioc_err(ioc, "%s: buffer with unique_id(0x%08x) not found\n",
+                   __func__, karg.unique_id);
+               return -EINVAL;
+       }
+
        if (!_ctl_diag_capability(ioc, buffer_type)) {
                ioc_err(ioc, "%s: doesn't have capability for buffer_type(0x%02x)\n",
                        __func__, buffer_type);
                } else
                        diag_register.requested_buffer_size = (1024 * 1024);
 
-               diag_register.unique_id = 0x7075900;
+               diag_register.unique_id =
+                   (ioc->hba_mpi_version_belonged == MPI2_VERSION) ?
+                   (MPT2DIAGBUFFUNIQUEID):(MPT3DIAGBUFFUNIQUEID);
                ioc->diag_buffer_status[MPI2_DIAG_BUF_TYPE_TRACE] = 0;
                _ctl_diag_register_2(ioc,  &diag_register);
        } else if (!strcmp(str, "release")) {