github.com/openfga/openfga@v1.5.4-rc1/internal/mocks/mock_oidc_server.go (about)

     1  package mocks
     2  
     3  import (
     4  	"crypto/rand"
     5  	"crypto/rsa"
     6  	"encoding/base64"
     7  	"encoding/json"
     8  	"fmt"
     9  	"log"
    10  	"math/big"
    11  	"net/http"
    12  	"strings"
    13  
    14  	"github.com/golang-jwt/jwt/v4"
    15  )
    16  
    17  type mockOidcServer struct {
    18  	issuerURL  string
    19  	privateKey *rsa.PrivateKey
    20  	publicKey  *rsa.PublicKey
    21  }
    22  
    23  const kidHeader = "1"
    24  
    25  func NewMockOidcServer(issuerURL string) (*mockOidcServer, error) {
    26  	privateKey, err := rsa.GenerateKey(rand.Reader, 4096)
    27  	if err != nil {
    28  		return nil, err
    29  	}
    30  
    31  	mockServer := &mockOidcServer{
    32  		issuerURL:  issuerURL,
    33  		privateKey: privateKey,
    34  		publicKey:  privateKey.Public().(*rsa.PublicKey),
    35  	}
    36  
    37  	mockServer.start()
    38  	return mockServer, nil
    39  }
    40  
    41  func (server *mockOidcServer) NewAliasMockServer(aliasURL string) *mockOidcServer {
    42  	mockServer := &mockOidcServer{
    43  		issuerURL:  aliasURL,
    44  		privateKey: server.privateKey,
    45  		publicKey:  server.privateKey.Public().(*rsa.PublicKey),
    46  	}
    47  
    48  	mockServer.start()
    49  	return mockServer
    50  }
    51  
    52  func (server mockOidcServer) start() {
    53  	port := strings.Split(server.issuerURL, ":")[2]
    54  
    55  	mockHandler := http.NewServeMux()
    56  
    57  	mockHandler.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
    58  		err := json.NewEncoder(w).Encode(map[string]string{
    59  			"issuer":   server.issuerURL,
    60  			"jwks_uri": fmt.Sprintf("%s/jwks.json", server.issuerURL),
    61  		})
    62  		if err != nil {
    63  			log.Fatalf("failed to json encode the openid configurations: %v", err)
    64  		}
    65  	})
    66  
    67  	mockHandler.HandleFunc("/jwks.json", func(w http.ResponseWriter, r *http.Request) {
    68  		err := json.NewEncoder(w).Encode(map[string]interface{}{
    69  			"keys": []map[string]string{
    70  				{
    71  					"kid": kidHeader,
    72  					"kty": "RSA",
    73  					"n":   base64.RawURLEncoding.EncodeToString(server.publicKey.N.Bytes()),
    74  					"e":   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(server.publicKey.E)).Bytes()),
    75  				},
    76  			},
    77  		})
    78  		if err != nil {
    79  			log.Fatalf("failed to json encode the jwks keys: %v", err)
    80  		}
    81  	})
    82  
    83  	go func() {
    84  		log.Fatal(http.ListenAndServe(":"+port, mockHandler))
    85  	}()
    86  }
    87  
    88  func (server mockOidcServer) GetToken(audience, subject string) (string, error) {
    89  	token := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.RegisteredClaims{
    90  		Issuer:   server.issuerURL,
    91  		Audience: []string{audience},
    92  		Subject:  subject,
    93  	})
    94  	token.Header["kid"] = kidHeader
    95  	return token.SignedString(server.privateKey)
    96  }