Fix Memory Leaks (#734)

awslabs · Dec 13, 2023 · a8bec86 · a8bec86
1 parent 09c8afe
commit a8bec86
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 13 deletions.
diff --git a/crt/aws-c-s3 b/crt/aws-c-s3
diff --git a/crt/aws-c-sdkutils b/crt/aws-c-sdkutils
diff --git a/crt/aws-lc b/crt/aws-lc
diff --git a/crt/s2n b/crt/s2n
diff --git a/src/native/credentials.c b/src/native/credentials.c
@@ -24,6 +24,7 @@ struct aws_credentials *aws_credentials_new_from_java_credentials(JNIEnv *env, j
     if (java_credentials == NULL) {
         return NULL;
     }
+    struct aws_credentials *credentials = NULL;
 
     jbyteArray access_key_id =
         (*env)->GetObjectField(env, java_credentials, credentials_properties.access_key_id_field_id);
@@ -36,7 +37,8 @@ struct aws_credentials *aws_credentials_new_from_java_credentials(JNIEnv *env, j
         (*env)->GetLongField(env, java_credentials, credentials_properties.expiration_field_id);
 
     if (access_key_id == NULL && secret_access_key == NULL) {
-        return aws_credentials_new_anonymous(aws_jni_get_allocator());
+        credentials = aws_credentials_new_anonymous(aws_jni_get_allocator());
+        goto done;
     }
 
     if (access_key_id == NULL || secret_access_key == NULL) {
@@ -45,11 +47,9 @@ struct aws_credentials *aws_credentials_new_from_java_credentials(JNIEnv *env, j
             env,
             "Aws_credentials_new_from_java_credentials: Both access_key_id and secret_access_key must be either null "
             "or non-null.");
-        return NULL;
+        goto done;
     }
 
-    struct aws_credentials *credentials = NULL;
-
     struct aws_byte_cursor access_key_id_cursor = aws_jni_byte_cursor_from_jbyteArray_acquire(env, access_key_id);
     struct aws_byte_cursor secret_access_key_cursor =
         aws_jni_byte_cursor_from_jbyteArray_acquire(env, secret_access_key);
@@ -73,6 +73,17 @@ struct aws_credentials *aws_credentials_new_from_java_credentials(JNIEnv *env, j
         aws_jni_byte_cursor_from_jbyteArray_release(env, session_token, session_token_cursor);
     }
 
+done:
+    /* When local references are created by a thread from C, the JVM does not clean them up promptly. */
+    if (access_key_id) {
+        (*env)->DeleteLocalRef(env, access_key_id);
+    }
+    if (secret_access_key) {
+        (*env)->DeleteLocalRef(env, secret_access_key);
+    }
+    if (session_token) {
+        (*env)->DeleteLocalRef(env, session_token);
+    }
     return credentials;
 }
 

diff --git a/src/native/s3_client.c b/src/native/s3_client.c
@@ -152,8 +152,10 @@ static int s_s3express_get_creds_java(
         /* If we can't get an environment, then the JVM is probably shutting down. Don't crash. */
         return AWS_OP_SUCCESS;
     }
+    jobject properties_object = NULL;
+    jobject original_credentials_object = NULL;
 
-    jobject properties_object = (*env)->NewObject(
+    properties_object = (*env)->NewObject(
         env,
         s3express_credentials_properties_properties.s3express_credentials_properties_class,
         s3express_credentials_properties_properties.constructor_method_id);
@@ -164,7 +166,7 @@ static int s_s3express_get_creds_java(
             "object.");
         goto done;
     }
-    jobject original_credentials_object = aws_java_credentials_from_native_new(env, original_credentials);
+    original_credentials_object = aws_java_credentials_from_native_new(env, original_credentials);
     if ((*env)->ExceptionCheck(env) || original_credentials_object == NULL) {
         aws_jni_throw_runtime_exception(
             env, "S3ExpressCredentialsProvider.getS3ExpressCredentials: Failed to create Credentials object.");
@@ -195,9 +197,6 @@ static int s_s3express_get_creds_java(
         original_credentials_object,
         (jlong)callback_data);
 
-    (*env)->DeleteLocalRef(env, properties_object);
-    (*env)->DeleteLocalRef(env, original_credentials_object);
-
     if (aws_jni_check_and_clear_exception(env)) {
         /* Check if any exception raised */
         AWS_LOGF_ERROR(
@@ -209,6 +208,12 @@ static int s_s3express_get_creds_java(
     }
     result = AWS_OP_SUCCESS;
 done:
+    if (properties_object) {
+        (*env)->DeleteLocalRef(env, properties_object);
+    }
+    if (original_credentials_object) {
+        (*env)->DeleteLocalRef(env, original_credentials_object);
+    }
     aws_jni_release_thread_env(impl->jvm, env);
     /********** JNI ENV RELEASE **********/
     return result;
@@ -227,6 +232,8 @@ static void s_s3express_destroy_java(struct aws_s3express_credentials_provider *
     (*env)->CallVoidMethod(
         env, impl->java_s3express_provider, s3express_credentials_provider_properties.destroyProvider);
 
+    (*env)->DeleteGlobalRef(env, impl->java_s3express_provider);
+
     aws_jni_release_thread_env(impl->jvm, env);
     /********** JNI ENV RELEASE **********/
     /* Once the java call returns, the java resource should be cleaned up already. We can finish up the shutdown