From 1e33edeefd8c15ee7b2f8ef0195963abb3cacdc9 Mon Sep 17 00:00:00 2001
From: Siddhanta Rath <siddhanta.rath@gojek.com>
Date: Mon, 18 Dec 2023 20:57:37 +0530
Subject: [PATCH] feat: enable tls on grpc server

---
 .env.sample              |  1 +
 .env.test                |  1 +
 config/server.go         | 14 +++++++++++---
 services/grpc/service.go | 25 ++++++++++++++++++++++++-
 4 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/.env.sample b/.env.sample
index 004016e0..f36ab4c8 100644
--- a/.env.sample
+++ b/.env.sample
@@ -10,6 +10,7 @@ SERVER_WEBSOCKET_WRITE_WAIT_INTERVAL_MS=5000
 SERVER_WEBSOCKET_PINGER_SIZE=1
 
 SERVER_GRPC_PORT=8081
+SERVER_GRPC_TLS_ENABLED=false
 
 WORKER_BUFFER_CHANNEL_SIZE=5
 WORKER_BUFFER_FLUSH_TIMEOUT_MS=5000
diff --git a/.env.test b/.env.test
index 4d6aad85..1136269e 100644
--- a/.env.test
+++ b/.env.test
@@ -11,6 +11,7 @@ SERVER_WEBSOCKET_WRITE_WAIT_INTERVAL_MS=1000
 SERVER_WEBSOCKET_PINGER_SIZE=1
 
 SERVER_GRPC_PORT=8081
+SERVER_GRPC_TLS_ENABLED=false
 
 WORKER_BUFFER_CHANNEL_SIZE=5
 WORKER_BUFFER_FLUSH_TIMEOUT_MS=5000
diff --git a/config/server.go b/config/server.go
index ace3fdd5..f0ece0e5 100644
--- a/config/server.go
+++ b/config/server.go
@@ -31,7 +31,10 @@ type serverWs struct {
 }
 
 type serverGRPC struct {
-	Port string
+	Port         string
+	TLSEnabled   bool
+	TLSCertPath  string
+	TLSPublicKey string
 }
 
 func serverConfigLoader() {
@@ -71,9 +74,14 @@ func serverWsConfigLoader() {
 }
 
 func serverGRPCConfigLoader() {
-
 	viper.SetDefault("SERVER_GRPC_PORT", "8081")
+	viper.SetDefault("SERVER_GRPC_TLS_ENABLED", false)
+	viper.SetDefault("SERVER_GRPC_TLS_CERT_PATH", "cert/server.crt")
+	viper.SetDefault("SERVER_GRPC_TLS_PUBLIC_KEY", "cert/server.key")
 	ServerGRPC = serverGRPC{
-		Port: util.MustGetString("SERVER_GRPC_PORT"),
+		Port:         util.MustGetString("SERVER_GRPC_PORT"),
+		TLSEnabled:   util.MustGetBool("SERVER_GRPC_TLS_ENABLED"),
+		TLSCertPath:  util.MustGetString("SERVER_GRPC_TLS_CERT_PATH"),
+		TLSPublicKey: util.MustGetString("SERVER_GRPC_TLS_PUBLIC_KEY"),
 	}
 }
diff --git a/services/grpc/service.go b/services/grpc/service.go
index 3b4a018b..fc7abfd9 100644
--- a/services/grpc/service.go
+++ b/services/grpc/service.go
@@ -2,7 +2,9 @@ package grpc
 
 import (
 	"context"
+	"crypto/tls"
 	"fmt"
+	"google.golang.org/grpc/credentials"
 	"net"
 
 	pbgrpc "buf.build/gen/go/gotocompany/proton/grpc/go/gotocompany/raccoon/v1beta1/raccoonv1beta1grpc"
@@ -17,7 +19,7 @@ type Service struct {
 }
 
 func NewGRPCService(c collection.Collector) *Service {
-	server := grpc.NewServer()
+	server := newGRPCServer()
 	pbgrpc.RegisterEventServiceServer(server, &Handler{C: c})
 	return &Service{
 		s:         server,
@@ -41,3 +43,24 @@ func (s *Service) Shutdown(context.Context) error {
 	s.s.GracefulStop()
 	return nil
 }
+
+func newGRPCServer() *grpc.Server {
+	if config.ServerGRPC.TLSEnabled {
+		return grpc.NewServer(grpc.Creds(loadTLSCredentials()))
+	}
+	return grpc.NewServer()
+}
+
+func loadTLSCredentials() credentials.TransportCredentials {
+	serverCert, err := tls.LoadX509KeyPair(config.ServerGRPC.TLSCertPath, config.ServerGRPC.TLSPublicKey)
+	if err != nil {
+		panic("failed to load TLS credentials to start grpc server with TLS")
+	}
+
+	config := &tls.Config{
+		Certificates: []tls.Certificate{serverCert},
+		ClientAuth:   tls.NoClientCert,
+	}
+
+	return credentials.NewTLS(config)
+}