main.go

package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"time"

	"github.com/nats-io/jwt/v2"
	"github.com/nats-io/nats-server/v2/server"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/micro"
	"github.com/nats-io/nkeys"
)

func main() {
	issuerKeyPair, err := nkeys.CreateAccount()
	if err != nil {
		log.Fatalf("Failed to create issuer keypair: %v", err)
	}

	issuerPubKey, err := issuerKeyPair.PublicKey()
	if err != nil {
		log.Fatalf("Failed to get issuer public key: %v", err)
	}

	ns, err := createServer(issuerPubKey)
	if err != nil {
		log.Fatalf("Failed to create server: %v", err)
	}

	go ns.Start()

	if !ns.ReadyForConnections(10 * time.Second) {
		log.Fatalf("NATS server not ready for connections")
	}

	natsConn, err := nats.Connect("",
		nats.InProcessServer(ns),
		nats.UserInfo("auth", ""),
	)
	if err != nil {
		log.Fatalf("Failed to create NATS connection: %v", err)
	}

	if err := setupAuthHandler(natsConn, issuerKeyPair); err != nil {
		log.Fatalf("Failed to set up auth handler: %v", err)
	}

	log.Println("Ready")

	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, os.Interrupt)
	<-sigch
}

func createServer(issuerPubKey string) (*server.Server, error) {
	var (
		authAccount = server.NewAccount("AUTH")
		sysAccount  = server.NewAccount("SYS")
	)

	opts := &server.Options{
		Host:          "0.0.0.0",
		Port:          4222,
		HTTPPort:      8222,
		SystemAccount: sysAccount.Name,
		Users: []*server.User{
			{
				Username: "auth",
				Account:  authAccount,
			},
		},
		Accounts: []*server.Account{
			authAccount,
			sysAccount,
		},
		AuthCallout: &server.AuthCallout{
			Issuer:    issuerPubKey,
			AuthUsers: []string{"auth"},
			Account:   authAccount.Name,
		},
	}

	return server.NewServer(opts)
}

func setupAuthHandler(nc *nats.Conn, issuerKeyPair nkeys.KeyPair) error {
	srv, err := micro.AddService(nc, micro.Config{
		Name:    "auth-callout",
		Version: "0.1.0",
	})
	if err != nil {
		return fmt.Errorf("failed to create micro service: %v", err)
	}

	g := srv.
		AddGroup("$SYS").
		AddGroup("REQ").
		AddGroup("USER")

	handler := func(r micro.Request) {
		rc, err := jwt.DecodeAuthorizationRequestClaims(string(r.Data()))
		if err != nil {
			log.Printf("Failed to decode auth request: %v", err)
			return
		}

		log.Printf("Auth request for nkey: %s", rc.UserNkey)

		uc := jwt.NewUserClaims(rc.UserNkey)
		uc.Name = rc.Name
		uc.Audience = "SYS"

		userJwt, err := uc.Encode(issuerKeyPair)
		if err != nil {
			log.Printf("Failed to encode user JWT: %v", err)
			return
		}

		authResp := jwt.NewAuthorizationResponseClaims(rc.UserNkey)
		authResp.Audience = rc.Server.ID
		authResp.Jwt = userJwt

		token, err := authResp.Encode(issuerKeyPair)
		if err != nil {
			log.Printf("Failed to encode response: %v", err)
			return
		}

		r.Respond([]byte(token))
	}

	if err := g.AddEndpoint("AUTH", micro.HandlerFunc(handler)); err != nil {
		return fmt.Errorf("failed to add endpoint: %v", err)
	}

	return nil
}