Add support for external auth providers in code search

sourcegraph · Feb 7, 2025 · 4a47767 · 4a47767
1 parent 7bdd2d6
commit 4a47767
Show file tree

Hide file tree

Showing 10 changed files with 40 additions and 24 deletions.
diff --git a/agent/scripts/reverse-proxy.py b/agent/scripts/reverse-proxy.py
@@ -22,11 +22,12 @@ async def proxy_handler(request):
             del headers['Transfer-Encoding']
 
         # Use value of 'Authorization: Bearer' to fill 'X-Forwarded-User' and remove 'Authorization' header
-        if 'Authorization' in headers:
-            match = re.match('Bearer (.*)', headers['Authorization'])
-            if match:
-                headers['X-Forwarded-User'] = match.group(1)
-            del headers['Authorization']
+
+        match = re.match('Bearer (.*)', headers['Authorization'])
+        if match:
+            headers['X-Forwarded-User'] = match.group(1)
+            if 'Authorization' in headers:
+                del headers['Authorization']
 
         # Forward the request to target
         async with session.request(

diff --git a/jetbrains/build.gradle.kts b/jetbrains/build.gradle.kts
@@ -343,7 +343,9 @@ tasks {
       return destinationDir
     }
 
-    val sourcegraphDir = unzipCodeSearch()
+    val codeSearchDirOverride = System.getenv("CODE_SEARCH_DIR_OVERRIDE")
+    val sourcegraphDir: File =
+        if (codeSearchDirOverride != null) file(codeSearchDirOverride) else unzipCodeSearch()
     exec {
       workingDir(sourcegraphDir.toString())
       commandLine(*pnpmPath, "install", "--frozen-lockfile", "--fix-lockfile")

diff --git a/jetbrains/src/main/java/com/sourcegraph/find/FindPopupPanel.java b/jetbrains/src/main/java/com/sourcegraph/find/FindPopupPanel.java
@@ -67,7 +67,8 @@ public FindPopupPanel(@NotNull Project project, @NotNull FindService findService
         new JSToJavaBridgeRequestHandler(project, this, findService);
     String endpointUrl = CodyAuthService.getInstance(project).getEndpoint().getUrl();
 
-    browser = JBCefApp.isSupported() ? new SourcegraphJBCefBrowser(requestHandler, endpointUrl) : null;
+    browser =
+        JBCefApp.isSupported() ? new SourcegraphJBCefBrowser(requestHandler, endpointUrl) : null;
     if (browser == null) {
       showNoBrowserErrorNotification();
       Logger logger = Logger.getInstance(JSToJavaBridgeRequestHandler.class);

diff --git a/jetbrains/src/main/java/com/sourcegraph/find/FindService.java b/jetbrains/src/main/java/com/sourcegraph/find/FindService.java
@@ -11,6 +11,7 @@
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.wm.ex.WindowManagerEx;
 import com.intellij.util.ui.UIUtil;
+import com.sourcegraph.config.ConfigUtil;
 import com.sourcegraph.find.browser.BrowserAndLoadingPanel;
 import com.sourcegraph.find.browser.JavaToJSBridge;
 import java.awt.*;
@@ -31,6 +32,10 @@ public FindService(@NotNull Project project) {
     mainPanel = new FindPopupPanel(project, this);
   }
 
+  public static FindService getInstance(@NotNull Project project) {
+    return project.getService(FindService.class);
+  }
+
   public synchronized void showPopup() {
     createOrShowPopup();
   }
@@ -40,6 +45,15 @@ public void hidePopup() {
     hideMaterialUiOverlay();
   }
 
+  public void refreshConfiguration() {
+    JavaToJSBridge javaToJSBridge = mainPanel.getJavaToJSBridge();
+    if (javaToJSBridge != null) {
+      mainPanel
+          .getJavaToJSBridge()
+          .callJS("pluginSettingsChanged", ConfigUtil.getConfigAsJson(project));
+    }
+  }
+
   private void createOrShowPopup() {
     if (popup != null) {
       if (!popup.isVisible()) {

diff --git a/jetbrains/src/main/java/com/sourcegraph/find/browser/JSToJavaBridgeRequestHandler.java b/jetbrains/src/main/java/com/sourcegraph/find/browser/JSToJavaBridgeRequestHandler.java
@@ -46,6 +46,8 @@ public JBCefJSQuery.Response handle(@NotNull JsonObject request) {
         case "getTheme":
           JsonObject currentThemeAsJson = ThemeUtil.getCurrentThemeAsJson();
           return createSuccessResponse(currentThemeAsJson);
+        case "getCustomRequestHeaders":
+          return createSuccessResponse(ConfigUtil.getAuthorizationHeadersAsJson(project));
         case "indicateSearchError":
           arguments = request.getAsJsonObject("arguments");
           // This must run on EDT (Event Dispatch Thread) because it changes the UI.

diff --git a/jetbrains/src/main/java/com/sourcegraph/find/browser/SourcegraphJBCefBrowser.java b/jetbrains/src/main/java/com/sourcegraph/find/browser/SourcegraphJBCefBrowser.java
@@ -2,7 +2,6 @@
 
 import com.intellij.openapi.util.Disposer;
 import com.intellij.ui.jcef.JBCefBrowser;
-import com.sourcegraph.cody.config.notification.CodySettingChangeListener;
 import com.sourcegraph.config.ThemeUtil;
 import javax.swing.*;
 import org.cef.CefApp;
@@ -23,11 +22,6 @@ public SourcegraphJBCefBrowser(
     Disposer.register(this, jsToJavaBridge);
     javaToJSBridge = new JavaToJSBridge(this);
 
-    requestHandler
-        .getProject()
-        .getService(CodySettingChangeListener.class)
-        .setJavaToJSBridge(javaToJSBridge);
-
     UIManager.addPropertyChangeListener(
         propertyChangeEvent -> {
           if (propertyChangeEvent.getPropertyName().equals("lookAndFeel")) {

diff --git a/jetbrains/src/main/kotlin/com/sourcegraph/cody/auth/CodyAuthService.kt b/jetbrains/src/main/kotlin/com/sourcegraph/cody/auth/CodyAuthService.kt
@@ -4,6 +4,7 @@ import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.sourcegraph.config.ConfigUtil
+import com.sourcegraph.find.FindService
 
 @Service(Service.Level.PROJECT)
 class CodyAuthService(val project: Project) {
@@ -18,6 +19,7 @@ class CodyAuthService(val project: Project) {
 
   fun setActivated(isActivated: Boolean) {
     this.isActivated = isActivated
+    if (isActivated) FindService.getInstance(project).refreshConfiguration()
   }
 
   fun getEndpoint(): SourcegraphServerPath {

diff --git a/jetbrains/src/main/kotlin/com/sourcegraph/cody/config/notification/ChangeListener.kt b/jetbrains/src/main/kotlin/com/sourcegraph/cody/config/notification/ChangeListener.kt
@@ -4,11 +4,9 @@ import com.intellij.openapi.Disposable
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.project.Project
 import com.intellij.util.messages.MessageBusConnection
-import com.sourcegraph.find.browser.JavaToJSBridge
 
 abstract class ChangeListener(protected val project: Project) : Disposable {
   protected val connection: MessageBusConnection = project.messageBus.connect()
-  var javaToJSBridge: JavaToJSBridge? = null
   protected val logger = Logger.getInstance(ChangeListener::class.java)
 
   override fun dispose() {

diff --git a/...ins/src/main/kotlin/com/sourcegraph/cody/config/notification/CodySettingChangeListener.kt b/...ins/src/main/kotlin/com/sourcegraph/cody/config/notification/CodySettingChangeListener.kt
@@ -19,8 +19,6 @@ class CodySettingChangeListener(project: Project) : ChangeListener(project) {
         CodySettingChangeActionNotifier.TOPIC,
         object : CodySettingChangeActionNotifier {
           override fun afterAction(context: CodySettingChangeContext) {
-            // Notify JCEF about the config changes
-            javaToJSBridge?.callJS("pluginSettingsChanged", ConfigUtil.getConfigAsJson(project))
 
             if (context.oldCodyEnabled != context.newCodyEnabled) {
               if (context.newCodyEnabled) {

diff --git a/jetbrains/src/main/kotlin/com/sourcegraph/config/ConfigUtil.kt b/jetbrains/src/main/kotlin/com/sourcegraph/config/ConfigUtil.kt
@@ -111,21 +111,25 @@ object ConfigUtil {
   }
 
   @JvmStatic
-  fun getConfigAsJson(project: Project): JsonObject {
+  fun getAuthorizationHeadersAsJson(project: Project): JsonObject {
     val endpoint = CodyAuthService.getInstance(project).getEndpoint()
-    val authHeadersString = CompletableFuture<String>()
-
+    val authHeaders = CompletableFuture<Map<String, String>>()
     CodyAgentService.withAgent(project) { agent ->
       agent.server.internal_getAuthHeaders(endpoint.url).thenAccept { headers ->
-        val x = headers.map { (key, value) -> "$key,$value" }.joinToString(",")
-        authHeadersString.complete(x)
+        authHeaders.complete(headers)
       }
     }
 
+    val jsonObject = JsonObject()
+    authHeaders.get().forEach { (key, value) -> jsonObject.addProperty(key, value) }
+    return jsonObject
+  }
+
+  @JvmStatic
+  fun getConfigAsJson(project: Project): JsonObject {
     return JsonObject().apply {
-      addProperty("instanceURL", endpoint.url)
+      addProperty("instanceURL", CodyAuthService.getInstance(project).getEndpoint().url)
       addProperty("pluginVersion", getPluginVersion())
-      addProperty("customRequestHeadersAsString", authHeadersString.get())
       addProperty("anonymousUserId", CodyApplicationSettings.instance.anonymousUserId)
     }
   }