2 files changed, 16 insertions, 5 deletions
diff --git a/src/lib/modes/mode_pad/mode_pad.cpp b/src/lib/modes/mode_pad/mode_pad.cpp
index 5c949e9cf..be3ecf7dc 100644
--- a/src/lib/modes/mode_pad/mode_pad.cpp
+++ b/src/lib/modes/mode_pad/mode_pad.cpp
@@ -53,7 +53,7 @@ void PKCS7_Padding::add_padding(secure_vector<uint8_t>& buffer,
 */
 size_t PKCS7_Padding::unpad(const uint8_t input[], size_t input_length) const
    {
-   if(input_length <= 2)
+   if(!valid_blocksize(input_length))
       return input_length;
 
    CT::poison(input, input_length);
@@ -104,7 +104,7 @@ void ANSI_X923_Padding::add_padding(secure_vector<uint8_t>& buffer,
 */
 size_t ANSI_X923_Padding::unpad(const uint8_t input[], size_t input_length) const
    {
-   if(input_length <= 2)
+   if(!valid_blocksize(input_length))
       return input_length;
 
    CT::poison(input, input_length);
@@ -146,7 +146,7 @@ void OneAndZeros_Padding::add_padding(secure_vector<uint8_t>& buffer,
 */
 size_t OneAndZeros_Padding::unpad(const uint8_t input[], size_t input_length) const
    {
-   if(input_length <= 2)
+   if(!valid_blocksize(input_length))
       return input_length;
 
    CT::poison(input, input_length);
@@ -170,7 +170,8 @@ size_t OneAndZeros_Padding::unpad(const uint8_t input[], size_t input_length) co
    bad_input |= ~seen_0x80;
 
    CT::unpoison(input, input_length);
-   return bad_input.select_and_unpoison(input_length, pad_pos);
+
+   return CT::Mask<size_t>::expand(bad_input).select_and_unpoison(input_length, pad_pos);
    }
 
 /*
@@ -193,7 +194,7 @@ void ESP_Padding::add_padding(secure_vector<uint8_t>& buffer,
 */
 size_t ESP_Padding::unpad(const uint8_t input[], size_t input_length) const
    {
-   if(input_length <= 2)
+   if(!valid_blocksize(input_length))
       return input_length;
 
    CT::poison(input, input_length);
diff --git a/src/lib/utils/ct_utils.h b/src/lib/utils/ct_utils.h
index eb510baa2..9243d6701 100644
--- a/src/lib/utils/ct_utils.h
+++ b/src/lib/utils/ct_utils.h
@@ -125,6 +125,16 @@ class Mask
          }
 
       /**
+      * Return a Mask<T> which is set if m is set
+      */
+      template<typename U>
+      static Mask<T> expand(Mask<U> m)
+         {
+         static_assert(sizeof(U) < sizeof(T), "sizes ok");
+         return ~Mask<T>::is_zero(m.value());
+         }
+
+      /**
       * Return a Mask<T> which is set if v is == 0 or cleared otherwise
       */
       static Mask<T> is_zero(T x)