Support other byte lengths in bytestring_to_count

src/zeek.bif

@@ -2961,26 +2961,32 @@ function bytestring_to_count%(s: string, is_le: bool &default=F%): count
 	static const bool host_bigendian = false;
 #endif
 	const u_char *p = s->Bytes();
-	unsigned int i;
+	unsigned int i = 0;
 
-	switch ( s->Len() ) {
+	if ( s->Len() == 0 )
-	case sizeof(uint8_t):
+		{
+		return zeek::val_mgr->Count(0);
+		}
+	if ( s->Len() == sizeof(uint8_t) )
 		{
 		uint8_t value = 0;
 		memcpy(&value, p, sizeof(uint8_t));
 		return zeek::val_mgr->Count(value);
 		}
-
+	else if ( s->Len() <= sizeof(uint16_t) )
-	case sizeof(uint16_t):
 		{
 		uint16_t value = 0;
 
 		if ( (host_bigendian && is_le) || (! host_bigendian && ! is_le) )
 			{
 			char buf[sizeof(uint16_t)];
+			memset(buf, 0, sizeof(uint16_t));
 			char *d = &buf[sizeof(uint16_t)-1];
 
-			for ( i = 0; i < sizeof(uint16_t); i++ )
+			if ( is_le )
+				d -= sizeof(uint16_t) - s->Len();
+
+			for ( i = 0; i < s->Len(); i++ )
 				*d-- = *p++;
 
 			memcpy(&value, buf, sizeof(uint16_t));
@@ -2990,46 +2996,65 @@ function bytestring_to_count%(s: string, is_le: bool &default=F%): count
 
 		return zeek::val_mgr->Count(value);
 		}
-
+	else if ( s->Len() <= sizeof(uint32_t) )
-	case sizeof(uint32_t):
 		{
 		uint32_t value = 0;
+		char buf[sizeof(uint32_t)];
+		memset(buf, 0, sizeof(uint32_t));
 
 		if ( (host_bigendian && is_le) || (! host_bigendian && ! is_le) )
 			{
-			char buf[sizeof(uint32_t)];
 			char *d = &buf[sizeof(uint32_t)-1];
 
-			for ( i = 0; i < sizeof(uint32_t); i++ )
+			if ( ! is_le )
-				*d-- = *p++;
+				d -= sizeof(uint32_t) - s->Len();
 
-			memcpy(&value, buf, sizeof(uint32_t));
+			for ( i = 0; i < s->Len(); i++ )
+				*d-- = *p++;
 			}
 		else
-			memcpy(&value, p, sizeof(uint32_t));
+			{
+			char* d = &buf[0];
 
-		return zeek::val_mgr->Count(value);
+			if ( ! is_le )
+				d += sizeof(uint32_t) - s->Len();
+
+			for ( i = 0; i < s->Len(); i++ )
+				*d++ = *p++;
 			}
 
-	case sizeof(uint64_t):
+		memcpy(&value, buf, sizeof(uint32_t));
+		return zeek::val_mgr->Count(value);
+		}
+	else if ( s->Len() <= sizeof(uint64_t) )
 		{
 		uint64_t value = 0;
+		char buf[sizeof(uint64_t)];
+		memset(buf, 0, sizeof(uint64_t));
 
 		if ( (host_bigendian && is_le) || (! host_bigendian && ! is_le) )
 			{
-			char buf[sizeof(uint64_t)];
 			char *d = &buf[sizeof(uint64_t)-1];
 
-			for ( i = 0; i < sizeof(uint64_t); i++ )
+			if ( ! is_le )
-				*d-- = *p++;
+				d -= sizeof(uint64_t) - s->Len();
 
-			memcpy(&value, buf, sizeof(uint64_t));
+			for ( i = 0; i < s->Len(); i++ )
+				*d-- = *p++;
 			}
 		else
-			memcpy(&value, p, sizeof(uint64_t));
+			{
+			char* d = &buf[0];
 
-		return zeek::val_mgr->Count(value);
+			if ( ! is_le )
+				d += sizeof(uint64_t) - s->Len();
+
+			for ( i = 0; i < s->Len(); i++ )
+				*d++ = *p++;
 			}
+
+		memcpy(&value, buf, sizeof(uint64_t));
+		return zeek::val_mgr->Count(value);
 		}
 
 	zeek::emit_builtin_error("unsupported byte length for bytestring_to_count");

testing/btest/Baseline/bifs.bytestring_to_count/out

@@ -1,10 +1,6 @@
 ### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
 0
 0
-0
-0
-0
-0
 255
 255
 0
@@ -27,6 +23,10 @@
 2864429994
 0
 0
+65535
+65535
+16777215
+16777215
 18446744073709551615
 18446744073709551615
 18446742974214701055
@@ -35,4 +35,10 @@
 65535
 0
 0
+1099511627775
+1099511627775
+281474976710655
+281474976710655
+72057594037927935
+72057594037927935
 65535

testing/btest/bifs/bytestring_to_count.zeek

@@ -9,10 +9,6 @@ event zeek_init()
 	# unsupported byte lengths
 	print bytestring_to_count("", T); # 0
 	print bytestring_to_count("", F); # 0
-	print bytestring_to_count("\xAA\xBB\xCC", T); # 0
-	print bytestring_to_count("\xAA\xBB\xCC", F); # 0
-	print bytestring_to_count("\xAA\xBB\xCC\xDD\xEE", T); # 0
-	print bytestring_to_count("\xAA\xBB\xCC\xDD\xEE", F); # 0
 
 	# 8 bit
 	print bytestring_to_count("\xff", T); # 255
@@ -42,6 +38,11 @@ event zeek_init()
 	print bytestring_to_count("\x00\x00\x00\x00", F); # 0
 	print bytestring_to_count("\x00\x00\x00\x00", T); # 0
 
+	print bytestring_to_count("\x00\xff\xff", F); # 65535
+	print bytestring_to_count("\xff\xff\x00", T); # 65535
+	print bytestring_to_count("\xff\xff\xff", F); # 16777215
+	print bytestring_to_count("\xff\xff\xff", T); # 16777215
+
 	# 64 bit
 	print bytestring_to_count("\xff\xff\xff\xff\xff\xff\xff\xff", F); # 18446744073709551615
 	print bytestring_to_count("\xff\xff\xff\xff\xff\xff\xff\xff", T); # 18446744073709551615
@@ -52,6 +53,13 @@ event zeek_init()
 	print bytestring_to_count("\x00\x00\x00\x00\x00\x00\x00\x00", T); # 0
 	print bytestring_to_count("\x00\x00\x00\x00\x00\x00\x00\x00", F); # 0
 
+	print bytestring_to_count("\xff\xff\xff\xff\xff", F);         # 1099511627775
+	print bytestring_to_count("\xff\xff\xff\xff\xff", T);         # 1099511627775
+	print bytestring_to_count("\xff\xff\xff\xff\xff\xff", F);     # 281474976710655
+	print bytestring_to_count("\xff\xff\xff\xff\xff\xff", T);     # 281474976710655
+	print bytestring_to_count("\xff\xff\xff\xff\xff\xff\xff", F); # 72057594037927935
+	print bytestring_to_count("\xff\xff\xff\xff\xff\xff\xff", T); # 72057594037927935
+
 	# test the default endianness parameter
 	print bytestring_to_count("\x00\x00\x00\x00\x00\x00\xff\xff"); # 65535