Add helper function for compiling TLS arguments (#1610)

Summary:
Pull Request resolved: #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

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 "",
+    }

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"

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)