Keycloak Go RP Client Example with PKCE

main.go

/*
 This is an example about how to use a public client written in Golang to authenticate using Keycloak.
 This example is only for demonstration purposes and lacks important 
*/
package main

import (
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"time"
	"fmt"

	oidc "github.com/coreos/go-oidc"
	"github.com/abourget/getting-started-with-golang/pkce"
	"github.com/google/uuid"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

var oidcProvider oidc.Provider
var oidcConfig oidc.Config
var oauth2Config oauth2.Config
var idTokenVerifier oidc.IDTokenVerifier

func init() {
	oidcProvider = *createOidcProvider(context.Background())
	oidcConfig, oauth2Config = createConfig(oidcProvider)
	idTokenVerifier = *oidcProvider.Verifier(&oidcConfig)
}

func createOidcProvider(ctx context.Context) *oidc.Provider {
	// provider, err := oidc.NewProvider(ctx, "http://localhost:8180/auth/realms/myrealm")
	provider, err := oidc.NewProvider(ctx, "http://localhost:8180/realms/myrealm")

	if err != nil {
		log.Fatal("Failed to fetch discovery document: ", err)
	}

	return provider
}

func createConfig(provider oidc.Provider) (oidc.Config, oauth2.Config) {
	oidcConfig := &oidc.Config{
		ClientID: "mywebapp",
	}

	config := oauth2.Config{
		ClientID:    oidcConfig.ClientID,
		// ClientSecret: "CLIENT_SECRET", // change CLIENT_SECRET with the secret generated by Keycloak for the mywebapp client. This option is a string.
		ClientSecret: "EVKTysjpeMD31P1rUGEDxerIqlezAo6F",
		Endpoint:    provider.Endpoint(),
		RedirectURL: "http://localhost:8080/auth/callback",
		Scopes:      []string{oidc.ScopeOpenID, "profile", "email"},
	
	}

	return *oidcConfig, config
}

func main() {
	http.HandleFunc("/", redirectHandler)
	http.HandleFunc("/auth/callback", callbackHandler)

	log.Printf("To authenticate go to http://%s/", "localhost:8080")
	log.Fatal(http.ListenAndServe("localhost:8080", nil))
}

func redirectHandler(resp http.ResponseWriter, r *http.Request) {
	pkceCode, err := pkce.Generate()
	addPkceCookie(pkceCode, resp)
	state := addStateCookie(resp)

	if err != nil {
		http.Error(resp, "Failed to generate pkce challenge", http.StatusInternalServerError)
		return
	}

	http.Redirect(resp, r, oauth2Config.AuthCodeURL(state, pkceCode.Challenge(), pkceCode.Method()), http.StatusFound)
}

func callbackHandler(resp http.ResponseWriter, req *http.Request) {
	err := checkStateAndExpireCookie(req, resp)

	if err != nil {
		redirectHandler(resp, req)
		return
	}

	tokenResponse, err := exchangeCode(req)

	if err != nil {
		http.Error(resp, "Failed to exchange code", http.StatusBadRequest)
		return
	}

	idToken, err := validateIDToken(tokenResponse, req)

	if err != nil {
		http.Error(resp, "Failed to validate id_token", http.StatusUnauthorized)
		return
	}

	handleSuccessfulAuthentication(tokenResponse, *idToken, resp)
}

func addStateCookie(resp http.ResponseWriter) string {
	expire := time.Now().Add(1 * time.Minute)
	value := uuid.New().String()

	cookie := http.Cookie{
		Name:     "p_state",
		Value:    value,
		Expires:  expire,
		HttpOnly: true, // TODO: Add Secure: true
	}

	http.SetCookie(resp, &cookie)

	return value
}

func addPkceCookie(code pkce.Code, resp http.ResponseWriter) {
	expire := time.Now().Add(1 * time.Minute)
	value := string(code) // TODO: encrypt pkce value or store it on server-side session

	cookie := http.Cookie{
		Name:     "p_pkce",
		Value:    value,
		Expires:  expire,
		HttpOnly: true, // TODO: Add Secure: true
	}

	http.SetCookie(resp, &cookie)
}

func expireCookie(name string, resp http.ResponseWriter) {
	cookie := &http.Cookie{
		Name:     "p_state",
		Value:    "",
		MaxAge:   -1,
		HttpOnly: true,
	}

	http.SetCookie(resp, cookie)
}

func checkStateAndExpireCookie(req *http.Request, resp http.ResponseWriter) error {
	state, err := req.Cookie("p_state")

	expireCookie("p_state", resp)

	if err != nil {
		return errors.New("state cookie not set")
	}

	if req.URL.Query().Get("state") != state.Value {
		return errors.New("invalid state")
	}

	return nil
}

func exchangeCode(req *http.Request) (*oauth2.Token, error) {
	httpClient := &http.Client{Timeout: 2 * time.Second}
	ctx := context.WithValue(req.Context(), oauth2.HTTPClient, httpClient)
	pkceCookie, err := req.Cookie("p_pkce")
	pkceCode := pkce.Code(pkceCookie.Value)

	tokenResponse, err := oauth2Config.Exchange(ctx, req.URL.Query().Get("code"), pkceCode.Verifier())

	if err != nil {
		return nil, err
	}

	return tokenResponse, nil
}

func validateIDToken(tokenResponse *oauth2.Token, req *http.Request) (*oidc.IDToken, error) {
	rawIDToken, ok := tokenResponse.Extra("id_token").(string)

	if !ok {
		return nil, errors.New("id_token is not in the token response")
	}

	idToken, err := idTokenVerifier.Verify(req.Context(), rawIDToken)

	if err != nil {
		return nil, err
	}

	return idToken, nil
}

func handleSuccessfulAuthentication(tokenResponse *oauth2.Token, idToken oidc.IDToken, resp http.ResponseWriter) {
	payload := struct {
		TokenResponse *oauth2.Token
		IDToken       *json.RawMessage
	}{tokenResponse, new(json.RawMessage)}

	if err := idToken.Claims(&payload.IDToken); err != nil {
		return
	}

	data, err := json.MarshalIndent(&payload, "", "    ")

	if err != nil {
		http.Error(resp, err.Error(), http.StatusInternalServerError)
		return
	}

	resp.Write(data)
}

pkce.go

// adopted from https://github.com/vmware-tanzu/pinniped/blob/v0.21.0/pkg/oidcclient/pkce/pkce.go
package pkce

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"

	"golang.org/x/oauth2"
)

