Updating JWT to use certificates

assume-role.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"crypto/rsa"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -26,6 +27,8 @@ type AssumeRoleResponse struct {
 	Token string `json:"token"`
 }
 
+var privateKey *rsa.PrivateKey
+
 func assumeRole(w http.ResponseWriter, r *http.Request) {
 
 	log.Println("Assume Role Request")
@@ -131,7 +134,7 @@ func assumeRole(w http.ResponseWriter, r *http.Request) {
 
 	// Prevent duplication of roles with the same name for the same account
 
-	checkExisting := pool.QueryRow(context.Background(),
+	checkExisting := libshared.Pool.QueryRow(context.Background(),
 		"SELECT id FROM roles WHERE accountid = $1 AND rolename = $2",
 		claims.Account, role.Rolename)
 
@@ -146,7 +149,7 @@ func assumeRole(w http.ResponseWriter, r *http.Request) {
 
 	fmt.Println("Role ID", existingRoleID)
 
-	roleToken, err := createJWT(secret, claims.Account, role.Rolename)
+	roleToken, err := libshared.CreateJWT(privateKey, claims.Account, role.Rolename, "role")
 	if err != nil {
 		log.Println("Error creating JWT:", err)
 		w.WriteHeader(http.StatusInternalServerError)

create-role.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"io"
+	"libshared"
 	"log"
 	"net/http"
 	"strings"
@@ -126,7 +127,7 @@ func createRole(w http.ResponseWriter, r *http.Request) {
 	log.Println("Description:", role.Description)
 
 	// Prevent duplication of roles with the same name for the same account
-	checkExisting := pool.QueryRow(context.Background(),
+	checkExisting := libshared.Pool.QueryRow(context.Background(),
 		"SELECT FROM roles WHERE accountid = $1 AND rolename = $2",
 		claims.Account, role.Rolename)
 	err = checkExisting.Scan()
@@ -136,7 +137,7 @@ func createRole(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	_, err = pool.Exec(
+	_, err = libshared.Pool.Exec(
 		context.Background(),
 		"INSERT INTO roles (accountid, rolename, description) VALUES ($1, $2, $3)",
 		claims.Account, role.Rolename, role.Description)

db.go

@@ -1,38 +0,0 @@
-package main
-
-import (
-	"context"
-	"fmt"
-	"log"
-	"os"
-
-	"github.com/jackc/pgx/v5/pgxpool"
-)
-
-var pool *pgxpool.Pool
-
-func getDbPool() *pgxpool.Pool {
-	// Construct the connection string
-	// Note: Ensure your Docker Compose env vars match these keys!
-	dburl := fmt.Sprintf("postgres://%s:%s@%s/%s?sslmode=disable",
-		os.Getenv("POSTGRES_USER"),
-		os.Getenv("POSTGRES_PASSWORD"),
-		os.Getenv("POSTGRES_HOSTNAME"),
-		os.Getenv("POSTGRES_DB"),
-	)
-
-	var err error
-	// Use pgxpool.New instead of Connect for v5
-	pool, err = pgxpool.New(context.Background(), dburl)
-	if err != nil {
-		log.Fatalf("Unable to create connection pool: %v", err)
-	}
-
-	// Ping the database to verify the connection is actually live
-	err = pool.Ping(context.Background())
-	if err != nil {
-		log.Fatalf("Unable to ping database: %v", err)
-	}
-
-	return pool
-}

get-role-token.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"io"
+	"libshared"
 	"log"
 	"net/http"
 	"strings"
@@ -90,7 +91,7 @@ func getRoleToken(w http.ResponseWriter, r *http.Request) {
 	log.Println("Authentication would taken place here")
 
 	// Check if policy with the same name already exists for the account
-	checkExisting := pool.QueryRow(context.Background(),
+	checkExisting := libshared.Pool.QueryRow(context.Background(),
 		"SELECT FROM roles WHERE accountid = $1 AND rolename = $2",
 		claims.Account, role.RoleName)
 	log.Println("checkExisting:", role.RoleName)
@@ -100,7 +101,7 @@ func getRoleToken(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	usertoken, err := createJWT(secret, claims.Account, role.RoleName)
+	usertoken, err := libshared.CreateJWT(privateKey, claims.Account, role.RoleName, "role")
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		return

jwt.go

@@ -1,22 +0,0 @@
-package main
-
-import (
-	"time"
-
-	"github.com/golang-jwt/jwt/v5"
-)
-
-// CreateJWT generates a signed JWT
-func createJWT(secret []byte, account string, role string) (string, error) {
-	claims := jwt.MapClaims{
-		"sub":     role,                             // subject (user id)
-		"exp":     time.Now().Add(time.Hour).Unix(), // expiration
-		"iat":     time.Now().Unix(),                // issued at
-		"purpose": "role",
-		"account": account,
-	}
-
-	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
-
-	return token.SignedString(secret)
-}

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"libshared"
 	"log"
 	"net/http"
 )
@@ -9,7 +10,14 @@ func LoginHandler(w http.ResponseWriter, r *http.Request) {
 }
 
 func main() {
-	pool = getDbPool()
+
+	var err error
+
+	libshared.Pool = libshared.GetDbPool()
+	privateKey, err = libshared.LoadPrivateKey("keys/private.pem")
+	if err != nil {
+		log.Fatal("Failed to load private key:", err)
+	}
 
 	http.HandleFunc("/role/create-role", createRole)
 	http.HandleFunc("/role/attach-role", attachRole)