From af31a6d72c738e6d384a097a73922501ea5b689a Mon Sep 17 00:00:00 2001
From: Aditya Shastri <adshastri@fb.com>
Date: Wed, 21 Sep 2022 15:27:01 -0700
Subject: [PATCH] Add helper function for compiling TLS arguments (#1610)

Summary:
Pull Request resolved: https://github.com/facebookresearch/fbpcs/pull/1610

All PCF2 stage services will need to put together TLS arguments. We define a simple helper function that can be reused by all stage services

input - if the TLS feature is enabled
output - the argument dict

Reviewed By: danbunnell

Differential Revision: D39636352

fbshipit-source-id: 8f601a67e47259ee6d7fecf58e3050151490b298
---
 .../service/argument_helper.py                | 24 ++++++++++++++
 .../private_computation/service/constants.py  |  4 +++
 .../test/service/test_argument_helpers.py     | 32 +++++++++++++++++++
 3 files changed, 60 insertions(+)
 create mode 100644 fbpcs/private_computation/service/argument_helper.py
 create mode 100644 fbpcs/private_computation/test/service/test_argument_helpers.py

diff --git a/fbpcs/private_computation/service/argument_helper.py b/fbpcs/private_computation/service/argument_helper.py
new file mode 100644
index 000000000..24f45711f
--- /dev/null
+++ b/fbpcs/private_computation/service/argument_helper.py
@@ -0,0 +1,24 @@
+#!/usr/bin/env fbpython
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from typing import Any, Dict
+
+from fbpcs.private_computation.service.constants import (
+    CA_CERT_PATH,
+    PRIVATE_KEY_PATH,
+    SERVER_CERT_PATH,
+)
+
+
+def get_tls_arguments(has_tls_feature: bool) -> Dict[str, Any]:
+    return {
+        "use_tls": has_tls_feature,
+        "ca_cert_path": CA_CERT_PATH if has_tls_feature else "",
+        "server_cert_path": SERVER_CERT_PATH if has_tls_feature else "",
+        "private_key_path": PRIVATE_KEY_PATH if has_tls_feature else "",
+    }
diff --git a/fbpcs/private_computation/service/constants.py b/fbpcs/private_computation/service/constants.py
index 7cb7ff511..03a16dfb6 100644
--- a/fbpcs/private_computation/service/constants.py
+++ b/fbpcs/private_computation/service/constants.py
@@ -38,3 +38,7 @@
 DEFAULT_SORT_STRATEGY = "sort"
 DEFAULT_MULTIKEY_PROTOCOL_MAX_COLUMN_COUNT = 6
 FBPCS_BUNDLE_ID = "FBPCS_BUNDLE_ID"
+
+CA_CERT_PATH = "tls/ca_cert.pem"
+SERVER_CERT_PATH = "tls/server_cert.pem"
+PRIVATE_KEY_PATH = "tls/private_key.pem"
diff --git a/fbpcs/private_computation/test/service/test_argument_helpers.py b/fbpcs/private_computation/test/service/test_argument_helpers.py
new file mode 100644
index 000000000..ae62ef039
--- /dev/null
+++ b/fbpcs/private_computation/test/service/test_argument_helpers.py
@@ -0,0 +1,32 @@
+#!/usr/bin/env fbpython
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from unittest import TestCase
+
+from fbpcs.private_computation.service.argument_helper import get_tls_arguments
+from fbpcs.private_computation.service.constants import (
+    CA_CERT_PATH,
+    PRIVATE_KEY_PATH,
+    SERVER_CERT_PATH,
+)
+
+
+class TestArgumentHelpers(TestCase):
+    def test_get_tls_arguments_no_feature(self):
+        tls_arguments = get_tls_arguments(False)
+
+        self.assertFalse(tls_arguments["use_tls"])
+        self.assertEqual(tls_arguments["ca_cert_path"], "")
+        self.assertEqual(tls_arguments["server_cert_path"], "")
+        self.assertEqual(tls_arguments["private_key_path"], "")
+
+    def test_get_tls_arguments_with_feature(self):
+        tls_arguments = get_tls_arguments(True)
+
+        self.assertTrue(tls_arguments["use_tls"])
+        self.assertEqual(tls_arguments["ca_cert_path"], CA_CERT_PATH)
+        self.assertEqual(tls_arguments["server_cert_path"], SERVER_CERT_PATH)
+        self.assertEqual(tls_arguments["private_key_path"], PRIVATE_KEY_PATH)