From 4fc698b0e06a5159b775fa4672c7c1e4ea9131e4 Mon Sep 17 00:00:00 2001
From: msukkari <michael.sukkarieh@mail.mcgill.ca>
Date: Tue, 14 Jan 2025 13:31:48 -0800
Subject: [PATCH 1/2] pipe tenant id to restful search interface

---
 internal/tenant/grpc.go        |  2 +-
 internal/tenant/rest_helper.go | 27 +++++++++++++++++++++++++++
 json/json.go                   |  7 +++++++
 3 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 internal/tenant/rest_helper.go

diff --git a/internal/tenant/grpc.go b/internal/tenant/grpc.go
index 774c216dd..d5c4f4086 100644
--- a/internal/tenant/grpc.go
+++ b/internal/tenant/grpc.go
@@ -19,7 +19,7 @@ import (
 
 const (
 	// headerKeyTenantID is the header key for the tenant ID.
-	headerKeyTenantID = "X-Sourcegraph-Tenant-ID"
+	headerKeyTenantID = "X-Tenant-ID"
 
 	// headerValueNoTenant indicates the request has no tenant.
 	headerValueNoTenant = "none"
diff --git a/internal/tenant/rest_helper.go b/internal/tenant/rest_helper.go
new file mode 100644
index 000000000..9166168f2
--- /dev/null
+++ b/internal/tenant/rest_helper.go
@@ -0,0 +1,27 @@
+package tenant
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net/http"
+
+	"github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func InjectTenantFromHeader(ctx context.Context, header http.Header) (context.Context, error) {
+	log.Printf("header: %v", header)
+	tenantID := header.Get(headerKeyTenantID)
+	log.Printf("tenantID: %s", tenantID)
+	if tenantID != "" {
+		tenant, err := tenanttype.Unmarshal(tenantID)
+		if err != nil {
+			return ctx, status.New(codes.InvalidArgument, fmt.Errorf("bad tenant value in header: %w", err).Error()).Err()
+		}
+
+		return tenanttype.WithTenant(ctx, tenant), nil
+	}
+	return ctx, nil
+}
diff --git a/json/json.go b/json/json.go
index 80b47348d..7936ddee4 100644
--- a/json/json.go
+++ b/json/json.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/sourcegraph/zoekt"
+	"github.com/sourcegraph/zoekt/internal/tenant"
 	"github.com/sourcegraph/zoekt/query"
 )
 
@@ -85,6 +86,12 @@ func (s *jsonSearcher) jsonSearch(w http.ResponseWriter, req *http.Request) {
 		defer cancel()
 	}
 
+	ctx, err = tenant.InjectTenantFromHeader(ctx, req.Header)
+	if err != nil {
+		jsonError(w, http.StatusBadRequest, err.Error())
+		return
+	}
+
 	if err := CalculateDefaultSearchLimits(ctx, q, s.Searcher, searchArgs.Opts); err != nil {
 		jsonError(w, http.StatusInternalServerError, err.Error())
 		return

From 01ad7a76fdc5f03ab3fd43553460c01d0782cb8f Mon Sep 17 00:00:00 2001
From: msukkari <michael.sukkarieh@mail.mcgill.ca>
Date: Tue, 14 Jan 2025 15:15:53 -0800
Subject: [PATCH 2/2] change tenant header convention

---
 internal/tenant/grpc.go        | 2 +-
 internal/tenant/rest_helper.go | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/internal/tenant/grpc.go b/internal/tenant/grpc.go
index d5c4f4086..774c216dd 100644
--- a/internal/tenant/grpc.go
+++ b/internal/tenant/grpc.go
@@ -19,7 +19,7 @@ import (
 
 const (
 	// headerKeyTenantID is the header key for the tenant ID.
-	headerKeyTenantID = "X-Tenant-ID"
+	headerKeyTenantID = "X-Sourcegraph-Tenant-ID"
 
 	// headerValueNoTenant indicates the request has no tenant.
 	headerValueNoTenant = "none"
diff --git a/internal/tenant/rest_helper.go b/internal/tenant/rest_helper.go
index 9166168f2..ad4076d07 100644
--- a/internal/tenant/rest_helper.go
+++ b/internal/tenant/rest_helper.go
@@ -13,7 +13,7 @@ import (
 
 func InjectTenantFromHeader(ctx context.Context, header http.Header) (context.Context, error) {
 	log.Printf("header: %v", header)
-	tenantID := header.Get(headerKeyTenantID)
+	tenantID := header.Get("X-TENANT-ID") // TODO: we don't use headerKeyTenantID here so we don't need to change it and potentially break future grpc changes
 	log.Printf("tenantID: %s", tenantID)
 	if tenantID != "" {
 		tenant, err := tenanttype.Unmarshal(tenantID)