From 67716a883d925b1d38ed494395dee79aeb782092 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torjus=20H=C3=A5kestad?= <torjus@usit.uio.no>
Date: Fri, 27 Oct 2023 21:46:09 +0200
Subject: [PATCH] Improve NATS

---
 go.mod           |  1 +
 mf-auth.toml     |  9 ++++++---
 server/config.go | 13 ++++++++++---
 server/server.go | 35 +++++++++++++++++++++++++++--------
 4 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/go.mod b/go.mod
index ef94242..87fff6f 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/golang-jwt/jwt/v5 v5.0.0
 	github.com/google/uuid v1.3.1
 	github.com/nats-io/nats.go v1.31.0
+	github.com/nats-io/nkeys v0.4.5
 	github.com/pelletier/go-toml/v2 v2.1.0
 	github.com/urfave/cli/v2 v2.25.7
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.45.0
diff --git a/mf-auth.toml b/mf-auth.toml
index 738abd2..d1fb685 100644
--- a/mf-auth.toml
+++ b/mf-auth.toml
@@ -1,5 +1,8 @@
 ListenAddr = ":8082"
-NATSAddr = "nats:4222"
-BaseSubject = "microfilm.auth.v1"
+UserServiceBaseURL = "http://mf-users:8080"
 
-UserServiceBaseURL = "http://mf-users:8080"
\ No newline at end of file
+[NATS]
+Enabled = true
+Addr = "nats://nats1:4222,nats://nats2:4222,nats://nats3:4222"
+NKeySeed = "SUAOUHJPINF4CK6TSNZMRR5G4DKGW5S76XRNIYURPEISNMWXJIXSVWIO7Y"
+Subject = "microfilm.auth.v1"
\ No newline at end of file
diff --git a/server/config.go b/server/config.go
index 923e23c..4c2f7b0 100644
--- a/server/config.go
+++ b/server/config.go
@@ -7,13 +7,20 @@ import (
 )
 
 type Config struct {
-	ListenAddr  string `toml:"ListenAddr"`
-	NATSAddr    string `toml:"NATSAddr"`
-	BaseSubject string `toml:"BaseSubject"`
+	ListenAddr  string      `toml:"ListenAddr"`
+	NATS        *NATSConfig `toml:"NATS"`
+	BaseSubject string      `toml:"BaseSubject"`
 
 	UserServiceBaseURL string `toml:"UserServiceBaseURL"`
 }
 
+type NATSConfig struct {
+	Enabled  bool   `toml:"Enabled"`
+	NKeySeed string `toml:"NKeySeed"`
+	Addr     string `toml:"Addr"`
+	Subject  string `toml:"Subject"`
+}
+
 func ConfigFromReader(r io.Reader) (*Config, error) {
 	decoder := toml.NewDecoder(r)
 	var c Config
diff --git a/server/server.go b/server/server.go
index 7b4ae70..0cb01ea 100644
--- a/server/server.go
+++ b/server/server.go
@@ -20,6 +20,7 @@ import (
 	"github.com/golang-jwt/jwt/v5"
 	"github.com/google/uuid"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nkeys"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
 	"go.opentelemetry.io/otel/propagation"
@@ -72,16 +73,34 @@ func NewServer(config *Config) (*Server, error) {
 
 	srv.store = store.NewMemoryAuthStore()
 
-	conn, err := nats.Connect(config.NATSAddr)
-	if err != nil {
-		return nil, err
-	}
-	encoded, err := nats.NewEncodedConn(conn, "json")
-	if err != nil {
-		return nil, err
+	if config.NATS.Enabled {
+		var opts []nats.Option
+		if config.NATS.NKeySeed != "" {
+			keys, err := nkeys.FromSeed([]byte(config.NATS.NKeySeed))
+			if err != nil {
+				return nil, err
+			}
+			pubkey, err := keys.PublicKey()
+			if err != nil {
+				return nil, err
+			}
+			srv.Logger.Debug("NATS enabled with NKeys", "pubkey", pubkey)
+			creds := nats.Nkey(pubkey, keys.Sign)
+			opts = append(opts, creds)
+		}
+
+		conn, err := nats.Connect(config.NATS.Addr, opts...)
+		if err != nil {
+			return nil, err
+		}
+		encoded, err := nats.NewEncodedConn(conn, "json")
+		if err != nil {
+			return nil, err
+		}
+
+		srv.nats = encoded
 	}
 
-	srv.nats = encoded
 	srv.userClient = NewUserClient(config.UserServiceBaseURL)
 
 	// Generate keys