Hey everyone,

welcome back, in this post we are going to learn how to implement the JWT authentication using Apex in Salesforce.

Sometimes you want to authorize servers to access data without interactively logging in each time the servers exchange information. For these cases, you can use the OAuth 2.0 JSON Web Token (JWT) bearer flow. This flow uses a certificate to sign the JWT request and doesn’t require explicit user interaction. However, this flow does require prior approval of the client app.

When we talk about JSON Web Token, it is consist of 3 parts

1. Headers – Which contains the algorithm which will be used to sign the request
2. Payload – The actual data and information about the users
3. Signature – Signature consists of 3 parts and the structure is given below.

HMACSHA256(
  base64UrlEncode(header) + "." +
  base64UrlEncode(payload),
  Private Key
)

In this blog post, we are going to authenticate Einstein Platform Services and you can get your account from here. Sign up using your salesforce account and download the key which we will be used for authentication purposes.

JWT can be authenticated using 2 ways

• Using Private Secret key – The platform which we are trying to access will provide the Key.
• Using Certification

let’s start preparing all three parts one by one

Step 1 – Prepare header

As we know that header contains the algorithm. So it will look like below

{“alg”:”RS256″}

Step 2 – Prepare Payload

Payload consist of 4 parts

1. iss – The issuer must contain the OAuth client_id or the connected app for which you registered the certificate. and in our case it is developer.force.com
2. aud – The audience identifies the authorization server as an intended audience. For Example, https://login.salesforce.com, https://test.salesforce.com, or https://community.force.com/customers if implementing for a community. or https://api.einstein.ai/v2/oauth2/token for einstein API
3. sub- is either email of the user or target username
4. exp – The validity must be the expiration time of the assertion within 3 minutes, expressed as the number of seconds from 1970-01-01T0:0:0Z measured in UTC.

Here is the Example

{
  "iss": "developer.force.com",
  "sub": "my@email.com",
  "aud": "https://api.einstein.ai/v2/oauth2/token",
  "exp": "1333685628"
}

Now, once we have prepared the payload, we need to url encode this as base64.

After we are done with the encoding the string now Create a string for the Encoded JWT Header and the encoded JWT Claims Set in this format. Like below

encoded_JWT_Header + "." + encoded_JWT_Claims_Set

Step 3 – Sign the result from Step 2 using Private Key or cert

encoded_JWT_Header + "." + encoded_JWT_Claims_Set

Now, as we have prepared all three parts of JWT the final step is to get the access token. To get the access token we need to make the post request. Below is sample request.

POST /services/oauth2/token HTTP/1.1
Host: login.example.com
Content-Type: application/x-www-form-urlencoded

grant_type= urn:ietf:params:oauth:grant-type:jwt-bearer&
assertion=eyJpc3MiOiAiM01WRz...[omitted for brevity]...ZT

Include these parameters in the post.

To prepare the JSON we are using JSON class and JSONGenerator which are provided by salesforce.

Remember, we signed up for the Einstein Platform Account and downloaded the key file. Upload that key files under files inside salesforce.

Below is the complete code for the authentication.

JWT Class to prepare the JWT header, payload.

public class JWT {
    
    public String alg {get;set;}
    public String iss {get;set;}
    public String sub {get;set;}
    public String aud {get;set;}
    public String exp {get;set;}
    public String iat {get;set;}
    public Map<String,String> claims {get;set;}
    public Integer validFor {get;set;}
    public String cert {get;set;}
    public String pkcs8 {get;set;}
    public String privateKey {get;set;}
    
    public static final String HS256 = 'HS256';
    public static final String RS256 = 'RS256';
    public static final String NONE = 'none';
    
    public JWT(String alg) {
        this.alg = alg;
        this.validFor = 3600;
    }
    
