Add support for external auth providers in code search

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentServer.kt

@@ -136,6 +136,8 @@ interface CodyAgentServer {
   fun testing_ignore_overridePolicy(params: ContextFilters?): CompletableFuture<Null?>
   @JsonRequest("extension/reset")
   fun extension_reset(params: Null?): CompletableFuture<Null?>
+  @JsonRequest("internal/getAuthHeaders")
+  fun internal_getAuthHeaders(params: String): CompletableFuture<Map<String, String>>
 
   // =============
   // Notifications

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(

agent/src/agent.ts

@@ -14,6 +14,7 @@ import {
     currentAuthStatusAuthed,
     firstNonPendingAuthStatus,
     firstResultFromOperation,
+    getAuthHeaders,
     resolvedConfig,
     telemetryRecorder,
     waitUntilComplete,
@@ -1411,6 +1412,11 @@ export class Agent extends MessageHandler implements ExtensionClient {
             contextFiltersProvider.setTestingContextFilters(contextFilters)
             return null
         })
+
+        this.registerAuthenticatedRequest('internal/getAuthHeaders', async url => {
+            const config = await firstResultFromOperation(resolvedConfig)
+            return await getAuthHeaders(config.auth, new URL(url))
+        })
     }
 
     private pushPendingPromise(pendingPromise: Promise<unknown>): void {

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")

jetbrains/src/main/java/com/sourcegraph/find/FindPopupDialog.java

@@ -48,8 +48,6 @@ public FindPopupDialog(@Nullable Project project, JComponent myMainPanel) {
     this.setModal(false);
     // Prevent the dialog from being cancelable by any default behaviors
     myCancelAction.setEnabled(false);
-
-    super.show();
   }
 
   // Overwrite the createPeer function that is being used in the super constructor so that we can

jetbrains/src/main/java/com/sourcegraph/find/FindPopupPanel.java

@@ -10,18 +10,17 @@
 import com.intellij.openapi.ide.CopyPasteManager;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.Splitter;
+import com.intellij.openapi.util.Disposer;
 import com.intellij.ui.OnePixelSplitter;
 import com.intellij.ui.PopupBorder;
 import com.intellij.ui.jcef.JBCefApp;
 import com.intellij.util.ui.JBUI;
 import com.intellij.util.ui.components.BorderLayoutPanel;
 import com.sourcegraph.Icons;
+import com.sourcegraph.cody.auth.CodyAuthService;
 import com.sourcegraph.common.NotificationGroups;
 import com.sourcegraph.common.ui.DumbAwareEDTAction;
-import com.sourcegraph.find.browser.BrowserAndLoadingPanel;
-import com.sourcegraph.find.browser.JSToJavaBridgeRequestHandler;
-import com.sourcegraph.find.browser.JavaToJSBridge;
-import com.sourcegraph.find.browser.SourcegraphJBCefBrowser;
+import com.sourcegraph.find.browser.*;
 import java.awt.*;
 import java.awt.datatransfer.StringSelection;
 import java.util.Date;
@@ -35,6 +34,7 @@
  * href="https://sourcegraph.com/github.com/JetBrains/intellij-community/-/blob/platform/lang-impl/src/com/intellij/find/impl/FindPopupPanel.java">FindPopupPanel.java</a>
  */
 public class FindPopupPanel extends BorderLayoutPanel implements Disposable {
+  private final JavaToJSBridge javaToJSBridge;
   private final SourcegraphJBCefBrowser browser;
   private final PreviewPanel previewPanel;
   private final BrowserAndLoadingPanel browserAndLoadingPanel;
@@ -64,14 +64,24 @@ public FindPopupPanel(@NotNull Project project, @NotNull FindService findService
     browserAndLoadingPanel = new BrowserAndLoadingPanel(project);
     JSToJavaBridgeRequestHandler requestHandler =
         new JSToJavaBridgeRequestHandler(project, this, findService);
-    browser = JBCefApp.isSupported() ? new SourcegraphJBCefBrowser(requestHandler) : null;
+    String endpointUrl = CodyAuthService.getInstance(project).getEndpoint().getUrl();
+
+    browser =
+        JBCefApp.isSupported() ? new SourcegraphJBCefBrowser(requestHandler, endpointUrl) : null;
+
     if (browser == null) {
+      javaToJSBridge = null;
       showNoBrowserErrorNotification();
       Logger logger = Logger.getInstance(JSToJavaBridgeRequestHandler.class);
       logger.warn("JCEF browser is not supported!");
     } else {
+      String initJSCode = "window.initializeSourcegraph();";
+      JSToJavaBridge jsToJavaBridge = new JSToJavaBridge(browser, requestHandler, initJSCode);
+      javaToJSBridge = new JavaToJSBridge(browser);
+      Disposer.register(this, jsToJavaBridge);
       browserAndLoadingPanel.setBrowser(browser);
     }
+
     // The border is needed on macOS because without it, window and splitter resize don't work
     // because the JCEF
     // doesn't properly pass the mouse events to Swing.
@@ -139,7 +149,7 @@ public PreviewPanel getPreviewPanel() {
 
   @Nullable
   public JavaToJSBridge getJavaToJSBridge() {
-    return browser != null ? browser.getJavaToJSBridge() : null;
+    return javaToJSBridge;
   }
 
   public BrowserAndLoadingPanel.ConnectionAndAuthState getConnectionAndAuthState() {

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()) {
@@ -57,6 +71,7 @@ private void createOrShowPopup() {
       }
     } else {
       popup = new FindPopupDialog(project, mainPanel);
+      popup.show();
 
       // We add a manual listener to the global key handler since the editor component seems to work
       // around the

jetbrains/src/main/java/com/sourcegraph/find/browser/HttpSchemeHandler.java

@@ -1,9 +1,11 @@
 package com.sourcegraph.find.browser;
 
 import com.google.common.collect.ImmutableMap;
+import com.intellij.openapi.diagnostic.Logger;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.URL;
 import java.util.Map;
 import java.util.Optional;
 import org.cef.callback.CefCallback;
@@ -21,11 +23,19 @@ public class HttpSchemeHandler extends CefResourceHandlerAdapter {
   private int responseHeader = 400;
   private int offset = 0;
 
+  Logger logger = Logger.getInstance(HttpSchemeHandler.class);
+
   public boolean processRequest(@NotNull CefRequest request, @NotNull CefCallback callback) {
     String extension = getExtension(request.getURL());
     mimeType = getMimeType(extension);
     String url = request.getURL();
-    String path = url.replace("http://sourcegraph", "");
+    String path;
+    try {
+      path = new URL(url).getPath();
+    } catch (Exception ignored) {
+      logger.error("Failed to parse request url: " + url);
+      return false;
+    }
 
     if (mimeType != null) {
       data = loadResource(path);

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.

jetbrains/src/main/java/com/sourcegraph/find/browser/JavaToJSBridge.java

@@ -7,11 +7,13 @@
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.ui.jcef.JBCefBrowserBase;
 import com.intellij.ui.jcef.JBCefJSQuery;
+import com.sourcegraph.config.ThemeUtil;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Function;
+import javax.swing.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -25,6 +27,13 @@ public JavaToJSBridge(JBCefBrowserBase browser) {
     this.browser = browser;
     this.query = JBCefJSQuery.create(browser);
     this.lock = new ReentrantLock();
+
+    UIManager.addPropertyChangeListener(
+        propertyChangeEvent -> {
+          if (propertyChangeEvent.getPropertyName().equals("lookAndFeel")) {
+            this.callJS("themeChanged", ThemeUtil.getCurrentThemeAsJson());
+          }
+        });
   }
 
   public void callJS(@NotNull String action, @Nullable JsonObject arguments) {

jetbrains/src/main/java/com/sourcegraph/find/browser/SourcegraphJBCefBrowser.java

@@ -1,43 +1,26 @@
 package com.sourcegraph.find.browser;
 
-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;
+import org.cef.browser.CefBrowser;
+import org.cef.handler.CefLifeSpanHandlerAdapter;
 import org.jetbrains.annotations.NotNull;
 
 public class SourcegraphJBCefBrowser extends JBCefBrowser {
-  private final JavaToJSBridge javaToJSBridge;
+  public SourcegraphJBCefBrowser(
+      @NotNull JSToJavaBridgeRequestHandler requestHandler, String endpointUrl) {
+    super(endpointUrl.replaceAll("/+$", "").replace("https://", "http://") + "/html/index.html");
 
-  public SourcegraphJBCefBrowser(@NotNull JSToJavaBridgeRequestHandler requestHandler) {
-    super("http://sourcegraph/html/index.html");
-    // Create and set up JCEF browser
-    CefApp.getInstance()
-        .registerSchemeHandlerFactory("http", "sourcegraph", new HttpSchemeHandlerFactory());
-
-    // Create bridges, set up handlers, then run init function
-    String initJSCode = "window.initializeSourcegraph();";
-    JSToJavaBridge jsToJavaBridge = new JSToJavaBridge(this, requestHandler, initJSCode);
-    Disposer.register(this, jsToJavaBridge);
-    javaToJSBridge = new JavaToJSBridge(this);
-
-    requestHandler
-        .getProject()
-        .getService(CodySettingChangeListener.class)
-        .setJavaToJSBridge(javaToJSBridge);
-
-    UIManager.addPropertyChangeListener(
-        propertyChangeEvent -> {
-          if (propertyChangeEvent.getPropertyName().equals("lookAndFeel")) {
-            javaToJSBridge.callJS("themeChanged", ThemeUtil.getCurrentThemeAsJson());
+    // Schema registration need to happen in a callback, or it may crash in IJ 2023.2:
+    // https://youtrack.jetbrains.com/issue/JBR-5853
+    CefLifeSpanHandlerAdapter lifeSpanHandler =
+        new CefLifeSpanHandlerAdapter() {
+          public void onAfterCreated(CefBrowser browser) {
+            CefApp.getInstance()
+                .registerSchemeHandlerFactory("http", null, new HttpSchemeHandlerFactory());
           }
-        });
-  }
+        };
 
-  @NotNull
-  public JavaToJSBridge getJavaToJSBridge() {
-    return javaToJSBridge;
+    this.getJBCefClient().addLifeSpanHandler(lifeSpanHandler, this.getCefBrowser());
   }
 }

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 {

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() {

jetbrains/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) {

jetbrains/src/main/kotlin/com/sourcegraph/config/ConfigUtil.kt

@@ -11,6 +11,7 @@ import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.ProjectManager
+import com.sourcegraph.cody.agent.CodyAgentService
 import com.sourcegraph.cody.agent.protocol_generated.ExtensionConfiguration
 import com.sourcegraph.cody.auth.CodyAuthService
 import com.sourcegraph.cody.auth.CodySecureStore
@@ -21,6 +22,7 @@ import com.typesafe.config.ConfigRenderOptions
 import com.typesafe.config.ConfigValueFactory
 import java.nio.file.Path
 import java.nio.file.Paths
+import java.util.concurrent.CompletableFuture
 import kotlin.io.path.readText
 import org.jetbrains.annotations.Contract
 import org.jetbrains.annotations.VisibleForTesting
@@ -110,12 +112,32 @@ object ConfigUtil {
   }
 
   @JvmStatic
-  fun getConfigAsJson(project: Project): JsonObject {
+  fun getAuthorizationHeadersAsJson(project: Project): JsonObject {
     val endpoint = CodyAuthService.getInstance(project).getEndpoint()
+    val authHeaders = CompletableFuture<Map<String, String>>()
+
+    if (isCodyEnabled()) {
+      CodyAgentService.withAgent(project) { agent ->
+        agent.server.internal_getAuthHeaders(endpoint.url).thenAccept { headers ->
+          authHeaders.complete(headers)
+        }
+      }
+    } else {
+      val token = CodySecureStore.getFromSecureStore(endpoint.url)
+      if (token != null) {
+        authHeaders.complete(mapOf("Authorization" to "token $token"))
+      }
+    }
+
+    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("accessToken", CodySecureStore.getFromSecureStore(endpoint.url.toString()))
-      addProperty("customRequestHeadersAsString", "")
+      addProperty("instanceURL", CodyAuthService.getInstance(project).getEndpoint().url)
       addProperty("pluginVersion", getPluginVersion())
       addProperty("anonymousUserId", CodyApplicationSettings.instance.anonymousUserId)
     }