fix ssl pool

.github/workflows/run_test.yaml

@@ -33,18 +33,18 @@ jobs:
         run: |
           docker-compose up -d
           sleep 45
-          pytest -s -v -k "not TestSSLConnection and not TestSSLConnectionSelfSigned"
+          pytest -s -v -k "not SSL and not self_signed_SSL"
           docker-compose down -v
         working-directory: tests
       - name: Test SSL connection with pytest
         run: |
           enable_ssl=true docker-compose up -d
           sleep 45
-          pytest -s -v test_ssl_connection.py::TestSSLConnection
+          pytest -s -v -k "SSL"
         working-directory: tests
       - name: Test self-signed SSL connection with pytest
         run: |
           enable_ssl=true docker-compose up -d
           sleep 45
-          pytest -s -v test_ssl_connection.py::TestSSLConnectionSelfSigned
+          pytest -s -v -k "self_signed_SSL"
         working-directory: tests

nebula2/gclient/net/ConnectionPool.py

@@ -52,6 +52,7 @@ def init(self, addresses, configs, ssl_conf=None):
             logger.error('The pool has init or closed.')
             raise RuntimeError('The pool has init or closed.')
         self._configs = configs
+        self._ssl_configs = ssl_conf
         for address in addresses:
             if address not in self._addresses:
                 try:
@@ -77,7 +78,7 @@ def init(self, addresses, configs, ssl_conf=None):
 
         conns_per_address = int(self._configs.min_connection_pool_size / ok_num)
 
-        if ssl_conf is None:
+        if self._ssl_configs is None:
             for addr in self._addresses:
                 for i in range(0, conns_per_address):
                     connection = Connection()

tests/docker-compose.yaml

@@ -397,7 +397,7 @@ services:
       - sh
       - -c
       - |
-        sleep 3 &&
+        sleep 5 &&
         nebula-console -addr graphd0 -port 9669 -u root -p nebula -e 'ADD HOSTS "172.28.2.1":9779,"172.28.2.2":9779,"172.28.2.3":9779' &&
         sleep 36000
     depends_on:

tests/test_ssl_connection.py

@@ -9,6 +9,7 @@
 import os
 import time
 import ssl
+import pytest
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 root_dir = os.path.join(current_dir, '..')
@@ -45,6 +46,7 @@
 port = 9669
 
 
+@pytest.mark.SSL
 class TestSSLConnection(TestCase):
     def test_create(self):
         try:
@@ -86,6 +88,7 @@ def test_close(self):
             assert True
 
 
+@pytest.mark.self_signed_SSL
 class TestSSLConnectionSelfSigned(TestCase):
     def test_create_self_signed(self):
         try:

tests/test_ssl_pool.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python
+# --coding:utf-8--
+
+# Copyright (c) 2021 vesoft inc. All rights reserved.
+#
+# This source code is licensed under Apache 2.0 License.
+import sys
+import os
+import ssl
+import copy
+import pytest
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+root_dir = os.path.join(current_dir, '..')
+sys.path.insert(0, root_dir)
+
+from unittest import TestCase
+
+from nebula2.gclient.net import ConnectionPool
+from nebula2.Config import Config, SSL_config
+
+
+@pytest.mark.SSL
+class TestConnectionPool(TestCase):
+    @classmethod
+    def setup_class(self):
+        self.addresses = list()
+        self.addresses.append(('127.0.0.1', 9669))
+        self.configs = Config()
+        self.configs.min_connection_pool_size = 2
+        self.configs.max_connection_pool_size = 4
+        self.configs.idle_time = 2000
+        self.configs.interval_check = 2
+
+        # set SSL config
+        self.ssl_config = SSL_config()
+        self.ssl_config.cert_reqs = ssl.CERT_OPTIONAL
+        self.ssl_config.ca_certs = os.path.join(current_dir, 'secrets/test.ca.pem')
+        self.ssl_config.keyfile = os.path.join(current_dir, 'secrets/test.client.key')
+        self.ssl_config.certfile = os.path.join(current_dir, 'secrets/test.client.crt')
+        # self signed SSL config
+        self.ssl_selfs_signed_config = SSL_config()
+        self.ssl_selfs_signed_config.cert_reqs = ssl.CERT_OPTIONAL
+        self.ssl_selfs_signed_config.cert_reqs = ssl.CERT_OPTIONAL
+        self.ssl_selfs_signed_config.ca_certs = os.path.join(
+            current_dir, 'secrets/test.self-signed.pem'
+        )
+        self.ssl_selfs_signed_config.keyfile = os.path.join(
+            current_dir, 'secrets/test.self-signed.key'
+        )
+        self.ssl_selfs_signed_config.certfile = os.path.join(
+            current_dir, 'secrets/test.self-signed.pem'
+        )
+
+    def test_ssl_with_ca(self):
+        pool = ConnectionPool()
+        assert pool.init(self.addresses, self.configs, self.ssl_config)
+        session = pool.get_session("root", "nebula")
+        resp = session.execute("SHOW HOSTS")
+        assert resp.is_succeeded()
+
+    def test_ssl_with_invalid_ca(self):
+        pool = ConnectionPool()
+        config = copy.copy(self.ssl_config)
+        config.ca_certs = "invalid"
+
+        with self.assertRaises(Exception):
+            pool.init(self.addresses, self.configs, config)