--- a/mod_http_oauth2/mod_http_oauth2.lua	Sat Apr 29 11:26:04 2023 +0200
+++ b/mod_http_oauth2/mod_http_oauth2.lua	Sat Apr 29 13:09:46 2023 +0200
@@ -17,6 +17,10 @@
 local array = require "util.array";
 local st = require "util.stanza";
 
+local function b64url(s)
+	return (s:gsub("[+/=]", { ["+"] = "-", ["/"] = "_", ["="] = "" }))
+end
+
 local function read_file(base_path, fn, required)
 	local f, err = io.open(base_path .. "/" .. fn);
 	if not f then
@@ -69,6 +73,8 @@
 local registration_algo = module:get_option_string("oauth2_registration_algorithm", "HS256");
 local registration_options = module:get_option("oauth2_registration_options", { default_ttl = 60 * 60 * 24 * 90 });
 
+local pkce_required = module:get_option_boolean("oauth2_require_code_challenge", false);
+
 local verification_key;
 local jwt_sign, jwt_verify;
 if registration_key then
@@ -211,6 +217,7 @@
 
 local grant_type_handlers = {};
 local response_type_handlers = {};
+local verifier_transforms = {};
 
 function grant_type_handlers.password(params)
 	local request_jid = assert(params.username, oauth_error("invalid_request", "missing 'username' (JID)"));
@@ -236,12 +243,18 @@
 	end
 	local granted_scopes, granted_role = filter_scopes(request_username, params.scope);
 
+	if pkce_required and not params.code_challenge then
+		return oauth_error("invalid_request", "PKCE required");
+	end
+
 	local code = id.medium();
 	local ok = codes:set(params.client_id .. "#" .. code, {
 		expires = os.time() + 600;
 		granted_jid = granted_jid;
 		granted_scopes = granted_scopes;
 		granted_role = granted_role;
+		challenge = params.code_challenge;
+		challenge_method = params.code_challenge_method;
 		id_token = id_token;
 	});
 	if not ok then
@@ -340,6 +353,14 @@
 		return oauth_error("invalid_client", "incorrect credentials");
 	end
 
+	-- TODO Decide if the code should be removed or not when PKCE fails
+	local transform = verifier_transforms[code.challenge_method or "plain"];
+	if not transform then
+		return oauth_error("invalid_request", "unknown challenge transform method");
+	elseif transform(params.code_verifier) ~= code.challenge then
+		return oauth_error("invalid_grant", "incorrect credentials");
+	end
+
 	return json.encode(new_access_token(code.granted_jid, code.granted_role, code.granted_scopes, client, code.id_token));
 end
 
@@ -371,6 +392,18 @@
 	));
 end
 
+-- RFC 7636 Proof Key for Code Exchange by OAuth Public Clients
+
+function verifier_transforms.plain(code_verifier)
+	-- code_challenge = code_verifier
+	return code_verifier;
+end
+
+function verifier_transforms.S256(code_verifier)
+	-- code_challenge = BASE64URL-ENCODE(SHA256(ASCII(code_verifier)))
+	return code_verifier and b64url(hashes.SHA256(code_verifier));
+end
+
 -- Used to issue/verify short-lived tokens for the authorization process below
 local new_user_token, verify_user_token = jwt.init("HS256", random.bytes(32), nil, { default_ttl = 600 });
 
@@ -903,6 +936,7 @@
 				registration_endpoint = handle_register_request and module:http_url() .. "/register" or nil;
 				scopes_supported = usermanager.get_all_roles and array(it.keys(usermanager.get_all_roles(module.host))):append(array(openid_claims:items()));
 				response_types_supported = array(it.keys(response_type_handlers));
+				code_challenge_methods_supported = array(it.keys(verifier_transforms));
 				authorization_response_iss_parameter_supported = true;
 
 				-- OpenID