setup support for authorization

5 files changed, 34 insertions, 8 deletions

diff --git a/config/config.go b/config/config.go
index 57c8aa8..9eee8e2 100644
--- a/config/config.go
+++ b/config/config.go
@@ -4,6 +4,7 @@ import (
 	"log"
 	"os"
 	"strconv"
+	"thunderbird-ai-compose-server/crypto"
 	"thunderbird-ai-compose-server/types"
 
 	"github.com/joho/godotenv"
@@ -15,10 +16,11 @@ func init() {
 	godotenv.Load()
 
 	Config = types.ServerConfig{
-		Port:     getEnvAsInt("PORT"),
-		Provider: types.Provider(getEnv("PROVIDER")),
-		Model:    getEnv("MODEL"),
-		APIKey:   getEnv("API_KEY"),
+		AuthorizationKey: crypto.GenerateSecretKey(),
+		Port:             getEnvAsInt("PORT"),
+		Provider:         types.Provider(getEnv("PROVIDER")),
+		Model:            getEnv("MODEL"),
+		APIKey:           getEnv("API_KEY"),
 	}
 
 	if Config.Port == 0 {
diff --git a/crypto/secret.go b/crypto/secret.go
new file mode 100644
index 0000000..cc8c9b0
--- /dev/null
+++ b/crypto/secret.go
@@ -0,0 +1,12 @@
+package crypto
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+)
+
+func GenerateSecretKey() string {
+	key := make([]byte, 48)
+	rand.Read(key)
+	return hex.EncodeToString(key)
+}
diff --git a/main.go b/main.go
index aa5ae3c..937507a 100644
--- a/main.go
+++ b/main.go
@@ -18,5 +18,9 @@ func main() {
 	routers.Setup(app)
 
 	log.Printf("Starting server on port %d\n", config.Config.Port)
+	log.Println("Configure your Extension with the following details:")
+	log.Printf("Endpoint URL: http://localhost:%d\n", config.Config.Port)
+	log.Printf("Authorization Key: %s\n", config.Config.AuthorizationKey)
+	log.Println("Note: Keep the Authorization Key secure and do not share it publicly. The Authorization Key will change each time the server restarts. Use this to reset the key if needed.")
 	log.Fatal(app.Listen(":3000"))
 }
diff --git a/routers/router.go b/routers/router.go
index cae38db..ec39f5b 100644
--- a/routers/router.go
+++ b/routers/router.go
@@ -12,6 +12,13 @@ import (
 
 func Setup(router *fiber.App) {
 	router.Post("/generate", func(c *fiber.Ctx) error {
+		authHeader := c.Get("Authorization")
+		if authHeader != "Bearer "+config.Config.AuthorizationKey {
+			return c.Status(fiber.StatusUnauthorized).JSON(types.ErrorResponse{
+				Error: "Invalid or missing Authorization header",
+			})
+		}
+
 		var payload types.Payload
 
 		if err := c.BodyParser(&payload); err != nil {
diff --git a/types/types.go b/types/types.go
index 955d085..0d7ca44 100644
--- a/types/types.go
+++ b/types/types.go
@@ -8,10 +8,11 @@ const (
 )
 
 type ServerConfig struct {
-	Port     int
-	Provider Provider
-	Model    string
-	APIKey   string
+	AuthorizationKey string
+	Port             int
+	Provider         Provider
+	Model            string
+	APIKey           string
 }
 
 type Identity struct {