feat: support customized http header in gateway

Author: wenfengwang

Makefile

@@ -30,6 +30,9 @@ release-npi-cli:
 docker-build:
 	docker buildx build --platform ${DOCKER_PLATFORM} -t npiai/npi:${IMAGE_TAG} . --push
 
+docker-build-gateway:
+	docker buildx build --platform linux/amd64 -t npiai/gateway:latest -f build/gateway.Dockerfile proto/go --push
+
 docker-build-base:
 	docker buildx build --platform ${DOCKER_PLATFORM} -t npiai/base:3.10 -f build/base.Dockerfile build --push
 

build/gateway.Dockerfile

@@ -0,0 +1,15 @@
+FROM --platform=$BUILDPLATFORM golang as builder
+
+LABEL maintainer="Wenfeng Wang <w@npi.ai>"
+
+COPY . /npiai
+
+WORKDIR /npiai
+
+RUN GOOS=linux GOARCH=amd64 go build -o gateway main.go
+
+FROM --platform=$BUILDPLATFORM ubuntu:22.04
+
+COPY --from=builder /npiai/gateway /npiai/gateway
+
+CMD ["/npiai/gateway"]

proto/go/test.go proto/go/main.go

@@ -3,12 +3,14 @@ package main
 import (
 	"context"
 	"flag"
-	"net/http"
-
+	"fmt"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/grpclog"
+	"net/http"
+	"os"
+	"strings"
 
 	gw "github.com/npi-ai/npi/proto/go/api" // Update
 )
@@ -23,17 +25,38 @@ func run() error {
 	ctx := context.Background()
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
-
+	headerMatcher := func(key string) (string, bool) {
+		switch strings.ToUpper(key) {
+		case "X-NPI-TOKEN", "USER-AGENT", "ACCEPT":
+			return key, true
+		default:
+			return key, false
+		}
+	}
 	// Register gRPC server endpoint
 	// Note: Make sure the gRPC server is running properly and accessible
-	mux := runtime.NewServeMux()
+	mux := runtime.NewServeMux(
+		runtime.WithIncomingHeaderMatcher(headerMatcher),
+	)
 	opts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
+	opts = append(opts, grpc.WithUnaryInterceptor(
+		func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+			// just for debug
+			return invoker(ctx, method, req, reply, cc, opts...)
+		}))
+
+	endpoint := os.Getenv("GRPC_SERVER_ENDPOINT")
+	if endpoint != "" && *grpcServerEndpoint == "localhost:9140" {
+		*grpcServerEndpoint = endpoint
+	}
+
 	err := gw.RegisterAppServerHandlerFromEndpoint(ctx, mux, *grpcServerEndpoint, opts)
 	if err != nil {
 		return err
 	}
 
 	// Start HTTP server (and proxy calls to gRPC server endpoint)
+	fmt.Printf("Starting HTTP server on port 8081, endpoint: %s", *grpcServerEndpoint)
 	return http.ListenAndServe(":8081", mux)
 }