Handle authtokens in the go code same as ts

Author: rculbertson

modal-go/client.go

@@ -75,6 +75,10 @@ var client pb.ModalClientClient
 // The us-east client talks to the control plane; all other clients talk to input planes.
 var clients = map[string]pb.ModalClientClient{}
 
+// authToken is the auth token received from the control plane on the first request, and sent with all
+// subsequent requests to both the control plane and the input plane.
+var authToken string
+
 func init() {
 	var err error
 	defaultConfig, _ = readConfigFile()
@@ -126,6 +130,7 @@ func createClient(serverURL string) (*grpc.ClientConn, pb.ModalClientClient, err
 			grpc.MaxCallSendMsgSize(maxMessageSize),
 		),
 		grpc.WithChainUnaryInterceptor(
+			authTokenInterceptor(),
 			retryInterceptor(),
 			timeoutInterceptor(),
 		),
@@ -148,6 +153,36 @@ func clientContext(ctx context.Context) context.Context {
 	)
 }
 
+// We receive an auth token from the control plane on our first request. We then include that auth token in every
+// subsequent request to both the control plane and the input plane.
+func authTokenInterceptor() grpc.UnaryClientInterceptor {
+	return func(
+		ctx context.Context,
+		method string,
+		req, reply any,
+		cc *grpc.ClientConn,
+		inv grpc.UnaryInvoker,
+		opts ...grpc.CallOption,
+	) error {
+		var headers, trailers metadata.MD
+		// Add authToken to outgoing context if it's set
+		if authToken != "" {
+			ctx = metadata.AppendToOutgoingContext(ctx, "x-modal-auth-token", authToken)
+		}
+		opts = append(opts, grpc.Header(&headers), grpc.Trailer(&trailers))
+		err := inv(ctx, method, req, reply, cc, opts...)
+		// If we're talking to the control plane, and no auth token was sent, it will return one.
+		// The python server returns it in the trailers, the worker returns it in the headers.
+		if val, ok := headers["x-modal-auth-token"]; ok {
+			authToken = val[0]
+		} else if val, ok := trailers["x-modal-auth-token"]; ok {
+			authToken = val[0]
+		}
+
+		return err
+	}
+}
+
 func timeoutInterceptor() grpc.UnaryClientInterceptor {
 	return func(
 		ctx context.Context,

modal-go/function.go

@@ -10,8 +10,6 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/metadata"
 	"net/http"
 	"time"
 
@@ -49,13 +47,12 @@ func FunctionLookup(ctx context.Context, appName string, name string, options *L
 	}
 	ctx = clientContext(ctx)
 
-	var header, trailer metadata.MD
 	resp, err := client.FunctionGet(ctx, pb.FunctionGetRequest_builder{
 		AppName:         appName,
 		ObjectTag:       name,
 		Namespace:       pb.DeploymentNamespace_DEPLOYMENT_NAMESPACE_WORKSPACE,
 		EnvironmentName: environmentName(options.Environment),
-	}.Build(), grpc.Header(&header), grpc.Trailer(&trailer))
+	}.Build())
 
 	if status, ok := status.FromError(err); ok && status.Code() == codes.NotFound {
 		return nil, NotFoundError{fmt.Sprintf("function '%s/%s' not found", appName, name)}
@@ -64,15 +61,6 @@ func FunctionLookup(ctx context.Context, appName string, name string, options *L
 		return nil, err
 	}
 
-	// Attach x-modal-auth-token to all future requests.
-	authTokenArray := header.Get("x-modal-auth-token")
-	if len(authTokenArray) == 0 {
-		authTokenArray = trailer.Get("x-modal-auth-token")
-	}
-	if len(authTokenArray) > 0 {
-		authToken := authTokenArray[0]
-		ctx = metadata.AppendToOutgoingContext(ctx, "x-modal-auth-token", authToken)
-	}
 	var inputPlaneUrl *string
 	if meta := resp.GetHandleMetadata(); meta != nil {
 		if url := meta.GetInputPlaneUrl(); url != "" {

modal-js/src/client.ts

@@ -33,6 +33,9 @@ function authMiddleware(profile: Profile): ClientMiddleware {
       options.metadata.set("x-modal-auth-token", modalAuthToken);
     }
 
+    // We receive an auth token from the control plane on our first request. We then include that auth token in every
+    // subsequent request to both the control plane and the input plane. The python server returns it in the trailers,
+    // the worker returns it in the headers.
     const prevOnHeader = options.onHeader;
     options.onHeader = (header) => {
       const token = header.get("x-modal-auth-token");