// Generate generates a new random PKCE code.
func Generate() (Code, error) { return generate(rand.Reader) }

func generate(rand io.Reader) (Code, error) {
	// From https://tools.ietf.org/html/rfc7636#section-4.1:
	//   code_verifier = high-entropy cryptographic random STRING using the
	//   unreserved characters [A-Z] / [a-z] / [0-9] / "-" / "." / "_" / "~"
	//   from Section 2.3 of [RFC3986], with a minimum length of 43 characters
	//   and a maximum length of 128 characters.
	var buf [32]byte
	if _, err := io.ReadFull(rand, buf[:]); err != nil {
		return "", fmt.Errorf("could not generate PKCE code: %w", err)
	}
	return Code(hex.EncodeToString(buf[:])), nil
}

// Code implements the basic options required for RFC 7636: Proof Key for Code Exchange (PKCE).
type Code string

// Challenge returns the OAuth2 auth code parameter for sending the PKCE code challenge.
func (p *Code) Challenge() oauth2.AuthCodeOption {
	b := sha256.Sum256([]byte(*p))
	return oauth2.SetAuthURLParam("code_challenge", base64.RawURLEncoding.EncodeToString(b[:]))
}

// Method returns the OAuth2 auth code parameter for sending the PKCE code challenge method.
func (p *Code) Method() oauth2.AuthCodeOption {
	return oauth2.SetAuthURLParam("code_challenge_method", "S256")
}

// Verifier returns the OAuth2 auth code parameter for sending the PKCE code verifier.
func (p *Code) Verifier() oauth2.AuthCodeOption {
	return oauth2.SetAuthURLParam("code_verifier", string(*p))
}