Merge pull request #333 from DDoSolitary/patch-radiotap-writer

Overhaul RadioTapWriter

2 files changed, 27 insertions, 3 deletions

diff --git a/src/utils/radiotap_writer.cpp b/src/utils/radiotap_writer.cpp
index 3a18187..4bb3465 100644
--- a/src/utils/radiotap_writer.cpp
+++ b/src/utils/radiotap_writer.cpp
@@ -41,11 +41,12 @@ namespace Tins {
 namespace Utils {
 
 uint32_t calculate_padding(uint32_t alignment, uint32_t offset) {
-    return offset % alignment;
+    uint32_t extra = offset % alignment;
+    return extra == 0 ? 0 : alignment - extra;
 }
 
 uint32_t get_bit(uint32_t value) {
-    return log(value) / log(2);
+    return round(log2(value));
 }
 
 RadioTapWriter::RadioTapWriter(vector<uint8_t>& buffer)
@@ -54,7 +55,7 @@ RadioTapWriter::RadioTapWriter(vector<uint8_t>& buffer)
 
 void RadioTapWriter::write_option(const RadioTap::option& option) {
     const uint32_t bit = get_bit(option.option());
-    if (bit > RadioTapParser::MAX_RADIOTAP_FIELD) {
+    if (bit >= RadioTapParser::MAX_RADIOTAP_FIELD) {
         throw malformed_option();
     }
     const bool is_empty = buffer_.empty();
diff --git a/tests/src/radiotap_test.cpp b/tests/src/radiotap_test.cpp
index 04d6c6b..7670178 100644
--- a/tests/src/radiotap_test.cpp
+++ b/tests/src/radiotap_test.cpp
@@ -759,4 +759,27 @@ TEST_F(RadioTapTest, RadioTapWritingEmptyBuffer) {
 
 }
 
+TEST_F(RadioTapTest, RadioTapWritingInvalidOption) {
+    vector<uint8_t> buffer(4, 0);
+    RadioTapWriter writer(buffer);
+    uint8_t foo = 0;
+    RadioTap::option option((RadioTap::PresentFlags)(1 << RadioTapParser::MAX_RADIOTAP_FIELD), sizeof(foo), &foo);
+    EXPECT_THROW(writer.write_option(option), malformed_option);
+}
+
+TEST_F(RadioTapTest, RadioTapWriterAlignment) {
+    vector<uint8_t> buffer(4, 0);
+    RadioTapWriter writer(buffer);
+    uint8_t flags = 10;
+    uint8_t xchannel[sizeof(RadioTap::xchannel_type)] = {
+        1, 2, 3, 4, 5, 6, 7, 8
+    };
+    writer.write_option(RadioTap::option(RadioTap::FLAGS, sizeof(flags), &flags));
+    writer.write_option(RadioTap::option(RadioTap::XCHANNEL, sizeof(xchannel), xchannel));
+    vector<uint8_t> expected = {
+        2, 0, 4, 0, 10, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8
+    };
+    EXPECT_EQ(buffer, expected);
+}
+
 #endif // TINS_HAVE_DOT11