From 402e87d2f9ca094d2a05a489dde38be3b7fadb16 Mon Sep 17 00:00:00 2001
From: Amogh-Bharadwaj <amogh@peerdb.io>
Date: Tue, 28 Nov 2023 00:17:58 +0530
Subject: [PATCH] refactor cert handling

---
 flow/cmd/api.go             | 17 +++--------------
 flow/cmd/cert.go            | 29 +++++++++++++++++++++++++++++
 flow/cmd/snapshot_worker.go | 18 +++---------------
 flow/cmd/worker.go          | 18 +++---------------
 4 files changed, 38 insertions(+), 44 deletions(-)
 create mode 100644 flow/cmd/cert.go

diff --git a/flow/cmd/api.go b/flow/cmd/api.go
index 21ad34f27b..9675163b27 100644
--- a/flow/cmd/api.go
+++ b/flow/cmd/api.go
@@ -3,11 +3,9 @@ package main
 import (
 	"context"
 	"crypto/tls"
-	"encoding/base64"
 	"fmt"
 	"net"
 	"net/http"
-	"strings"
 	"time"
 
 	utils "github.com/PeerDB-io/peer-flow/connectors/utils/catalog"
@@ -93,23 +91,14 @@ func APIMain(args *APIServerParams) error {
 	}
 	if args.TemporalCert != "" && args.TemporalKey != "" {
 		log.Info("Using temporal certificate/key for authentication")
-		certBytes, err := base64.StdEncoding.DecodeString(strings.TrimSpace(args.TemporalCert))
-		if err != nil {
-			return fmt.Errorf("unable to decode temporal certificate: %w", err)
-		}
-
-		keyBytes, err := base64.StdEncoding.DecodeString(strings.TrimSpace(args.TemporalKey))
-		if err != nil {
-			return fmt.Errorf("unable to decode temporal key: %w", err)
-		}
 
-		cert, err := tls.X509KeyPair(certBytes, keyBytes)
+		certs, err := ProcessCertAndKey(args.TemporalCert, args.TemporalKey)
 		if err != nil {
-			return fmt.Errorf("unable to obtain temporal key pair: %w", err)
+			return fmt.Errorf("unable to process certificate and key: %w", err)
 		}
 
 		connOptions := client.ConnectionOptions{
-			TLS: &tls.Config{Certificates: []tls.Certificate{cert}},
+			TLS: &tls.Config{Certificates: certs},
 		}
 		clientOptions.ConnectionOptions = connOptions
 	}
diff --git a/flow/cmd/cert.go b/flow/cmd/cert.go
new file mode 100644
index 0000000000..1fe45a4e99
--- /dev/null
+++ b/flow/cmd/cert.go
@@ -0,0 +1,29 @@
+package main
+
+import (
+	"crypto/tls"
+	"encoding/base64"
+	"fmt"
+	"strings"
+)
+
+func ProcessCertAndKey(cert string, key string) ([]tls.Certificate, error) {
+	temporalCert := strings.TrimSpace(cert)
+	certBytes, err := base64.StdEncoding.DecodeString(temporalCert)
+	if err != nil {
+		return nil, fmt.Errorf("unable to decode temporal certificate: %w", err)
+	}
+
+	temporalKey := strings.TrimSpace(key)
+	keyBytes, err := base64.StdEncoding.DecodeString(temporalKey)
+	if err != nil {
+		return nil, fmt.Errorf("unable to decode temporal key: %w", err)
+	}
+
+	keyPair, err := tls.X509KeyPair(certBytes, keyBytes)
+	if err != nil {
+		return nil, fmt.Errorf("unable to obtain temporal key pair: %w", err)
+	}
+
+	return []tls.Certificate{keyPair}, nil
+}
diff --git a/flow/cmd/snapshot_worker.go b/flow/cmd/snapshot_worker.go
index ebd7609c05..a0d88ad080 100644
--- a/flow/cmd/snapshot_worker.go
+++ b/flow/cmd/snapshot_worker.go
@@ -2,9 +2,7 @@ package main
 
 import (
 	"crypto/tls"
-	"encoding/base64"
 	"fmt"
-	"strings"
 
 	"github.com/PeerDB-io/peer-flow/activities"
 	"github.com/PeerDB-io/peer-flow/shared"
@@ -28,23 +26,13 @@ func SnapshotWorkerMain(opts *SnapshotWorkerOptions) error {
 	}
 
 	if opts.TemporalCert != "" && opts.TemporalKey != "" {
-		certBytes, err := base64.StdEncoding.DecodeString(strings.TrimSpace(opts.TemporalCert))
+		certs, err := ProcessCertAndKey(opts.TemporalCert, opts.TemporalKey)
 		if err != nil {
-			return fmt.Errorf("unable to decode temporal certificate: %w", err)
-		}
-
-		keyBytes, err := base64.StdEncoding.DecodeString(strings.TrimSpace(opts.TemporalKey))
-		if err != nil {
-			return fmt.Errorf("unable to decode temporal key: %w", err)
-		}
-
-		cert, err := tls.X509KeyPair(certBytes, keyBytes)
-		if err != nil {
-			return fmt.Errorf("unable to obtain temporal key pair: %w", err)
+			return fmt.Errorf("unable to process certificate and key: %w", err)
 		}
 
 		connOptions := client.ConnectionOptions{
-			TLS: &tls.Config{Certificates: []tls.Certificate{cert}},
+			TLS: &tls.Config{Certificates: certs},
 		}
 		clientOptions.ConnectionOptions = connOptions
 	}
diff --git a/flow/cmd/worker.go b/flow/cmd/worker.go
index da6c03f122..5fe41a4a1a 100644
--- a/flow/cmd/worker.go
+++ b/flow/cmd/worker.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"crypto/tls"
-	"encoding/base64"
 	"fmt"
 	"os"
 	"os/signal"
@@ -95,23 +94,12 @@ func WorkerMain(opts *WorkerOptions) error {
 
 	if opts.TemporalCert != "" && opts.TemporalKey != "" {
 		log.Info("Using temporal certificate/key for authentication")
-		certBytes, err := base64.StdEncoding.DecodeString(opts.TemporalCert)
+		certs, err := ProcessCertAndKey(opts.TemporalCert, opts.TemporalKey)
 		if err != nil {
-			return fmt.Errorf("unable to decode temporal certificate: %w", err)
+			return fmt.Errorf("unable to process certificate and key: %w", err)
 		}
-
-		keyBytes, err := base64.StdEncoding.DecodeString(opts.TemporalKey)
-		if err != nil {
-			return fmt.Errorf("unable to decode temporal key: %w", err)
-		}
-
-		cert, err := tls.X509KeyPair(certBytes, keyBytes)
-		if err != nil {
-			return fmt.Errorf("unable to obtain temporal key pair: %w", err)
-		}
-
 		connOptions := client.ConnectionOptions{
-			TLS: &tls.Config{Certificates: []tls.Certificate{cert}},
+			TLS: &tls.Config{Certificates: certs},
 		}
 		clientOptions.ConnectionOptions = connOptions
 	}