[Hadron] Always run databricks auth hook (#27) (#695)

Change to always allow auth hook to run
because we want to reject password based login
for databricks identities.

Corresponding hadron PR for CI:
https://github.com/databricks-eng/hadron/pull/752

Co-authored-by: Vikas Jain <[email protected]>

Author: dimitri

src/backend/libpq/auth.c

@@ -791,32 +791,59 @@ CheckPasswordAuth(Port *port, char **logdetail)
 	int			result;
 	char	   *shadow_pass;
 
+	/* BEGIN NEON */
+
+	/*
+	 * This flag is passed to databricks auth hook and is updated by the hook
+	 * to false if we should continue with password auth. This is by default
+	 * true so that we don't accidentally do password auth if there is some
+	 * bug in the hook. It's better to rely on the hook to set it explicitly
+	 * false to continue with password auth.
+	 */
+	bool		skip_password_auth = true;
+
+	/* END NEON */
+
 	sendAuthRequest(port, AUTH_REQ_PASSWORD, NULL, 0);
 
 	passwd = recv_password_packet(port);
 	if (passwd == NULL)
 		return STATUS_EOF;		/* client wouldn't send password */
 
-	shadow_pass = get_role_password(port->user_name, logdetail);
-	if (shadow_pass)
+	/* BEGIN NEON */
+	elog(DEBUG1, "Databricks: before authentication hook");
+
+	if (DatabricksAuthentication_hook)
 	{
-		result = plain_crypt_verify(port->user_name, shadow_pass, passwd,
-									logdetail);
+		result = (*DatabricksAuthentication_hook)
+			(port, passwd, &skip_password_auth, (const char **) logdetail);
 	}
 	else
+	{
+		/* If hook is not set, do the password auth by default */
+		skip_password_auth = false;
 		result = STATUS_ERROR;
+	}
 
-	if (result != STATUS_OK && DatabricksAuthentication_hook)
-	{
-		elog(LOG, "Calling DatabricksAuthentication_hook");
+	elog(DEBUG1, "Databricks: after authentication hook");
 
-		result = (*DatabricksAuthentication_hook)(port, passwd);
+	/* only try PG password auth if the hook didn't return STATUS_OK and */
+	/* the hook set the skip_password_auth flag to false */
+	if (result != STATUS_OK && !skip_password_auth)
+	{
+		shadow_pass = get_role_password(port->user_name, logdetail);
+		if (shadow_pass)
+		{
+			result = plain_crypt_verify(port->user_name, shadow_pass, passwd,
+										logdetail);
+		}
+		else
+			result = STATUS_ERROR;
 
-		elog(LOG, "DatabricksAuthentication_hook returned: %d", result);
+		if (shadow_pass)
+			pfree(shadow_pass);
 	}
-
-	if (shadow_pass)
-		pfree(shadow_pass);
+	/* END NEON */
 	pfree(passwd);
 
 	if (result == STATUS_OK)

src/include/libpq/auth.h

@@ -32,8 +32,14 @@ typedef char *(*auth_password_hook_typ) (char *input);
 /* Default LDAP password mutator hook, can be overridden by a shared library */
 extern PGDLLIMPORT auth_password_hook_typ ldap_password_hook;
 
-/* Hook for databricks authentication */
-typedef int (*DatabricksAuthentication_hook_type) (Port *, char *);
+/* Hook for databricks authentication
+ * returns STATUS_OK on success, STATUS_ERROR on failure
+ * skip_passwd_auth is set to true/false if password authentication should be tried or not on STATUS_ERROR
+ * */
+typedef int (*DatabricksAuthentication_hook_type) (Port *port,
+												   const char *passwd,
+												   bool *skip_passwd_auth,
+												   const char **logdetail);
 extern PGDLLIMPORT DatabricksAuthentication_hook_type DatabricksAuthentication_hook;
 
 #endif							/* AUTH_H */