Support SSL connection

.github/workflows/test.yaml

@@ -16,12 +16,25 @@ jobs:
       - name: Test with GO test
         run: |
           pushd nebula-docker-compose/
-          docker-compose up -d
+          enable_ssl=false docker-compose up -d
           sleep 10
           popd
           go test -v -race
+      - name: Test SSL connection with CA
+        run: |
+          pushd nebula-docker-compose/
+          enable_ssl=true docker-compose up -d
+          sleep 10
+          popd
+          ssl_test=true go test -v -run TestSslConnection
+      - name: Test SSL connection self-signed
+        run: |
+          pushd nebula-docker-compose/
+          enable_ssl=true docker-compose up -d
+          sleep 10
+          popd
+          self_signed=true go test -v -run TestSslConnectionSelfSigned
       - name: Run examples
         run: |
           go run basic_example/graph_client_basic_example.go
           go run gorountines_example/graph_client_goroutines_example.go
-

Makefile

@@ -1,4 +1,4 @@
-.PHONY: build test fmt ci run-examples
+.PHONY: build test fmt ci ssl-test run-examples
 
 default: build
 
@@ -18,6 +18,20 @@ ci:
 	sleep 5 && \
 	cd .. && \
 	go test -v -race; \
+	cd ./nebula-docker-compose && docker-compose down -v
+
+ssl-test:
+	cd ./nebula-docker-compose && enable_ssl=true docker-compose up -d && \
+	sleep 5 && \
+	cd .. && \
+	go test -v -run TestSslConnection; \
+	cd ./nebula-docker-compose && docker-compose down -v 
+
+ssl-test-self-signed:
+	cd ./nebula-docker-compose && enable_ssl=true docker-compose up -d && \
+	sleep 5 && \
+	cd .. && \
+	ssl_test=true go test -v -run TestSslConnection; \
 	cd ./nebula-docker-compose && docker-compose down -v 
 
 run-examples:

connection.go

@@ -7,6 +7,7 @@
 package nebula_go
 
 import (
+	"crypto/tls"
 	"fmt"
 	"math"
 	"time"
@@ -59,6 +60,33 @@ func (cn *connection) open(hostAddress HostAddress, timeout time.Duration) error
 	return nil
 }
 
+func (cn *connection) openSSL(hostAddress HostAddress, timeout time.Duration, sslConfig *tls.Config) error {
+	ip := hostAddress.Host
+	port := hostAddress.Port
+	newAdd := fmt.Sprintf("%s:%d", ip, port)
+	cn.timeout = timeout
+	bufferSize := 128 << 10
+	frameMaxLength := uint32(math.MaxUint32)
+
+	SSLSocket, err := thrift.NewSSLSocketTimeout(newAdd, sslConfig, timeout)
+	if err != nil {
+		return fmt.Errorf("failed to create a net.Conn-backed Transport,: %s", err.Error())
+	}
+
+	// Set transport buffer
+	bufferedTranFactory := thrift.NewBufferedTransportFactory(bufferSize)
+	transport := thrift.NewFramedTransportMaxLength(bufferedTranFactory.GetTransport(SSLSocket), frameMaxLength)
+	pf := thrift.NewBinaryProtocolFactoryDefault()
+	cn.graph = graph.NewGraphServiceClientFactory(transport, pf)
+	if err = cn.graph.Open(); err != nil {
+		return fmt.Errorf("failed to open transport, error: %s", err.Error())
+	}
+	if !cn.graph.IsOpen() {
+		return fmt.Errorf("transport is off")
+	}
+	return nil
+}
+
 // reopen reopens the current connection.
 // Because the code generated by Fbthrift does not handle the seqID,
 // the message will be dislocated when the timeout occurs, resulting in unexpected response.

connection_pool.go

