Jade/login.go

package main

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"

	"github.com/edgedb/edgedb-go"
	"github.com/gofiber/fiber/v2"
)

const EDGEDB_AUTH_BASE_URL = "http://127.0.0.1:10700/db/main/ext/auth"

func generatePKCE() (string, string) {
	verifier_source := make([]byte, 32)
	_, err := rand.Read(verifier_source)
	if err != nil {
		panic(err)
	}

	verifier := base64.RawURLEncoding.EncodeToString(verifier_source)
	challenge := sha256.Sum256([]byte(verifier))
	return verifier, base64.RawURLEncoding.EncodeToString(challenge[:])
}

func handleUiSignIn(c *fiber.Ctx) error {
	verifier, challenge := generatePKCE()

	fmt.Println("Challenge: ", challenge)

	c.Cookie(&fiber.Cookie{
		Name:     "jade-edgedb-pkce-verifier",
		Value:    verifier,
		HTTPOnly: true,
		Path:     "/",
		Secure:   true,
		SameSite: "Strict",
	})

	return c.Redirect(fmt.Sprintf("%s/ui/signup?challenge=%s", EDGEDB_AUTH_BASE_URL, challenge), fiber.StatusTemporaryRedirect)
}

func handleCallbackSignup(c *fiber.Ctx) error {
	code := c.Query("code")
	if code == "" {
		error := c.Query("error")
		return c.Status(fiber.StatusBadRequest).SendString(fmt.Sprintf("OAuth callback is missing 'code'. OAuth provider responded with error: %s", error))
	}

	verifier := c.Cookies("jade-edgedb-pkce-verifier")
	if verifier == "" {
		return c.Status(fiber.StatusBadRequest).SendString("Could not find 'verifier' in the cookie store. Is this the same user agent/browser that started the authorization flow?")
	}

	codeExchangeURL := fmt.Sprintf("%s/token?code=%s&verifier=%s", EDGEDB_AUTH_BASE_URL, code, verifier)
	resp, err := http.Get(codeExchangeURL)
	if err != nil {
		return c.Status(fiber.StatusBadRequest).SendString(fmt.Sprintf("Error from the auth server: %s", err.Error()))
	}
	defer resp.Body.Close()

	if resp.StatusCode != fiber.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return c.Status(fiber.StatusBadRequest).SendString(fmt.Sprintf("Error from the auth server: %s", string(body)))
	}

	var tokenResponse struct {
		AuthToken  string `json:"auth_token"`
		IdentityID string `json:"identity_id"`
	}
	err = json.NewDecoder(resp.Body).Decode(&tokenResponse)
	if err != nil {
		return c.Status(fiber.StatusBadRequest).SendString(fmt.Sprintf("Error decoding auth server response: %s", err.Error()))
	}

	c.Cookie(&fiber.Cookie{
		Name:     "jade-edgedb-auth-token",
		Value:    tokenResponse.AuthToken,
		HTTPOnly: true,
		Path:     "/",
		Secure:   true,
		SameSite: "Strict",
	})

	// Create a new User and attach the identity
	var identityUUID edgedb.UUID
	identityUUID, err = edgedb.ParseUUID(tokenResponse.IdentityID)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).SendString(fmt.Sprintf("Error in edgedb.ParseUUID: in handleCallbackSignup: %s", err.Error()))
	}

	err = edgeClient.Execute(edgeCtx, `
	INSERT User {
		setting := (
			INSERT Setting {
				default_model := "gpt-3.5-turbo"
			}
		),
		identity := (SELECT ext::auth::Identity FILTER .id = <uuid>$0)
	  }
	`, identityUUID)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).SendString(fmt.Sprintf("Error in edgedb.QuerySingle: in handleCallbackSignup: %s", err.Error()))
	}

	edgeClient = edgeClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": tokenResponse.AuthToken})

	return c.Redirect("/", fiber.StatusTemporaryRedirect)
}

func handleCallback(c *fiber.Ctx) error {
	code := c.Query("code")
	if code == "" {
		error := c.Query("error")
		return c.Status(fiber.StatusBadRequest).SendString(fmt.Sprintf("OAuth callback is missing 'code'. OAuth provider responded with error: %s", error))
	}

	verifier := c.Cookies("jade-edgedb-pkce-verifier")
	if verifier == "" {
		fmt.Println("Could not find 'verifier' in the cookie store. Is this the same user agent/browser that started the authorization flow?")
		return c.Status(fiber.StatusBadRequest).SendString("Could not find 'verifier' in the cookie store. Is this the same user agent/browser that started the authorization flow?")
	}

	codeExchangeURL := fmt.Sprintf("%s/token?code=%s&verifier=%s", EDGEDB_AUTH_BASE_URL, code, verifier)
	resp, err := http.Get(codeExchangeURL)
	if err != nil {
		return c.Status(fiber.StatusBadRequest).SendString(fmt.Sprintf("Error from the auth server: %s", err.Error()))
	}
	defer resp.Body.Close()

	if resp.StatusCode != fiber.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return c.Status(fiber.StatusBadRequest).SendString(fmt.Sprintf("Error from the auth server: %s", string(body)))
	}

	var tokenResponse struct {
		AuthToken string `json:"auth_token"`
	}
	err = json.NewDecoder(resp.Body).Decode(&tokenResponse)
	if err != nil {
		return c.Status(fiber.StatusBadRequest).SendString(fmt.Sprintf("Error decoding auth server response: %s", err.Error()))
	}

	c.Cookie(&fiber.Cookie{
		Name:     "jade-edgedb-auth-token",
		Value:    tokenResponse.AuthToken,
		HTTPOnly: true,
		Path:     "/",
		Secure:   true,
		SameSite: "Strict",
	})

	edgeClient = edgeClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": tokenResponse.AuthToken})

	return c.Redirect("/", fiber.StatusTemporaryRedirect)
}

func handleSignOut(c *fiber.Ctx) error {
	c.ClearCookie("jade-edgedb-auth-token")
	edgeClient.Close()

	var ctx = context.Background()
	client, err := edgedb.CreateClient(ctx, edgedb.Options{})
	if err != nil {
		fmt.Println("Error in edgedb.CreateClient: in init")
		log.Fatal(err)
	}

	edgeCtx = ctx
	edgeClient = client

	return c.Redirect("/", fiber.StatusTemporaryRedirect)
}