diff --git a/src/tbf.cpp b/src/tbf.cpp
index ccf7362..2afe257 100644
--- a/src/tbf.cpp
+++ b/src/tbf.cpp
@@ -51,6 +51,11 @@
 	m_imsi[sizeof(m_imsi) - 1] = '\0';
 }
 
+void gprs_rlcmac_tbf::set_new_tbf(gprs_rlcmac_tbf *tbf)
+{
+	m_new_tbf = tbf;
+}
+
 gprs_rlcmac_ul_tbf *tbf_alloc_ul(struct gprs_rlcmac_bts *bts,
 	int8_t use_trx, uint8_t ms_class,
 	uint32_t tlli, uint8_t ta, struct gprs_rlcmac_tbf *dl_tbf)
@@ -557,10 +562,10 @@
 				"finished.\n");
 			return NULL;
 		}
-		#warning "THIS should probably go over the IMSI too"
-		new_dl_tbf = ul_tbf->bts->dl_tbf_by_tlli(m_tlli);
-	} else
-		new_dl_tbf = static_cast<gprs_rlcmac_dl_tbf *>(this);
+	}
+
+	new_dl_tbf = static_cast<gprs_rlcmac_dl_tbf *>(m_new_tbf);
+	new_dl_tbf->was_releasing = was_releasing;
 	if (!new_dl_tbf) {
 		LOGP(DRLCMACDL, LOGL_ERROR, "We have a schedule for downlink "
 			"assignment at uplink %s, but there is no downlink "
@@ -625,13 +630,7 @@
 			return NULL;
 	}
 
-	/* on down TBF we get the uplink TBF to be assigned. */
-#warning "Probably want to find by IMSI too"
-	if (direction == GPRS_RLCMAC_DL_TBF)
-		new_tbf = bts->ul_tbf_by_tlli(m_tlli);
-	else
-		new_tbf = static_cast<gprs_rlcmac_ul_tbf *>(this);
-
+	new_tbf = static_cast<gprs_rlcmac_ul_tbf *>(m_new_tbf);
 	if (!new_tbf) {
 		LOGP(DRLCMACUL, LOGL_ERROR, "We have a schedule for uplink "
 			"assignment at downlink %s, but there is no uplink "