diff --git a/flash_size.hpp b/flash_size.hpp
index 4e9e2d0..9ae1c5f 100644
--- a/flash_size.hpp
+++ b/flash_size.hpp
@@ -18,13 +18,21 @@
 
 #include <ipmid/api.h>
 
+#include <ipmid/api-types.hpp>
+#include <vector>
+
 namespace google
 {
 namespace ipmi
 {
 
-ipmi_ret_t getFlashSize(const uint8_t* reqBuf, uint8_t* replyBuf,
-                        size_t* dataLen, HandlerInterface* handler);
+struct GetFlashSizeReply
+{
+    uint32_t flashSize;
+} __attribute__((packed));
+
+Resp getFlashSize(const std::vector<std::uint8_t>& data,
+                  HandlerInterface* handler);
 
 } // namespace ipmi
 } // namespace google