[JWT Auth] Support JWT Authentication

bayeux.go

@@ -3,6 +3,7 @@ package bayeux
 import (
 	"bytes"
 	"context"
+	"crypto/rsa"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -14,6 +15,8 @@ import (
 	"os"
 	"sync"
 	"time"
+
+	"github.com/golang-jwt/jwt"
 )
 
 type MaybeMsg struct {
@@ -100,6 +103,9 @@ type AuthenticationParameters struct {
 	Username     string // Salesforce user email (e.g. salesforce.user@email.com)
 	Password     string // Salesforce password
 	TokenURL     string // Salesforce token endpoint (e.g. https://login.salesforce.com/services/oauth2/token)
+	Path         string // Salesforce private key path
+	IsJwt        bool   // Salesforce auth method identifier
+	Audience     string // Salesforce authorization server’s URL for the audience value (e.g. https://login.salesforce.com or https://test.salesforce.com or https://site.force.com/customers)
 }
 
 // Bayeux struct allow for centralized storage of creds, ids, and cookies
@@ -108,6 +114,11 @@ type Bayeux struct {
 	id    clientIDAndCookies
 }
 
+type Authentication struct {
+	URLValues      *url.Values
+	AuthParameters *AuthenticationParameters
+}
+
 var wg sync.WaitGroup
 var logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile)
 var st = status{false, "", []string{}, 0}
@@ -299,19 +310,37 @@ func GetConnectedCount() int {
 	return st.connectCount
 }
 
-func GetSalesforceCredentials(ap AuthenticationParameters) (creds *Credentials, err error) {
-	params := url.Values{"grant_type": {"password"},
-		"client_id":     {ap.ClientID},
-		"client_secret": {ap.ClientSecret},
-		"username":      {ap.Username},
-		"password":      {ap.Password}}
-	res, err := http.PostForm(ap.TokenURL, params)
+func GetSalesforceCredentials(auth Authentication) (creds *Credentials, err error) {
+	var params *Authentication
+
+	if auth.AuthParameters == nil {
+		return nil, fmt.Errorf("auth parameters are empty")
+	}
+	if auth.AuthParameters.IsJwt {
+		params, err = GetJWTAuthentication(*auth.AuthParameters)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+
+		if auth.AuthParameters.ClientID == "" || auth.AuthParameters.ClientSecret == "" || auth.AuthParameters.Username == "" || auth.AuthParameters.Password == "" {
+			return nil, fmt.Errorf("missing required authentication parameters")
+		}
+		params, _ = GetClientCredentialAuthentication(auth.AuthParameters.ClientID, auth.AuthParameters.ClientSecret, auth.AuthParameters.Username, auth.AuthParameters.Password, auth.AuthParameters.TokenURL)
+	}
+	if auth.AuthParameters.TokenURL == "" {
+		return nil, fmt.Errorf("missing required authentication parameter: token_url")
+	}
+
+	res, err := http.PostForm(auth.AuthParameters.TokenURL, *params.URLValues)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("error posting form: %w", err)
 	}
+	defer res.Body.Close()
+
 	decoder := json.NewDecoder(res.Body)
 	if err := decoder.Decode(&creds); err == io.EOF {
-		return nil, err
+		return nil, fmt.Errorf("error decoding response: %w", err)
 	} else if err != nil {
 		return nil, err
 	} else if creds.AccessToken == "" {
@@ -320,6 +349,79 @@ func GetSalesforceCredentials(ap AuthenticationParameters) (creds *Credentials,
 	return creds, nil
 }
 
+// GetJWTAuthentication prepares the authentication parameters for JWT-based authentication
+func GetJWTAuthentication(ap AuthenticationParameters) (*Authentication, error) {
+	// Define your JWT claims (payload)
+	claims := jwt.MapClaims{
+		"iss": ap.ClientID,                          // Issuer
+		"sub": ap.Username,                          // Subject
+		"aud": ap.Audience,                          // Audience
+		"exp": time.Now().Add(1 * time.Hour).Unix(), // Expiration time (1 hour)
+	}
+
+	// Load your private key (RSA key) used for signing
+	privateKey, err := loadPrivateKey(ap.Path)
+	if err != nil {
+		return nil, fmt.Errorf("Error loading private key: %w", err)
+	}
+
+	// Create a new JWT token
+	tokenString, err := jwt.NewWithClaims(jwt.SigningMethodRS256, claims).SignedString(privateKey)
+	if err != nil {
+		return nil, fmt.Errorf("Error signing JWT token: %w", err)
+	}
+
+	return &Authentication{
+		URLValues: &url.Values{
+			"grant_type": {"urn:ietf:params:oauth:grant-type:jwt-bearer"},
+			"assertion":  {tokenString},
+		},
+		AuthParameters: &AuthenticationParameters{
+			ClientSecret: ap.ClientSecret,
+			Username:     ap.Username,
+			Audience:     ap.Audience,
+			Path:         ap.Path,
+		},
+	}, nil
+}
+
+// GetClientCredentialAuthentication prepares the authentication parameters for client credential-based authentication
+func GetClientCredentialAuthentication(clientId, clientSecret, username, password, tokenUrl string) (*Authentication, error) {
+	if clientId != "" && clientSecret != "" && username != "" && password != "" && tokenUrl != "" {
+		return nil, errors.New("all authentication parameters must be set")
+	}
+
+	return &Authentication{
+		URLValues: &url.Values{
+			"grant_type":    {"password"},
+			"client_id":     {clientId},
+			"client_secret": {clientSecret},
+			"username":      {username},
+			"password":      {password},
+		},
+		AuthParameters: &AuthenticationParameters{
+			ClientID:     clientId,
+			ClientSecret: clientSecret,
+			Username:     username,
+			Password:     password,
+			TokenURL:     tokenUrl,
+		},
+	}, nil
+}
+func loadPrivateKey(keyFile string) (*rsa.PrivateKey, error) {
+	keyBytes, err := os.ReadFile(keyFile)
+	if err != nil {
+		return nil, err
+	}
+
+	privateKey, err := jwt.ParseRSAPrivateKeyFromPEM(keyBytes)
+	if err != nil {
+		return nil, err
+	}
+
+	return privateKey, nil
+}
+
 func (b *Bayeux) Channel(ctx context.Context, out chan MaybeMsg, r string, creds Credentials, channel string) chan MaybeMsg {
 	b.creds = creds
 	err := b.getClientID(ctx)

example_test.go

@@ -17,7 +17,13 @@ func Example() {
 	ap.Username = "salesforce.user@email.com"
 	ap.Password = "foobar"
 	ap.TokenURL = "https://login.salesforce.com/services/oauth2/token"
-	creds, _ := GetSalesforceCredentials(ap)
+
+	//Create a variable of type Authentication
+	auth := Authentication{
+		AuthParameters: &ap,
+	}
+
+	creds, _ := GetSalesforceCredentials(auth)
 	c := b.Channel(ctx, out, replay, *creds, "channel")
 	for {
 		select {

examples/main.go

@@ -12,12 +12,19 @@ func Example() {
 	out := make(chan bay.MaybeMsg)
 	b := bay.Bayeux{}
 	var ap bay.AuthenticationParameters
+	// make it false to use user-password flow
+	ap.IsJwt = true
 	ap.ClientID = "3MVG9pRsdbjsbdjfm1I.fz3f7zBuH4xdKCJcM9B5XLgxXh2AFTmQmr8JMn1vsadjsadjjsadakd_C"
 	ap.ClientSecret = "E9FE118633BC7SGDADUHUE81F19C1D4529D09CB7231754AD2F2CA668400619"
 	ap.Username = "salesforce.user@email.com"
 	ap.Password = "foobar"
 	ap.TokenURL = "https://login.salesforce.com/services/oauth2/token"
-	creds, _ := bay.GetSalesforceCredentials(ap)
+	ap.Path = "/path/to/private.key"
+	ap.Audience = "https://login.salesforce.com"
+	auth := bay.Authentication{
+		AuthParameters: &ap,
+	}
+	creds, _ := bay.GetSalesforceCredentials(auth)
 	replay := "-1"
 	c := b.Channel(ctx, out, replay, *creds, "channel")
 	for {

go.mod

@@ -1,3 +1,5 @@
 module github.com/elastic/bayeux
 
 go 1.17
+
+require github.com/golang-jwt/jwt v3.2.2+incompatible

go.sum

@@ -0,0 +1,2 @@
+github.com/golang-jwt/jwt v3.2.2+incompatible h1:IfV12K8xAKAnZqdXVzCZ+TOjboZ2keLg81eXfW3O+oY=
+github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=