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