    public String assertion() {
        
        String jwt = '';
        JSONGenerator header = JSON.createGenerator(false);
        header.writeStartObject();
        header.writeStringField('alg', this.alg);
        header.writeEndObject();
        String encodedHeader = base64URLencode(Blob.valueOf(header.getAsString()));
        
        JSONGenerator body = JSON.createGenerator(false);
        body.writeStartObject();
        body.writeStringField('iss', this.iss);
        body.writeStringField('sub', this.sub);
        body.writeStringField('aud', this.aud);
        Long rightNow = (dateTime.now().getTime()/1000)+1;
        body.writeNumberField('iat', rightNow);
        body.writeNumberField('exp', (rightNow + validFor));
        if (claims != null) {
            for (String claim : claims.keySet()) {
                body.writeStringField(claim, claims.get(claim));
            }
        }
        body.writeEndObject();
        jwt = encodedHeader + '.' + base64URLencode(Blob.valueOf(body.getAsString()));
        
        if ( this.alg == HS256 ) {
            Blob key = EncodingUtil.base64Decode(privateKey);
            Blob signature = Crypto.generateMac('hmacSHA256',Blob.valueof(jwt),key);
            jwt += '.' + base64URLencode(signature);  
        } else if ( this.alg == RS256 ) {
            Blob signature = null;
            
            if (cert != null ) {
                signature = Crypto.signWithCertificate('rsa-sha256', Blob.valueOf(jwt), cert);
            } else {
                Blob privateKey = EncodingUtil.base64Decode(pkcs8);
                signature = Crypto.sign('rsa-sha256', Blob.valueOf(jwt), privateKey);
            }
            jwt += '.' + base64URLencode(signature);  
        } else if ( this.alg == NONE ) {
            jwt += '.';
        }
        return jwt;
    }
    
    public String base64URLencode(Blob input){ 
        String output = encodingUtil.base64Encode(input);
        output = output.replace('+', '-');
        output = output.replace('/', '_');
        while ( output.endsWith('=')){
            output = output.subString(0,output.length()-1);
        }
        return output;
    }
}

JWTBearerFlow Class for getting the Access Token

public class JWTBearerFlow {
    
    public static String getAccessToken(String tokenEndpoint, JWT jwt) {
        
        String grantType = 'urn:ietf:params:oauth:grant-type:jwt-bearer';
        String access_token = null;
        String body = 'grant_type='+EncodingUtil.urlEncode(grantType, 'UTF-8')+'&assertion=' + jwt.assertion();
        HttpRequest req = new HttpRequest();                            
        req.setMethod('POST');
        req.setEndpoint(tokenEndpoint);
        req.setHeader('Content-type', 'application/x-www-form-urlencoded');
        req.setBody(body);
        Http http = new Http();               
        try{
            HTTPResponse res = http.send(req);
            if ( res.getStatusCode() == 200 ) {
                Map<String, Object> responseMap = (Map<String, Object>)JSON.deserializeUntyped(res.getBody());
                access_token = (String)responseMap.get('access_token');
            }else{
                System.debug('JWTBearerFlow Error Occurred '+res.getBody());
            }
        }catch(Exception ex){
            if(String.valueOf(ex).startsWith('Unauthorized endpoint')){
                System.debug('JWTBearerFlow Please check Setup->Security->Remote site settings and add '+tokenEndpoint);
            }else{
                System.debug('JWTBearerFlow  '+ex.getStackTraceString());
                System.debug('JWTBearerFlow '+ex);
            }
        }
        return access_token;
    }
}

EinsteinOCRService Class which is used for reading the API Key and getting the access token.

public class EinsteinOCRService {
    
    public static FINAL STRING END_POINT = 'https://api.einstein.ai/v2/oauth2/token';
    public static String getAccessToken() {
        
        ContentVersion base64Content = [SELECT Title, VersionData 
                                        FROM ContentVersion
                                        WHERE Title='einstein_platform' OR 
                                        Title='predictive_services'
                                        ORDER BY Title 
                                        LIMIT 1];
        String keyContents = base64Content.VersionData.tostring();
        keyContents = keyContents.replace('-----BEGIN RSA PRIVATE KEY-----', '');
        keyContents = keyContents.replace('-----END RSA PRIVATE KEY-----', '');
        keyContents = keyContents.replace('\n', '');
        
        // Get a new token
        JWT jwt = new JWT('RS256');
        // jwt.cert = 'JWTCert'; 
        // Uncomment this if you used a Salesforce certificate to sign up for an Einstein Platform account
        jwt.pkcs8 = keyContents; // Comment this if you are using jwt.cert
        jwt.iss = 'developer.force.com';
        jwt.sub = 'youremail@gmail.com';
        jwt.aud =  END_POINT;
        jwt.exp = '3600';
        String access_token = JWTBearerFlow.getAccessToken(END_POINT, jwt);
        return access_token;    
    }
}

References

• JSON Class
• JSONGenerator Class
• Salesforce Help doc
• JWT.io