improving shortstr overflow checking

2 files changed, 18 insertions, 6 deletions

diff --git a/codegen.py b/codegen.py
index 6f39574f26..648983f1f2 100644
--- a/codegen.py
+++ b/codegen.py
@@ -214,6 +214,8 @@ def genErl(spec):
             elif type == 'table':
                 print "  F%d = rabbit_binary_parser:parse_table(F%dTab)," % \
                       (f.index, f.index)
+            elif type == 'shortstr':
+                print "  if F%dLen > 255 -> exit(method_field_shortstr_overflow); true -> ok end," % (f.index)
             else:
                 pass
 
@@ -246,7 +248,10 @@ def genErl(spec):
             elif type == 'table':
                 print "  F%dTab = rabbit_binary_generator:generate_table(F%d)," % (f.index, f.index)
                 print "  F%dLen = size(F%dTab)," % (f.index, f.index)
-            elif type in ['shortstr', 'longstr']:
+            elif type == 'shortstr':
+                print "  F%dLen = size(F%d)," % (f.index, f.index)
+                print "  if F%dLen > 255 -> exit(method_field_shortstr_overflow); true -> ok end," % (f.index)
+            elif type == 'longstr':
                 print "  F%dLen = size(F%d)," % (f.index, f.index)
             else:
                 pass
diff --git a/src/rabbit_binary_generator.erl b/src/rabbit_binary_generator.erl
index b8e161a6bd..b903a6eee7 100644
--- a/src/rabbit_binary_generator.erl
+++ b/src/rabbit_binary_generator.erl
@@ -196,12 +196,16 @@ generate_array(Array) when is_list(Array) ->
                      fun ({Type, Value}) -> field_value_to_binary(Type, Value) end,
                      Array)).
 
-short_string_to_binary(String) when is_binary(String) and (size(String) < 256) ->
-    [<<(size(String)):8>>, String];
+short_string_to_binary(String) when is_binary(String) ->
+    Len = size(String),
+    if Len < 256 -> [<<(size(String)):8>>, String];
+       true      -> exit(content_properties_shortstr_overflow)
+    end;
 short_string_to_binary(String) ->
     StringLength = length(String),
-    true = (StringLength < 256), % assertion
-    [<<StringLength:8>>, String].
+    if StringLength < 256 -> [<<StringLength:8>>, String];
+       true               -> exit(content_properties_shortstr_overflow)
+    end.
 
 long_string_to_binary(String) when is_binary(String) ->
     [<<(size(String)):32>>, String];
@@ -239,7 +243,10 @@ encode_properties(Bit, [T | TypeList], [Value | ValueList], FirstShortAcc, Flags
     end.
 
 encode_property(shortstr, String) ->
-    Len = size(String), <<Len:8/unsigned, String:Len/binary>>;
+    Len = size(String),
+    if Len < 256 -> <<Len:8/unsigned, String:Len/binary>>;
+       true      -> exit(content_properties_shortstr_overflow)
+    end;
 encode_property(longstr, String) ->
     Len = size(String), <<Len:32/unsigned, String:Len/binary>>;
 encode_property(octet, Int) ->