Merge pull request #4 from dysnix/predictkube-scaler-tls

Predictkube scaler tls Signed-off-by: alex60217101990 <alex6021710@gmail.com>
kedacore · Jan 10, 2022 · 9f374a8 · 9f374a8
2 parents 8ee4fb3 + 11dc96c
commit 9f374a8
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 37 deletions.
diff --git a/pkg/scalers/predictkube_scaler.go b/pkg/scalers/predictkube_scaler.go
@@ -2,25 +2,20 @@ package scalers
 
 import (
 	"context"
+	"crypto/tls"
 	"errors"
 	"fmt"
 	"math"
 	"strconv"
 	"time"
 
-	libs "github.com/dysnix/predictkube-libs/external/configs"
-	pc "github.com/dysnix/predictkube-libs/external/grpc/client"
-	"github.com/dysnix/predictkube-libs/external/http_transport"
-	tc "github.com/dysnix/predictkube-libs/external/types_convertation"
-	"github.com/dysnix/predictkube-proto/external/proto/commonproto"
-	pb "github.com/dysnix/predictkube-proto/external/proto/services"
-	validator "github.com/go-playground/validator/v10"
-	kedautil "github.com/kedacore/keda/v2/pkg/util"
+	"github.com/go-playground/validator/v10"
 	"github.com/prometheus/client_golang/api"
 	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/prometheus/common/model"
-	str2duration "github.com/xhit/go-str2duration/v2"
+	"github.com/xhit/go-str2duration/v2"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
 	health "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/status"
 	"k8s.io/api/autoscaling/v2beta2"
@@ -29,17 +24,51 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/metrics/pkg/apis/external_metrics"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
+
+	libs "github.com/dysnix/predictkube-libs/external/configs"
+	pc "github.com/dysnix/predictkube-libs/external/grpc/client"
+	"github.com/dysnix/predictkube-libs/external/http_transport"
+	tc "github.com/dysnix/predictkube-libs/external/types_convertation"
+	"github.com/dysnix/predictkube-proto/external/proto/commonproto"
+	pb "github.com/dysnix/predictkube-proto/external/proto/services"
+
+	kedautil "github.com/kedacore/keda/v2/pkg/util"
 )
 
 const (
 	predictKubeMetricType = "External"
 )
 
 var (
-	mlEngineHost = "predictkube-dev.dysnix.org"
-	mlEnginePort = 8080
+	mlEngineHost = "api.predictkube.com"
+	mlEnginePort = 443
 
 	defaultStep = time.Minute * 5
+
+	grpcConf = &libs.GRPC{
+		Enabled:       true,
+		UseReflection: true,
+		Compression: libs.Compression{
+			Enabled: false,
+		},
+		Conn: &libs.Connection{
+			Host:            mlEngineHost,
+			Port:            uint16(mlEnginePort),
+			ReadBufferSize:  50 << 20,
+			WriteBufferSize: 50 << 20,
+			MaxMessageSize:  50 << 20,
+			Insecure:        false,
+			Timeout:         time.Second * 15,
+		},
+		Keepalive: &libs.Keepalive{
+			Time:    time.Minute * 5,
+			Timeout: time.Minute * 5,
+			EnforcementPolicy: &libs.EnforcementPolicy{
+				MinTime:             time.Minute * 20,
+				PermitWithoutStream: false,
+			},
+		},
+	}
 )
 
 type PredictKubeScaler struct {
@@ -67,45 +96,29 @@ type predictKubeMetadata struct {
 var predictKubeLog = logf.Log.WithName("predictkube_scaler")
 
 func (pks *PredictKubeScaler) setupClientConn() error {
-	clientOpt, err := pc.SetGrpcClientOptions(&libs.GRPC{
-		Enabled:       true,
-		UseReflection: true,
-		Compression: libs.Compression{
-			Enabled: false,
-		},
-		Conn: &libs.Connection{
-			Host:            mlEngineHost,
-			Port:            uint16(mlEnginePort),
-			ReadBufferSize:  50 << 20,
-			WriteBufferSize: 50 << 20,
-			MaxMessageSize:  50 << 20,
-			Insecure:        true,
-			Timeout:         time.Second * 15,
-		},
-		Keepalive: &libs.Keepalive{
-			Time:    time.Minute * 5,
-			Timeout: time.Minute * 5,
-			EnforcementPolicy: &libs.EnforcementPolicy{
-				MinTime:             time.Minute * 20,
-				PermitWithoutStream: false,
-			},
-		},
-	},
+	clientOpt, err := pc.SetGrpcClientOptions(grpcConf,
 		&libs.Base{
 			Monitoring: libs.Monitoring{
 				Enabled: false,
 			},
 			Profiling: libs.Profiling{
 				Enabled: false,
 			},
-			IsDebugMode: true,
 			Single: &libs.Single{
 				Enabled: false,
 			},
 		},
 		pc.InjectPublicClientMetadataInterceptor(pks.metadata.apiKey),
 	)
 
+	if !grpcConf.Conn.Insecure {
+		clientOpt = append(clientOpt, grpc.WithTransportCredentials(
+			credentials.NewTLS(&tls.Config{
+				ServerName: mlEngineHost,
+			}),
+		))
+	}
+
 	if err != nil {
 		return err
 	}
@@ -460,4 +473,4 @@ func (pks *PredictKubeScaler) initPredictKubePrometheusConn() (err error) {
 	pks.api = v1.NewAPI(pks.prometheusClient)
 
 	return pks.ping(context.Background())
-}
+}
diff --git a/pkg/scalers/predictkube_scaler_test.go b/pkg/scalers/predictkube_scaler_test.go
@@ -195,6 +195,8 @@ func TestPredictKubeGetMetricSpecForScaling(t *testing.T) {
 }
 
 func TestPredictKubeGetMetrics(t *testing.T) {
+	grpcConf.Conn.Insecure = true
+
 	mockPredictServer, grpcServer := runMockGrpcPredictServer()
 	<-time.After(time.Second * 3)
 	defer func() {