package proxy import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "net/http/httputil" "strconv" "cloud.google.com/go/container/apiv1/containerpb" "github.com/aws/aws-sdk-go-v2/service/ec2" awstypes "github.com/aws/aws-sdk-go-v2/service/ec2/types" "github.com/gin-gonic/gin" "github.com/go-logr/logr" regionalpb "github.com/tidbcloud/infra-provider/api/tidbcloud/regional-server/v1" "github.com/tidbcloud/infra-provider/pkg/cloud/aws" "github.com/tidbcloud/infra-provider/pkg/config" "github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1" bizerr "github.com/tidbcloud/infra-provider/regional-server/pkg/common/errors" "github.com/tidbcloud/infra-provider/regional-server/pkg/common/gcpcontainer" "github.com/tidbcloud/infra-provider/regional-server/pkg/common/message" "github.com/tidbcloud/infra-provider/regional-server/pkg/common/utils" rsconfig "github.com/tidbcloud/infra-provider/regional-server/pkg/config" "github.com/tidbcloud/infra-provider/regional-server/pkg/resourcemanager/approute" "github.com/tidbcloud/infra-provider/regional-server/pkg/resourcemanager/o11yroute" "github.com/tidbcloud/infra-provider/regional-server/pkg/router/interceptor" agentcontrolapi "github.com/tidbcloud/observability/agent-control-server/pkg/agentcontrol/api" "google.golang.org/protobuf/encoding/protojson" "k8s.io/utils/pointer" "sigs.k8s.io/controller-runtime/pkg/client" ) type o11yModifyAction string const ( o11yModifyActionPause o11yModifyAction = "pause" o11yModifyActionResume o11yModifyAction = "resume" o11yModifyActionScale o11yModifyAction = "scale" ) type ClusterServiceHijackProxy struct { logger logr.Logger appRouteR approute.APPRouteReader o11yRouteRW o11yroute.O11yRouteReaderWriter agentControlAPI agentcontrolapi.AgentControlAPI rp *ReverseProxyWrapper cfg *rsconfig.RootConfig marshalOpt protojson.MarshalOptions unmarshalOpt protojson.UnmarshalOptions } func NewClusterServiceHijackProxy(rp *ReverseProxyWrapper, client client.Client, cfg *rsconfig.RootConfig, agentControlAPI agentcontrolapi.AgentControlAPI) (*ClusterServiceHijackProxy, error) { p := &ClusterServiceHijackProxy{ logger: logr.Discard(), rp: rp, cfg: cfg, appRouteR: approute.NewAPPRouteReader(client), o11yRouteRW: o11yroute.NewO11yRouteReaderWriter(client), agentControlAPI: agentControlAPI, marshalOpt: protojson.MarshalOptions{ UseProtoNames: false, EmitUnpopulated: true, }, unmarshalOpt: protojson.UnmarshalOptions{ DiscardUnknown: true, }, } return p, nil } func (p *ClusterServiceHijackProxy) WithLogger(logger logr.Logger) { p.logger = logger } func (p *ClusterServiceHijackProxy) GetCluster(c *gin.Context) { // check approute before proxy id := c.Param("clusterId") cluster := new(regionalpb.Cluster) cluster.Id = id cluster.ClusterPhase = regionalpb.Cluster_SYNCING approutes := new(v1alpha1.APPRouteList) labelselector := client.MatchingLabels{ v1alpha1.ClusterIDKey: id, } if err := p.appRouteR.List(c, approutes, labelselector); err != nil { p.logger.Error(err, "get approute cr failed", "cluster_id", id) interceptor.Fail(c, bizerr.New(bizerr.K8SClientError, err.Error())) return } // if the approute is not existed yet, just proxy to cluster service if len(approutes.Items) == 0 { p.rp.ServeHTTP(c.Writer, c.Request) return } approute := approutes.Items[0] switch approute.Status.Phase { case "", v1alpha1.APPRoutePhaseCreating, v1alpha1.APPRoutePhasePending, v1alpha1.APPRoutePhaseFailed: interceptor.SucceedPb(c, http.StatusOK, cluster) return case v1alpha1.APPRoutePhaseDeleting, v1alpha1.APPRoutePhaseActive: // do nothing, just proxy to cluster service, let cluster service return the real data. default: } p.rp.ServeHTTP(c.Writer, c.Request) } func (p *ClusterServiceHijackProxy) ListClusters(c *gin.Context) { tenantID := c.Query("tenantId") projectIDs := c.QueryArray("projectIds") approutes := new(v1alpha1.APPRouteList) selector := client.MatchingLabels{ utils.GenTenantLabelKey(tenantID): "true", } err := p.appRouteR.List(c, approutes, selector) if err != nil { p.logger.Error(err, "list approute cr failed") interceptor.Fail(c, bizerr.New(bizerr.K8SClientError, err.Error())) return } if len(approutes.Items) == 0 { p.rp.ServeHTTP(c.Writer, c.Request) return } match := func(approute *v1alpha1.APPRoute, keys []string) bool { if len(keys) == 0 { return true } for _, key := range keys { if _, ok := approute.GetLabels()[key]; ok { return true } } return false } keys := make([]string, 0, len(projectIDs)) for _, id := range projectIDs { keys = append(keys, utils.GenProjectLabelKey(id)) } notReadyClusters := make([]*regionalpb.Cluster, 0, len(approutes.Items)) for _, item := range approutes.Items { if ok := match(&item, keys); !ok { continue } notready := false switch item.Status.Phase { case "", v1alpha1.APPRoutePhaseCreating, v1alpha1.APPRoutePhasePending, v1alpha1.APPRoutePhaseFailed: notready = true default: } if clusterID, ok := item.GetLabels()[v1alpha1.ClusterIDKey]; ok && notready { notReadyClusters = append(notReadyClusters, ®ionalpb.Cluster{ Id: clusterID, ClusterPhase: regionalpb.Cluster_SYNCING, }) } } if len(notReadyClusters) == 0 { p.rp.ServeHTTP(c.Writer, c.Request) return } // modify response shrp := httputil.NewSingleHostReverseProxy(p.rp.target) shrp.Director = p.rp.Director shrp.ModifyResponse = p.genModifyResponse(notReadyClusters) //nolint:bodyclose shrp.ServeHTTP(c.Writer, c.Request) } func (p *ClusterServiceHijackProxy) genModifyResponse(notReadyClusters []*regionalpb.Cluster) func(*http.Response) error { return func(resp *http.Response) error { b, err := io.ReadAll(resp.Body) if err != nil { return err } err = resp.Body.Close() if err != nil { return err } data := new(regionalpb.ListClustersResponse) if err != p.unmarshalOpt.Unmarshal(b, data) { return err } existed := map[string]struct{}{} for _, cluster := range data.Clusters { existed[cluster.Id] = struct{}{} } for _, cluster := range notReadyClusters { if _, ok := existed[cluster.Id]; ok { continue } data.Clusters = append(data.Clusters, cluster) } nb, err := p.marshalOpt.Marshal(data) if err != nil { return err } body := io.NopCloser(bytes.NewReader(nb)) resp.Body = body resp.ContentLength = int64(len(nb)) resp.Header.Set("Content-Length", strconv.Itoa(len(nb))) return nil } } func (p *ClusterServiceHijackProxy) PauseCluster(c *gin.Context) { rp := p.rp.WithActionFormat() rp.ModifyResponse = p.modifyO11yStatus(func(o *v1alpha1.O11yRoute) { //nolint:bodyclose o.Spec.Suspend = true }, c, o11yModifyActionPause) rp.ServeHTTP(c.Writer, c.Request) } func (p *ClusterServiceHijackProxy) ResumeCluster(c *gin.Context) { rp := p.rp.WithActionFormat() rp.ModifyResponse = p.modifyO11yStatus(func(o *v1alpha1.O11yRoute) { //nolint:bodyclose o.Spec.Suspend = false }, c, o11yModifyActionResume) rp.ServeHTTP(c.Writer, c.Request) } func (p *ClusterServiceHijackProxy) ScaleCluster(c *gin.Context) { rp := p.rp.WithActionFormat() req := ®ionalpb.ScaleClusterRequest{} if ok := interceptor.HandleJsonPbRequestFromBody(c, req); !ok { p.logger.Error(nil, "failed to handle json pb request from body") return } var body []byte // the value read from request.Body is stored as BodyBytesKey // by c.ShouldBindBodyWith() in HandleJsonPbRequestFromBody function if copiedBody, ok := c.Get(gin.BodyBytesKey); ok { if cbb, ok := copiedBody.([]byte); ok { body = cbb } } // Since the Body is already consumed in HandleJsonPbRequestFromBody(), // we need to reconstruct a new Body from the raw body to retain it. c.Request.Body = io.NopCloser(bytes.NewBuffer(body)) rp.ModifyResponse = p.modifyO11yStatus(func(o *v1alpha1.O11yRoute) { //nolint:bodyclose if req.NodeSpec != nil && o.Spec.Agent.TargetClusterTopology != nil { switch req.NodeSpec.ComponentType { case regionalpb.ComponentType_PD: o.Spec.Agent.TargetClusterTopology.PD.Replicas = int(req.NodeSpec.Count) case regionalpb.ComponentType_TIDB: o.Spec.Agent.TargetClusterTopology.TiDB.Replicas = int(req.NodeSpec.Count) case regionalpb.ComponentType_TIKV: o.Spec.Agent.TargetClusterTopology.TiKV.Replicas = int(req.NodeSpec.Count) case regionalpb.ComponentType_TIFLASH: o.Spec.Agent.TargetClusterTopology.TiFlash = &v1alpha1.TiFlash{Replicas: int(req.NodeSpec.Count)} } } }, c, o11yModifyActionScale) rp.ServeHTTP(c.Writer, c.Request) } // not touch http.Response at all, no need to check bodyclose func (p *ClusterServiceHijackProxy) modifyO11yStatus(mutate func(*v1alpha1.O11yRoute), c *gin.Context, action o11yModifyAction) func(*http.Response) error { //nolint:bodyclose return p.modifyResponseIfSucceeded(func() error { id := c.Param("clusterId") if err := p.tryModifyO11yStatusUsingAgentControlAPI(c, id, action); err != nil { p.logger.Error(err, "failed to update o11y status using agent control api", "cluster_id", id) // do not return for now } o11yroutes := new(v1alpha1.O11yRouteList) labelselector := client.MatchingLabels{ v1alpha1.ClusterIDKey: id, } if err := p.o11yRouteRW.List(c, o11yroutes, labelselector); err != nil { p.logger.Error(err, "get o11yroute cr failed", "cluster_id", id) return err } if len(o11yroutes.Items) == 0 { return nil } o11yroute := o11yroutes.Items[0] mutate(&o11yroute) if err := p.o11yRouteRW.Update(c, &o11yroute); err != nil { p.logger.Error(err, "update o11yroute cr failed", "cluster_id", id) return err } return nil }) } func (p *ClusterServiceHijackProxy) tryModifyO11yStatusUsingAgentControlAPI(ctx context.Context, id string, action o11yModifyAction) error { if p.agentControlAPI == nil || p.cfg == nil { return nil } approutes := new(v1alpha1.APPRouteList) labelselector := client.MatchingLabels{ v1alpha1.ClusterIDKey: id, } if err := p.appRouteR.List(ctx, approutes, labelselector); err != nil { p.logger.Error(err, "get approute cr failed", "cluster_id", id) return err } if len(approutes.Items) == 0 { return nil } approute := approutes.Items[0] labels := approute.GetLabels() if len(labels) == 0 { return fmt.Errorf("approute %s/%s has no labels", approute.Namespace, approute.Name) } tenantID, hasTenant := labels[v1alpha1.TagTenantKey] projectID, hasProject := labels[v1alpha1.TagProjectKey] if !hasTenant || !hasProject { return fmt.Errorf("approute missing required labels: tenant/project") } if !config.FeatureEnabled(p.cfg.FeatureGates, config.O11yNewAgentController, tenantID, projectID) { return nil } switch action { case o11yModifyActionPause: req := &agentcontrolapi.PauseTiDBAgentRequest{ClusterID: id} if _, err := p.agentControlAPI.PauseTiDBAgent(ctx, req); err != nil { p.logger.Error(err, "pause tidb agent failed on agent control api", "cluster_id", id) return err } case o11yModifyActionResume: req := &agentcontrolapi.ResumeTiDBAgentRequest{ClusterID: id} if _, err := p.agentControlAPI.ResumeTiDBAgent(ctx, req); err != nil { p.logger.Error(err, "resume tidb agent failed on agent control api", "cluster_id", id) return err } case o11yModifyActionScale: req := &agentcontrolapi.ScaleTiDBAgentRequest{ClusterID: id} if _, err := p.agentControlAPI.ScaleTiDBAgent(ctx, req); err != nil { p.logger.Error(err, "scale tidb agent failed on agent control api", "cluster_id", id) return err } } return nil } func (p *ClusterServiceHijackProxy) modifyResponseIfSucceeded(action func() error) func(*http.Response) error { return func(resp *http.Response) error { if resp.StatusCode/100 == 2 { return action() } p.logger.Info("get response code from cluster server", "status_code", resp.StatusCode) return nil } } type O11yRegionalServiceGCPHijackProxy struct { logger logr.Logger rp *ReverseProxyWrapper // this client is only set on gcp region gcpContainerCli gcpcontainer.ClusterGetter } func (p *O11yRegionalServiceGCPHijackProxy) WithLogger(logger logr.Logger) { p.logger = logger } func NewO11yRegionlSericeGCPHijackProxy(rp *ReverseProxyWrapper, gcpContainerCli gcpcontainer.ClusterGetter) (*O11yRegionalServiceGCPHijackProxy, error) { p := &O11yRegionalServiceGCPHijackProxy{ logger: logr.Discard(), rp: rp, gcpContainerCli: gcpContainerCli, } return p, nil } // o11y-regional-server needs gke caData to access the dataplene k8s cluster when reconciling `cluster` CR (one lives in o11y regional kube-apiserver). // However the caData is not existed on `cluster` CR. // So we need to gain this information via gcloud api and inject it into request body before proxying the request. // Thus o11y-regional-server can save the caData on `cluster` and fully managed the o11y agents in dataplane. func (p *O11yRegionalServiceGCPHijackProxy) ImportAgent(c *gin.Context) { param := new(message.O11yAgent) if err := c.BindJSON(param); err != nil { interceptor.Fail(c, bizerr.New(bizerr.ParamParseError, err.Error())) return } var gkecluster string var project string var location string if gkecluster = param.AgentEnvironment.Dataplane.ClusterName; gkecluster == "" { interceptor.Fail(c, bizerr.New(bizerr.ParameterInvalid, "agentEnvironment.datapane.clusterName is empty")) return } if project = param.AgentEnvironment.Gcp.GcpProject; project == "" { interceptor.Fail(c, bizerr.New(bizerr.ParameterInvalid, "agentEnvironment.gcp.gcpProject is empty")) return } if location = param.Region; location == "" { interceptor.Fail(c, bizerr.New(bizerr.ParameterInvalid, "region is empty")) return } name := fmt.Sprintf("projects/%s/locations/%s/clusters/%s", project, location, gkecluster) resp, err := p.gcpContainerCli.GetCluster(c, &containerpb.GetClusterRequest{ Name: name, }) if err != nil { p.logger.Error(err, "get cluster from gcp container service failed", "name", name) interceptor.Fail(c, bizerr.New(bizerr.InternalServerError, err.Error())) return } if resp.MasterAuth != nil { param.AgentEnvironment.Dataplane.ClusterCAData = resp.MasterAuth.ClusterCaCertificate } bs, err := json.Marshal(param) if err != nil { p.logger.Error(err, "json marshal failed", "name", name) interceptor.Fail(c, bizerr.New(bizerr.ParamParseError, err.Error())) return } p.logger.Info("inject cadata", "content", string(bs), "length", len(bs)) shrp := httputil.NewSingleHostReverseProxy(p.rp.target) shrp.Director = func(r *http.Request) { r.Body = io.NopCloser(bytes.NewBuffer(bs)) r.Header.Set("Content-Length", strconv.Itoa(len(bs))) r.ContentLength = int64(len(bs)) p.rp.Director(r) } shrp.ServeHTTP(c.Writer, c.Request) } func (p *O11yRegionalServiceGCPHijackProxy) PauseReconcile(c *gin.Context) { p.rp.ServeHTTP(c.Writer, c.Request) } type O11yRegionalServiceAWSHijackProxy struct { logger logr.Logger rp *ReverseProxyWrapper // this client is only set on gcp region awsClient aws.Client } func (p *O11yRegionalServiceAWSHijackProxy) WithLogger(logger logr.Logger) { p.logger = logger } func NewO11yRegionlSericeAWSHijackProxy(rp *ReverseProxyWrapper, awsClient aws.Client) (*O11yRegionalServiceAWSHijackProxy, error) { p := &O11yRegionalServiceAWSHijackProxy{ logger: logr.Discard(), rp: rp, awsClient: awsClient, } return p, nil } // The `Name` tag need be existed on vpcendpoint to let o11y-regional-server reconcile the vpcendpoint on aws correctly, // because the o11y-regional-server uses `Name` tag to identify the vpcendpoint. // However, we found some vpcendpoint used by o11y doesn't have the `Name` tag due to legacy bugs. // So we need to check the `Name` tag and set it if it is not existed via aws `CreateTags` api // before proxing request to o11y-regional-server. func (p *O11yRegionalServiceAWSHijackProxy) ImportAgent(c *gin.Context) { bodyCopy := new(bytes.Buffer) _, err := io.Copy(bodyCopy, c.Request.Body) if err != nil { interceptor.Fail(c, bizerr.New(bizerr.InternalServerError, err.Error())) return } param := new(message.O11yAgent) err = json.Unmarshal(bodyCopy.Bytes(), param) if err != nil { p.logger.Error(err, "json marshal failed") interceptor.Fail(c, bizerr.New(bizerr.ParamParseError, err.Error())) return } if param.EndpointName == "" { interceptor.Fail(c, bizerr.New(bizerr.ParamParseError, "endpointName is empty")) return } if param.AgentEnvironment.Aws.Account == "" { interceptor.Fail(c, bizerr.New(bizerr.ParamParseError, "agentEnvironment.aws.account is empty")) return } // set name tag by calling aws api before proxy request. if vpcEndpointID := param.AgentEnvironment.Aws.VpcEndpointID; vpcEndpointID != "" { account := param.AgentEnvironment.Aws.Account endpointName := param.EndpointName assumeRole := new(aws.AssumedRole) assumeRole.RoleARN = fmt.Sprintf("arn:aws:iam::%s:role/full-manager-service-role", account) assumeRole.ExternalID = pointer.String("123") ec2cli := p.awsClient.EC2(assumeRole) nameTagFound := false // descreibe vpcendpoint first then set tag if necessary { input := new(ec2.DescribeVpcEndpointsInput) input.VpcEndpointIds = []string{vpcEndpointID} out, err := ec2cli.DescribeVpcEndpoints(c, input) if err != nil { p.logger.Error(err, "call aws describe vpcendpoint api failed", "param", input) interceptor.Fail(c, bizerr.New(bizerr.AWSClientError, err.Error())) return } if len(out.VpcEndpoints) == 0 { interceptor.Fail(c, bizerr.New(bizerr.AWSClientError, fmt.Sprintf("vpcendpoint %s is not found", vpcEndpointID))) return } vpcEndpoint := out.VpcEndpoints[0] for _, tag := range vpcEndpoint.Tags { if tag.Key != nil && *tag.Key == "Name" { nameTagFound = true break } } } if !nameTagFound { input := new(ec2.CreateTagsInput) input.Tags = []awstypes.Tag{ { Key: pointer.String("Name"), Value: pointer.String(endpointName), }, } input.Resources = []string{vpcEndpointID} _, err = ec2cli.CreateTags(c, input) if err != nil { p.logger.Error(err, "call aws create tags api failed", "param", input) interceptor.Fail(c, bizerr.New(bizerr.AWSClientError, err.Error())) return } } } shrp := httputil.NewSingleHostReverseProxy(p.rp.target) shrp.Director = func(r *http.Request) { // since we not edit the payload, no need to change content-length r.Body = io.NopCloser(bodyCopy) p.rp.Director(r) } shrp.ServeHTTP(c.Writer, c.Request) } func (p *O11yRegionalServiceAWSHijackProxy) PauseReconcile(c *gin.Context) { p.rp.ServeHTTP(c.Writer, c.Request) }