From 176733e5dd379d51daf9b521ddd66753ecad9cbd Mon Sep 17 00:00:00 2001
From: Martin Mares <mj@ucw.cz>
Date: Fri, 28 Feb 2020 16:59:21 +0100
Subject: [PATCH] BSB: Check replies more carefully

---
 bsb/daemon/burrow-bsbd.c |  8 +++---
 bsb/firmware/interface.h |  2 +-
 bsb/firmware/main.c      | 53 ++++++++++++++++++++++++++++++++--------
 3 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/bsb/daemon/burrow-bsbd.c b/bsb/daemon/burrow-bsbd.c
index 6b83a5f..05e54ae 100644
--- a/bsb/daemon/burrow-bsbd.c
+++ b/bsb/daemon/burrow-bsbd.c
@@ -255,15 +255,15 @@ static void process_frame(time_t t, byte *pkt, uint len)
 
 	msg(L_DEBUG, "<< %s", hex);
 
-	byte *params = pkt + BF_PARAMS;
-	uint param_len = len - BF_PARAMS - 2;	// 2 for CRC
+	byte *body = pkt + BF_BODY;
+	uint body_len = len - BF_BODY - 2;	// 2 for CRC
 
 	switch (pkt[BF_OP]) {
 		case BSB_OP_INFO:
-			process_info(params, param_len);
+			process_info(body, body_len);
 			break;
 		case BSB_OP_ANSWER:
-			process_answer(params, param_len);
+			process_answer(body, body_len);
 			break;
 	}
 }
diff --git a/bsb/firmware/interface.h b/bsb/firmware/interface.h
index 8c032c4..74e57a7 100644
--- a/bsb/firmware/interface.h
+++ b/bsb/firmware/interface.h
@@ -84,7 +84,7 @@ enum bsb_frame {
 	BF_DEST,
 	BF_LEN,
 	BF_OP,
-	BF_PARAMS,
+	BF_BODY,
 };
 
 enum bsb_address {
diff --git a/bsb/firmware/main.c b/bsb/firmware/main.c
index bb0c48e..31a64a6 100644
--- a/bsb/firmware/main.c
+++ b/bsb/firmware/main.c
@@ -144,10 +144,10 @@ struct bsb_state {
 	byte tx_check_index;
 	byte tx_attempts;
 	byte tx_result;			// TX_RESULT_xxx
+	byte tx_want_reply;
 	u32 tx_start_timestamp;
 	u32 tx_ticks_to_send;
 	u32 tx_end_timestamp;
-	u32 reply_type_mask;		// Expected reply types (0=no reply)
 
 	byte rx_buf[BSB_MAX_SIZE];
 	byte tx_buf[BSB_MAX_SIZE];
@@ -337,7 +337,7 @@ static void bsb_tx_step(void)
 			bsb_attempt_tx();
 			return;
 		case TX_STATE_SENT:
-			if (!bsb.reply_type_mask) {
+			if (!bsb.tx_want_reply) {
 				// The frame needs no reply => confirm transmission
 				bsb_stats.tx_ok_no_reply++;
 				bsb.tx_state = TX_STATE_OFF;
@@ -400,7 +400,7 @@ static void bsb_send(byte *buf, uint len)
 	// Detect frame type and prepare expected reply types
 	if (buf[BF_OP] == BSB_OP_QUERY) {
 		// So far, we allow only QUERY frames
-		bsb.reply_type_mask = (1 << BSB_OP_ANSWER) | (1 << BSB_OP_ERROR);
+		bsb.tx_want_reply = 1;
 	} else {
 		bsb_stats.tx_rejects++;
 		bsb_tx_result = TX_RESULT_FORBIDDEN;
@@ -416,7 +416,7 @@ static void bsb_send(byte *buf, uint len)
 		crc = bsb_crc_update(crc, buf[i]);
 	buf[len-2] = crc >> 8;
 	buf[len-1] = crc;
-	debug_printf("TX: Prepared frame: type=%02x dest=%02x len=%u crc=%04x reply=%08x\n", buf[BF_OP], buf[BF_DEST], len, crc, (uint) bsb.reply_type_mask);
+	debug_printf("TX: Prepared frame: type=%02x dest=%02x len=%u crc=%04x want_reply=%u\n", buf[BF_OP], buf[BF_DEST], len, crc, bsb.tx_want_reply);
 
 	// Queue frame for transmission
 	memcpy(bsb.tx_buf, buf, len);
@@ -429,6 +429,42 @@ static void bsb_send(byte *buf, uint len)
 	gpio_clear(GPIOB, GPIO0);
 }
 
+static bool bsb_is_reply_for(byte *reply, byte *req) {
+	// Check source and destination address
+	if (reply[BF_DEST] != BSB_ADDR_GATEWAY ||
+	    reply[BF_SRC] != req[BF_DEST] ^ 0x80)
+		return 0;
+
+	// Extract body of request and reply
+	byte *rq = req + BF_BODY;
+	byte *rp = reply + BF_BODY;
+	uint rq_len = req[BF_LEN];
+	uint rp_len = reply[BF_LEN];
+	if (rq_len < BF_BODY + 2 || rp_len < BF_BODY + 2)
+		return 0;
+	rq_len -= BF_BODY + 2;
+	rp_len -= BF_BODY + 2;
+
+	switch (req[BF_OP]) {
+		case BSB_OP_QUERY:
+			switch (reply[BF_OP]) {
+				case BSB_OP_ANSWER:
+					if (rq_len >= 4 && rp_len >= 4 &&
+					    rq[0] == rp[1] &&		// This is byte-swapped. Why?
+					    rq[1] == rp[0] &&
+					    rq[2] == rp[2] &&
+					    rq[3] == rp[3])
+						return 1;
+					return 0;
+				case BSB_OP_ERROR:
+					return 1;
+			}
+			return 0;
+		default:
+			return 0;
+	}
+}
+
 static void bsb_check_reply(void)
 {
 	// Match received frame in bsb_rx_status_and_frame to expected replies
@@ -436,18 +472,15 @@ static void bsb_check_reply(void)
 
 	byte *frame = bsb_rx_status_and_frame + 1;
 	// XXX: This might be executed too fast, so we can still be in TX_STATE_SENT
-	if (bsb.tx_state == TX_STATE_WAITING_FOR_REPLY || bsb.tx_state == TX_STATE_SENT) {
-		if (frame[BF_DEST] == BSB_ADDR_GATEWAY &&
-		    frame[BF_SRC] == bsb.tx_buf[BF_DEST] ^ 0x80 &&
-		    frame[BF_OP] < 32 &&
-		    bsb.reply_type_mask & (1 << frame[BF_OP])) {
+	if ((bsb.tx_state == TX_STATE_WAITING_FOR_REPLY || bsb.tx_state == TX_STATE_SENT) &&
+	     bsb.tx_want_reply &&
+	     bsb_is_reply_for(frame, bsb.tx_buf)) {
 			debug_printf("BSB: Matched reply\n");
 			bsb.tx_state = TX_STATE_OFF;
 			bsb_stats.tx_ok_replied++;
 			bsb_rx_status_and_frame[0] = TX_RESULT_OK;
 			gpio_set(GPIOB, GPIO0);		// Turn the TX LED off
 			return;
-		}
 	}
 
 	bsb_rx_status_and_frame[0] = TX_RESULT_NONE;
-- 
2.39.2