@@ -8,6 +8,8 @@ package nebula_go
 
 import (
 	"container/list"
+	"crypto/tls"
+	"errors"
 	"fmt"
 	"sync"
 	"time"
@@ -25,6 +27,7 @@ type ConnectionPool struct {
 	rwLock                sync.RWMutex
 	cleanerChan           chan struct{} //notify when pool is close
 	closed                bool
+	sslConfig             *tls.Config
 }
 
 // NewConnectionPool constructs a new connection pool using the given addresses and configs
@@ -48,22 +51,62 @@ func NewConnectionPool(addresses []HostAddress, conf PoolConfig, log Logger) (*C
 		log:       log,
 		addresses: convAddress,
 		hostIndex: 0,
+		sslConfig: nil,
 	}
+
+	// Init pool with non-SSL socket
+	if err = newPool.initPool(); err != nil {
+		return nil, err
+	}
+	newPool.startCleaner()
+	return newPool, nil
+}
+
+// NewConnectionPool constructs a new SSL connection pool using the given addresses and configs
+func NewSslConnectionPool(addresses []HostAddress, conf PoolConfig, sslConfig *tls.Config, log Logger) (*ConnectionPool, error) {
+	// Process domain to IP
+	convAddress, err := DomainToIP(addresses)
+	if err != nil {
+		return nil, fmt.Errorf("failed to find IP, error: %s ", err.Error())
+	}
+
+	// Check input
+	if len(convAddress) == 0 {
+		return nil, fmt.Errorf("failed to initialize connection pool: illegal address input")
+	}
+
+	// Check config
+	conf.validateConf(log)
+
+	newPool := &ConnectionPool{
+		conf:      conf,
+		log:       log,
+		addresses: convAddress,
+		hostIndex: 0,
+		sslConfig: sslConfig,
+	}
+
+	// Init pool with SSL socket
 	if err = newPool.initPool(); err != nil {
 		return nil, err
 	}
 	newPool.startCleaner()
 	return newPool, nil
 }
 
-// initPool innitializes the connection pool
+// initPool initializes the connection pool
 func (pool *ConnectionPool) initPool() error {
 	for i := 0; i < pool.conf.MinConnPoolSize; i++ {
 		// Simple round-robin
 		newConn := newConnection(pool.addresses[i%len(pool.addresses)])
 
 		// Open connection to host
-		err := newConn.open(newConn.severAddress, pool.conf.TimeOut)
+		err := errors.New("")
+		if pool.sslConfig == nil {
+			err = newConn.open(newConn.severAddress, pool.conf.TimeOut)
+		} else {
+			err = newConn.openSSL(newConn.severAddress, pool.conf.TimeOut, pool.sslConfig)
+		}
 		if err != nil {
 			// If initialization failed, clean idle queue
 			idleLen := pool.idleConnectionQueue.Len()
@@ -167,8 +210,14 @@ func (pool *ConnectionPool) release(conn *connection) {
 func (pool *ConnectionPool) Ping(host HostAddress, timeout time.Duration) error {
 	newConn := newConnection(host)
 	// Open connection to host
-	if err := newConn.open(newConn.severAddress, timeout); err != nil {
-		return err
+	if pool.sslConfig == nil {
+		if err := newConn.open(newConn.severAddress, timeout); err != nil {
+			return err
+		}
+	} else {
+		if err := newConn.openSSL(newConn.severAddress, timeout, pool.sslConfig); err != nil {
+			return err
+		}
 	}
 	newConn.close()
 	return nil
@@ -220,9 +269,14 @@ func (pool *ConnectionPool) newConnToHost() (*connection, error) {
 	host := pool.getHost()
 	newConn := newConnection(host)
 	// Open connection to host
-	err := newConn.open(newConn.severAddress, pool.conf.TimeOut)
-	if err != nil {
-		return nil, err
+	if pool.sslConfig == nil {
+		if err := newConn.open(newConn.severAddress, pool.conf.TimeOut); err != nil {
+			return nil, err
+		}
+	} else {
+		if err := newConn.openSSL(newConn.severAddress, pool.conf.TimeOut, pool.sslConfig); err != nil {
+			return nil, err
+		}
 	}
 	// Add connection to active queue
 	pool.activeConnectionQueue.PushBack(newConn)

nebula-docker-compose/.env

@@ -0,0 +1,5 @@
+enable_ssl=false
+ca_path=/secrets/test.ca.pem
+password_path=/secrets/test.ca.password
+cert_path=/secrets/test.client.crt
+key_path=/secrets/test.client.key