From 4cbf728910a816a6d03e31ab701eef439752f485 Mon Sep 17 00:00:00 2001
From: "Zheng, Qi" <huaiqing.zq@antgroup.com>
Date: Mon, 17 Apr 2023 19:35:19 +0800
Subject: [PATCH] [toolchain] Add get key to buffer API for grpc_ratls

---
 .../examples/cpp/ratls/grpc_ratls_client.cc   | 82 ++++++++++++++++---
 .../examples/cpp/ratls/grpc_ratls_client.h    |  8 ++
 2 files changed, 78 insertions(+), 12 deletions(-)

diff --git a/tools/toolchains/grpc_ratls/ra_tls/grpc/v1.38.1/examples/cpp/ratls/grpc_ratls_client.cc b/tools/toolchains/grpc_ratls/ra_tls/grpc/v1.38.1/examples/cpp/ratls/grpc_ratls_client.cc
index f0d264e1..1c2fd75a 100644
--- a/tools/toolchains/grpc_ratls/ra_tls/grpc/v1.38.1/examples/cpp/ratls/grpc_ratls_client.cc
+++ b/tools/toolchains/grpc_ratls/ra_tls/grpc/v1.38.1/examples/cpp/ratls/grpc_ratls_client.cc
@@ -35,6 +35,15 @@ using ratls::GrSecret;
 using ratls::SecretRequest;
 using ratls::SecretReply;
 
+typedef enum {
+    GRPC_RATLS_SUCCESS = 0,             /// Success
+    GRPC_RATLS_ERR = -1,                /// General error
+    GRPC_RATLS_INVALID_PARAM = -2,      /// Invalid parameter
+    GRPC_RATLS_BUF_ERR = -3,            /// Invalid buffer or buffer allocation failure
+    GRPC_RATLS_NO_SECRET = -4,          /// No valid secret
+    GRPC_RATLS_BUF_TOO_SMALL = -5       /// Buffer is too small
+} grpc_ratls_result_t;
+
 // Client
 class GrSecretClient {
     public:
@@ -115,11 +124,11 @@ void base64_decode(const char *b64input, unsigned char *dest, size_t dest_len) {
     }
 }
 
-int grpc_ratls_get_secret(
+static grpc_ratls_result_t grpc_ratls_get_secret_string(
     const char *server_addr,
     const char *config_json,
     const char *name,
-    const char *secret_file
+    std::string* secret_string
 )
 {
     auto cred = grpc::sgx::TlsCredentials(config_json);
@@ -131,24 +140,73 @@ int grpc_ratls_get_secret(
     // std::cout << "secret received: " << secret << "len: " << secret.length() << std::endl;
 
     if (secret.empty()) {
-        return -1;
+        return GRPC_RATLS_NO_SECRET;
     } else {
         //Decode From Base64
         size_t len = base64_decode_len(secret.c_str());
         if (len) {
             char *secret_orig = (char *)malloc(len);
+            if (!secret_orig) {
+                return GRPC_RATLS_BUF_ERR;
+            }
             base64_decode(secret.c_str(), (unsigned char *)secret_orig, len);
-            std::string secret_string(secret_orig, secret_orig + len - 1);
+            secret_string->assign(secret_orig, secret_orig + len - 1);
+            free(secret_orig);
 
-            //write to file
-            std::ofstream myfile;
-            myfile.open(secret_file);
-            myfile << secret_string;
-            myfile.close();
-
-            return 0;
+            return GRPC_RATLS_SUCCESS;
         }
 
-        return -2;
+        return GRPC_RATLS_ERR;
     }
 }
+
+// Get secret to file
+int grpc_ratls_get_secret(
+    const char *server_addr,
+    const char *config_json,
+    const char *name,
+    const char *secret_file
+)
+{
+    std::string secret_string;
+    grpc_ratls_result_t ret = grpc_ratls_get_secret_string(
+        server_addr, config_json, name, &secret_string
+    );
+
+    if (ret == GRPC_RATLS_SUCCESS) {
+        //write to file
+        std::ofstream myfile;
+        myfile.open(secret_file);
+        myfile << secret_string;
+        myfile.close();
+    }
+
+    return ret;
+}
+
+// Get secret to buffer
+int grpc_ratls_get_secret_to_buf(
+    const char *server_addr,
+    const char *config_json,
+    const char *name,
+    char *secret_buf,
+    unsigned int *buf_len
+)
+{
+    std::string secret_string;
+    grpc_ratls_result_t ret = grpc_ratls_get_secret_string(
+        server_addr, config_json, name, &secret_string
+    );
+
+    if (ret == GRPC_RATLS_SUCCESS) {
+        if (*buf_len < secret_string.size()) {
+            std::cout << "buffer size is smaller than the secret string length " << secret_string.size() << std::endl;;
+            return GRPC_RATLS_BUF_TOO_SMALL;
+        }
+        //write to buffer
+        memcpy(secret_buf, secret_string.data(), secret_string.size());
+        *buf_len = secret_string.size();
+    }
+
+    return ret;
+}
diff --git a/tools/toolchains/grpc_ratls/ra_tls/grpc/v1.38.1/examples/cpp/ratls/grpc_ratls_client.h b/tools/toolchains/grpc_ratls/ra_tls/grpc/v1.38.1/examples/cpp/ratls/grpc_ratls_client.h
index 24a64a2a..10b39c74 100644
--- a/tools/toolchains/grpc_ratls/ra_tls/grpc/v1.38.1/examples/cpp/ratls/grpc_ratls_client.h
+++ b/tools/toolchains/grpc_ratls/ra_tls/grpc/v1.38.1/examples/cpp/ratls/grpc_ratls_client.h
@@ -13,6 +13,14 @@ extern int grpc_ratls_get_secret(
     const char *secret_file // secret file to be saved
 );
 
+extern int grpc_ratls_get_secret_to_buf(
+    const char *server_addr, // grpc server address+port, such as "localhost:50051"
+    const char *config_json, // ratls handshake config json file
+    const char *name, // secret name to be requested
+    char *secret_buf, // buffer to save secret
+    unsigned int *buf_len // buffer size
+);
+
 #ifdef __cplusplus
 }
 #endif