An Apex implementation of the OAuth 2.0 JWT Bearer Token Flow

This flow allows an access token (AKA a session ID) to be obtained for a user based on a certificate shared by the client and the authorization server. Unlike most other OAuth 2.0 flows, no password is needed. This avoids having to prompt for a password in a browser or having to have a stored password. So the flow works well for server to server interactions.

There is a Java sample implementation in the OAuth 2.0 JWT Bearer Token Flow help. Here is an Apex implementation of that. Having this in Apex allows e.g. one org to connect to another or a Force.com Site to obtain a session ID. This code could also be used to establish a connection from Salesforce to some other platform that supports the flow, but details like the encryption algorithm used might need changing.

public class Jwt {
    
    public class Configuration {
        public String jwtUsername {get; set;}
        public String jwtConnectedAppConsumerKey {get; set;}
        public String jwtSigningCertificateName {get; set;}
        public String jwtHostname {get; set;}
    }
    
    private class Header {
        String alg;
        Header(String alg) {
            this.alg = alg;
        }
    }
    
    private class Body {
        String iss;
        String prn;
        String aud;
        String exp;
        Body(String iss, String prn, String aud, String exp) {
            this.iss = iss;
            this.prn = prn;
            this.aud = aud;
            this.exp = exp;
        }
    }
    
    private class JwtException extends Exception {
    }
    
    private Configuration config;
    
    public Jwt(Configuration config) {
        
        this.config = config;
    }
    
    public String requestAccessToken() {

         Map<String, String> m = new Map<String, String>();
         m.put('grant_type', 'urn:ietf:params:oauth:grant-type:jwt-bearer');
         m.put('assertion', createToken());
    
         HttpRequest req = new HttpRequest();
         req.setHeader('Content-Type','application/x-www-form-urlencoded');
         req.setEndpoint('https://' + config.jwtHostname +'/services/oauth2/token');
         req.setMethod('POST');
         req.setTimeout(60 * 1000);
         req.setBody(formEncode(m));
         
         HttpResponse res = new Http().send(req);
         if (res.getStatusCode() >= 200 && res.getStatusCode() < 300) {
             return extractJsonField(res.getBody(), 'access_token');
         } else {
             throw new JwtException(res.getBody());
         }
    }
    
    private String formEncode(Map<String, String> m) {
        
         String s = '';
         for (String key : m.keySet()) {
            if (s.length() > 0) {
                s += '&';
            }
            s += key + '=' + EncodingUtil.urlEncode(m.get(key), 'UTF-8');
         }
         return s;
    }
    
    private String extractJsonField(String body, String field) {
        
        JSONParser parser = JSON.createParser(body);
        while (parser.nextToken() != null) {
            if (parser.getCurrentToken() == JSONToken.FIELD_NAME
                    && parser.getText() == field) {
                parser.nextToken();
                return parser.getText();
            }
        }
        throw new JwtException(field + ' not found in response ' + body);
    }
    
    private String createToken() {
        
        String alg = 'RS256';
        
        String iss = config.jwtConnectedAppConsumerKey;
        String prn = config.jwtUsername;
        String aud = 'https://' + config.jwtHostname;
        String exp = String.valueOf(System.currentTimeMillis() + 60 * 60 * 1000);
        
        String headerJson = JSON.serialize(new Header(alg));
        String bodyJson =  JSON.serialize(new Body(iss, prn, aud, exp));
        
        String token = base64UrlSafe(Blob.valueOf(headerJson))
                + '.' + base64UrlSafe(Blob.valueOf(bodyJson));
        String signature = base64UrlSafe(Crypto.signWithCertificate(
                'RSA-SHA256',
                Blob.valueOf(token),
                config.jwtSigningCertificateName
                ));
        token += '.' + signature;
        
        return token;
    }
    
    private String base64UrlSafe(Blob b) {
        
        return EncodingUtil.base64Encode(b).replace('+', '-').replace('/', '_');
    }
}

And a test for the class. Note that the certificate has to be manually created in the org because there is no API that allows the test to automatically create it:

@isTest
private class JwtTest {
    
    // No API for createing a certificate so must be manually pre-created
    // using Setup -> Security Controls -> Certificate and Key Management
    private static final String PRE_CREATED_CERTIFICATE_NAME = 'JWT';
    private static final String FAKE_TOKEN = 'fakeToken';
    
    private class Mock implements HttpCalloutMock {

        public HTTPResponse respond(HTTPRequest req) {
            
            HTTPResponse res = new HTTPResponse();
            System.assertEquals('POST', req.getMethod());
            System.assert(req.getBody().contains('grant_type'), req.getBody());
            System.assert(req.getBody().contains('assertion'), req.getBody());
            
            res.setStatusCode(200);
            res.setBody('{"scope":"api","access_token":"' + FAKE_TOKEN + '"}');

            return res;
        }
    }

    @isTest
    static void test() {
        
        Jwt.Configuration config = new Jwt.Configuration();
        config.jwtUsername = 'abcdef@ghijkl.com';
        config.jwtSigningCertificateName = PRE_CREATED_CERTIFICATE_NAME;
        config.jwtHostname = 'login.salesforce.com';
        config.jwtConnectedAppConsumerKey = '6MVG9ZsNvTsRRnx.BZjJLCHB.hXYNAVb_oM';
        
        Test.setMock(HttpCalloutMock.class, new Mock());
        Test.startTest();
        String accessToken = new Jwt(config).requestAccessToken();
        Test.stopTest();
        System.assertEquals(FAKE_TOKEN, accessToken);
    }
}
Advertisements

3 thoughts on “An Apex implementation of the OAuth 2.0 JWT Bearer Token Flow

  1. Hey Keith,

    I’m trying to integrate this with Ideas (http://support.aha.io/hc/en-us/articles/203636345-Using-JSON-Web-Token-JWT-for-Idea-Portal-single-sign-on-SSO-), but i’m getting grant type not supported error and I redirect the call from https://mycommunity.ideas.aha.io/auth/jwt to https://test.salesforce.com/services/oauth2/token, but I get an error 200.
    In other words where should I redirect my Ideas community so I can then redirect it to the token ?

    Thanks !

  2. Thanks for this post, very helpful. In implementing a similar Apex JWT generator, I had to remove the = characters that were automatically appended as padding when using base64encode, since the JWT spec (and decoder we were using) doesn’t allow the = character.

    So:
    private String base64UrlSafe(Blob b) {
    return EncodingUtil.base64Encode(b).replace(‘+’, ‘-‘).replace(‘/’, ‘_’).remove(‘=’);
    }

  3. Bit pedantic but it looks like your exp is in milliseconds not seconds. I believe it should be the seconds since the UNIX epoch, i.e. divide by 1000.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s