1 files changed, 129 insertions, 1 deletions

diff --git a/src/tests/test_keywrap.cpp b/src/tests/test_keywrap.cpp
index d710dfc68..a5b216562 100644
--- a/src/tests/test_keywrap.cpp
+++ b/src/tests/test_keywrap.cpp
@@ -10,6 +10,11 @@
    #include <botan/rfc3394.h>
 #endif
 
+#if defined(BOTAN_HAS_NIST_KEYWRAP)
+   #include <botan/nist_keywrap.h>
+   #include <botan/block_cipher.h>
+#endif
+
 namespace Botan_Tests {
 
 namespace {
@@ -18,7 +23,7 @@ namespace {
 class RFC3394_Keywrap_Tests final : public Text_Based_Test
    {
    public:
-      RFC3394_Keywrap_Tests() : Text_Based_Test("rfc3394.vec", "Key,KEK,Output") {}
+      RFC3394_Keywrap_Tests() : Text_Based_Test("keywrap/rfc3394.vec", "Key,KEK,Output") {}
 
       Test::Result run_one_test(const std::string&, const VarMap& vars) override
          {
@@ -50,6 +55,129 @@ class RFC3394_Keywrap_Tests final : public Text_Based_Test
 BOTAN_REGISTER_TEST("rfc3394", RFC3394_Keywrap_Tests);
 #endif
 
+#if defined(BOTAN_HAS_NIST_KEYWRAP) && defined(BOTAN_HAS_AES)
+
+class NIST_Keywrap_Tests final : public Text_Based_Test
+   {
+   public:
+      NIST_Keywrap_Tests() : Text_Based_Test("keywrap/nist_key_wrap.vec", "Input,Key,Output") {}
+
+      Test::Result run_one_test(const std::string& typ, const VarMap& vars) override
+         {
+         Test::Result result("NIST keywrap");
+
+         try
+            {
+            if(typ != "KW" && typ != "KWP")
+               throw Test_Error("Unknown type in NIST key wrap tests");
+
+            const std::vector<uint8_t> expected = get_req_bin(vars, "Output");
+            const std::vector<uint8_t> input = get_req_bin(vars, "Input");
+            const std::vector<uint8_t> key = get_req_bin(vars, "Key");
+
+            std::unique_ptr<Botan::BlockCipher> bc =
+               Botan::BlockCipher::create_or_throw("AES-" + std::to_string(key.size()*8));
+
+            bc->set_key(key);
+
+            std::vector<uint8_t> wrapped;
+
+            if(typ == "KW")
+               {
+               wrapped = nist_key_wrap(input.data(), input.size(), *bc);
+               }
+            else if(typ == "KWP")
+               {
+               wrapped = nist_key_wrap_padded(input.data(), input.size(), *bc);
+               }
+
+            result.test_eq("key wrap", wrapped, expected);
+
+            try
+               {
+               Botan::secure_vector<uint8_t> unwrapped;
+               if(typ == "KW")
+                  {
+                  unwrapped = nist_key_unwrap(expected.data(), expected.size(), *bc);
+                  }
+               else if(typ == "KWP")
+                  {
+                  unwrapped = nist_key_unwrap_padded(expected.data(), expected.size(), *bc);
+                  }
+
+               result.test_eq("key unwrap", unwrapped, input);
+               }
+            catch(Botan::Integrity_Failure& e)
+               {
+               result.test_failure("NIST key unwrap failed with integrity failure", e.what());
+               }
+            }
+         catch(std::exception& e)
+            {
+            result.test_failure("", e.what());
+            }
+
+         return result;
+         }
+
+   };
+
+BOTAN_REGISTER_TEST("nist_key_wrap", NIST_Keywrap_Tests);
+
+class NIST_Keywrap_Invalid_Tests final : public Text_Based_Test
+   {
+   public:
+      NIST_Keywrap_Invalid_Tests() : Text_Based_Test("keywrap/nist_key_wrap_invalid.vec", "Key,Input") {}
+
+      Test::Result run_one_test(const std::string& typ, const VarMap& vars) override
+         {
+         Test::Result result("NIST keywrap (invalid inputs)");
+
+         try
+            {
+            if(typ != "KW" && typ != "KWP")
+               throw Test_Error("Unknown type in NIST key wrap tests");
+
+            const std::vector<uint8_t> input = get_req_bin(vars, "Input");
+            const std::vector<uint8_t> key = get_req_bin(vars, "Key");
+
+            std::unique_ptr<Botan::BlockCipher> bc =
+               Botan::BlockCipher::create_or_throw("AES-" + std::to_string(key.size()*8));
+
+            bc->set_key(key);
+
+            try
+               {
+               Botan::secure_vector<uint8_t> unwrapped;
+               if(typ == "KW")
+                  {
+                  unwrapped = nist_key_unwrap(input.data(), input.size(), *bc);
+                  }
+               else if(typ == "KWP")
+                  {
+                  unwrapped = nist_key_unwrap_padded(input.data(), input.size(), *bc);
+                  }
+
+               result.test_failure("Was able to unwrap invalid keywrap input");
+               }
+            catch(Botan::Integrity_Failure)
+               {
+               result.test_success("Rejected invalid input");
+               }
+            }
+         catch(std::exception& e)
+            {
+            result.test_failure("", e.what());
+            }
+
+         return result;
+         }
+
+   };
+
+BOTAN_REGISTER_TEST("nist_key_wrap_invalid", NIST_Keywrap_Invalid_Tests);
+#endif
+
 }
 
 }