Sanitize input

awslabs · Nov 1, 2023 · 0fa9a80 · 0fa9a80
1 parent 221f449
commit 0fa9a80
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 8 deletions.
diff --git a/src/source/Crypto/Dtls_mbedtls.c b/src/source/Crypto/Dtls_mbedtls.c
@@ -10,9 +10,13 @@ mbedtls_ssl_srtp_profile DTLS_SRTP_SUPPORTED_PROFILES[] = {
 
 STATUS md5DigestCalculation(PBYTE inputStringBuff, UINT64 length, PBYTE outputBuff)
 {
-    mbedtls_md5_ret(inputStringBuff, length, outputBuff);
-    return STATUS_SUCCESS;
+    STATUS retStatus = STATUS_SUCCESS;
+    CHK_ERR(inputStringBuff != NULL && outputBuff != NULL, STATUS_INVALID_ARG, "Invalid input or output buffer");
+    CHK_ERR(!mbedtls_md5_ret(inputStringBuff, length, outputBuff), STATUS_INTERNAL_ERROR, "MD5 calculation failed");
+CleanUp:
+    return retStatus;
 }
+
 STATUS createDtlsSession(PDtlsSessionCallbacks pDtlsSessionCallbacks, TIMER_QUEUE_HANDLE timerQueueHandle, INT32 certificateBits,
                          BOOL generateRSACertificate, PRtcCertificate pRtcCertificates, PDtlsSession* ppDtlsSession)
 {

diff --git a/src/source/Crypto/Dtls_openssl.c b/src/source/Crypto/Dtls_openssl.c
@@ -50,8 +50,14 @@ STATUS md5DigestCalculation(PBYTE inputStringBuff, UINT64 length, PBYTE outputBu
 #if (OPENSSL_VERSION_NUMBER >= 0x30000000L)
     EVP_MD_CTX* mdctx = NULL;
     const EVP_MD* md = NULL;
+#endif
+
+    CHK_ERR(inputStringBuff != NULL && outputBuff != NULL, STATUS_INVALID_ARG, "Invalid input or output buffer");
+
+#if (OPENSSL_VERSION_NUMBER >= 0x30000000L)
+
     CHK_ERR(md = EVP_MD_fetch(NULL, "MD5", NULL), STATUS_INTERNAL_ERROR, "Failed to fetch MD5 provider");
-    mdctx = EVP_MD_CTX_new();
+    CHK_ERR(mdctx = EVP_MD_CTX_new(), STATUS_INTERNAL_ERROR, "Failed to create message digest context");
     CHK_ERR(EVP_DigestInit_ex(mdctx, md, NULL), STATUS_INTERNAL_ERROR, "Message digest initialization failed.");
     CHK_ERR(EVP_DigestUpdate(mdctx, inputStringBuff, length), STATUS_INTERNAL_ERROR, "Message digest update failed");
     CHK_ERR(EVP_DigestFinal_ex(mdctx, outputBuff, NULL), STATUS_INTERNAL_ERROR, "Message digest finalization failed");
@@ -61,11 +67,12 @@ STATUS md5DigestCalculation(PBYTE inputStringBuff, UINT64 length, PBYTE outputBu
 
 CleanUp:
 #if (OPENSSL_VERSION_NUMBER >= 0x30000000L)
-    EVP_MD_CTX_free(mdctx);
-    EVP_MD_free((EVP_MD*) md);
-// Adding else to get around Mac unused label error
-#else
-    retStatus = STATUS_SUCCESS;
+    if (mdctx != NULL) {
+        EVP_MD_CTX_free(mdctx);
+    }
+    if (md != NULL) {
+        EVP_MD_free((EVP_MD*) md);
+    }
 #endif
     return retStatus;
 }