MCPX MCP Gateway

package routing import ( "bytes" "fmt" "lunar/engine/actions" "lunar/engine/config" lunar_messages "lunar/engine/messages" "lunar/engine/runner" stream_config "lunar/engine/streams/config" lunar_context "lunar/engine/streams/lunar-context" stream_types "lunar/engine/streams/types" "lunar/engine/utils" context_manager "lunar/toolkit-core/context-manager" "lunar/toolkit-core/otel" "reflect" "github.com/negasus/haproxy-spoe-go/action" "github.com/negasus/haproxy-spoe-go/message" "github.com/negasus/haproxy-spoe-go/payload/kv" "github.com/negasus/haproxy-spoe-go/request" "github.com/rs/zerolog/log" ) type ( MessageHandler func(msgs *request.Request) MessageType int ) const ( requestType MessageType = iota responseType ) var sharedState = lunar_context.NewMemoryState[[]byte]() func Handler(data *HandlingDataManager) MessageHandler { data.RunDiagnosisWorker() ctxMng := context_manager.Get() handlerInner := func(req *request.Request) { var actions action.Actions if requestMessage, err := getMessageByType(requestType, req.Messages); err == nil { _, span := otel.Tracer(ctxMng.GetContext(), "routing#lunarOnRequestMessage") defer span.End() actions, err = processRequest(requestMessage, data) if ctxMng.GetContext().Err() != nil { actions = getShutdownActions() } else if err != nil { log.Error().Err(err).Msg("Error processing request") return } } if responseMessage, err := getMessageByType(responseType, req.Messages); err == nil { _, span := otel.Tracer(ctxMng.GetContext(), "routing#lunarOnResponseMessage") defer span.End() actions, err = processResponse(responseMessage, data) if err != nil { log.Error().Err(err).Msg("Error processing response") return } } req.Actions = actions } return handlerInner } func getMessageByType( messageType MessageType, incomingMessage *message.Messages, ) (msg *message.Message, err error) { switch messageType { case requestType: msg, err = incomingMessage.GetByName(lunar_messages.LunarRequest) if err != nil { msg, err = incomingMessage.GetByName(lunar_messages.LunarFullRequest) } return msg, err case responseType: msg, err = incomingMessage.GetByName(lunar_messages.LunarResponse) if err != nil { msg, err = incomingMessage.GetByName(lunar_messages.LunarFullResponse) } default: err = fmt.Errorf("unknown message type %v", messageType) } return msg, err } func getSPOEReqActions( args lunar_messages.OnRequest, lunarActions []actions.ReqLunarAction, ) action.Actions { var prioritizedAction actions.ReqLunarAction = &actions.NoOpAction{} for _, lunarAction := range lunarActions { lunarAction.EnsureRequestIsUpdated(&args) prioritizedAction = prioritizedAction.ReqPrioritize(lunarAction) } t := reflect.TypeOf(prioritizedAction) log.Trace().Msgf("Prioritized OnRequest action: %v", t.String()) prioritizedAction.EnsureRequestIsUpdated(&args) return prioritizedAction.ReqToSpoeActions() } func getShutdownActions() action.Actions { generateResp := action.Actions{} generateResp.SetVar(action.ScopeTransaction, actions.ReturnEarlyResponseActionName, true) generateResp.SetVar(action.ScopeTransaction, actions.StatusCodeActionName, 503) generateResp.SetVar(action.ScopeTransaction, actions.ResponseBodyActionName, []byte("Lunar Gateway is shutting down")) generateResp.SetVar(action.ScopeTransaction, actions.ResponseHeadersActionName, utils.DumpHeaders(map[string]string{})) return generateResp } func getSPOERespActions( args lunar_messages.OnResponse, lunarActions []actions.RespLunarAction, ) action.Actions { var prioritizedAction actions.RespLunarAction = &actions.NoOpAction{} for _, lunarAction := range lunarActions { lunarAction.EnsureResponseIsUpdated(&args) prioritizedAction = prioritizedAction.RespPrioritize(lunarAction) } t := reflect.TypeOf(prioritizedAction) log.Trace().Msgf("Prioritized OnResponse action: %v", t.String()) prioritizedAction.EnsureResponseIsUpdated(&args) return prioritizedAction.RespToSpoeActions() } func processRequest(msg *message.Message, data *HandlingDataManager) (action.Actions, error) { var err error var actions action.Actions args := readRequestArgs(msg) log.Trace().Msgf("On request args: %+v\n", args) if data.IsStreamsEnabled() { apiStream := stream_types.NewRequestAPIStream(args, sharedState) if args.IsFullRequest() { defer apiStream.StoreRequest() } data.GetMetricManager().UpdateMetricsProviderForAPICall(apiStream) flowActions := &stream_config.StreamActions{ Request: &stream_config.RequestStream{}, } if err = runner.RunFlow(data.stream, apiStream, flowActions); err == nil { actions = getSPOEReqActions(args, flowActions.Request.Actions) } data.GetMetricManager().UpdateMetricsProviderForFlow(data.stream) } else { // This is a patch for the legacy mode body parsing args.Body = bytes.NewBuffer(args.RawBody).String() policiesData := data.GetTxnPoliciesAccessor().GetTxnPoliciesData(config.TxnID(args.ID)) log.Trace().Msgf("On request policies: %+v\n", policiesData) actions, err = runner.DispatchOnRequest( args, &policiesData.EndpointPolicyTree, &policiesData.Config, data.policiesServices, data.diagnosisWorker, ) } log.Trace().Str("request-id", args.ID).Msg("On request finished") return actions, err } func processResponse(msg *message.Message, data *HandlingDataManager) (action.Actions, error) { var actions action.Actions var err error args := readResponseArgs(msg) log.Trace().Msgf("On response args: %+v\n", args) if data.IsStreamsEnabled() { apiStream := stream_types.NewResponseAPIStream(args, sharedState) if args.IsFullResponse() { defer apiStream.DiscardRequest() } data.GetMetricManager().UpdateMetricsProviderForAPICall(apiStream) flowActions := &stream_config.StreamActions{ Response: &stream_config.ResponseStream{}, } if err = runner.RunFlow(data.stream, apiStream, flowActions); err == nil { actions = getSPOERespActions(args, flowActions.Response.Actions) } data.GetMetricManager().UpdateMetricsProviderForFlow(data.stream) } else { // This is a patch for the legacy mode body parsing args.Body = bytes.NewBuffer(args.RawBody).String() policiesData := data.GetTxnPoliciesAccessor().GetTxnPoliciesData(config.TxnID(args.ID)) log.Trace().Msgf("On response policies: %+v\n", policiesData) actions, err = runner.DispatchOnResponse( args, &policiesData.EndpointPolicyTree, &policiesData.Config.Global, data.policiesServices, data.diagnosisWorker, ) } log.Trace().Str("response-id", args.ID).Msg("On response finished") return actions, err } func extractArg[T any](key string, arg *kv.KV) T { var res T rawValue, found := arg.Get(key) if !found { return res } if value, valid := rawValue.(T); !valid { log.Trace(). Msgf("Could not parse value %v (type: %T) from argument %v as type %T", rawValue, rawValue, key, res) } else { res = value } return res } func readRequestArgs(msg *message.Message) lunar_messages.OnRequest { onRequest := lunar_messages.OnRequest{LunarName: msg.Name} //nolint:exhaustruct onRequest.ID = extractArg[string]("id", msg.KV) onRequest.SequenceID = extractArg[string]("sequence_id", msg.KV) onRequest.Method = extractArg[string]("method", msg.KV) onRequest.Scheme = extractArg[string]("scheme", msg.KV) onRequest.URL = extractArg[string]("url", msg.KV) onRequest.Path = extractArg[string]("path", msg.KV) onRequest.Query = extractArg[string]("query", msg.KV) headerStr := extractArg[string]("headers", msg.KV) onRequest.Headers = utils.ParseHeaders(&headerStr) onRequest.RawBody = extractArg[[]byte]("body", msg.KV) onRequest.Time = context_manager.Get().GetClock().Now() return onRequest } func readResponseArgs(msg *message.Message) lunar_messages.OnResponse { onResponse := lunar_messages.OnResponse{LunarName: msg.Name} //nolint:exhaustruct onResponse.ID = extractArg[string]("id", msg.KV) onResponse.SequenceID = extractArg[string]("sequence_id", msg.KV) onResponse.Method = extractArg[string]("method", msg.KV) onResponse.URL = extractArg[string]("url", msg.KV) statusINT64 := extractArg[int64]("status", msg.KV) onResponse.Status = int(statusINT64) headerStr := extractArg[string]("headers", msg.KV) onResponse.Headers = utils.ParseHeaders(&headerStr) onResponse.RawBody = extractArg[[]byte]("body", msg.KV) onResponse.Time = context_manager.Get().GetClock().Now() return onResponse }