package rest import ( "bytes" "context" "encoding/json" "fmt" "io/ioutil" "os" "strconv" "strings" "github.com/davecgh/go-spew/spew" "github.com/elastic/go-elasticsearch/v7" "github.com/elastic/go-elasticsearch/v7/esapi" "github.com/go-chi/jwtauth" "go.uber.org/zap" ) type ( // EsSearchAggrTerms is aggregations parameter for es search api. EsSearchAggrTerms map[string]esSearchAggr EsSearchNestedAggrTerms map[string]esSearchNestedAggr EsSearchHighlight struct { RequireFieldMatch bool `json:"require_field_match"` Fields map[string]highlightSetting `json:"fields"` } esSearchParamsIndex struct { Prefix struct { Index struct { Value string `json:"value"` } `json:"_index"` } `json:"prefix"` } esSimpleQueryString struct { Wrap struct { Query string `json:"query"` } `json:"simple_query_string"` } esDisMax struct { TieBreaker float64 `json:"tie_breaker,omitempty"` Boost float64 `json:"boost,omitempty"` Queries []interface{} `json:"queries,omitempty"` } esDisMaxWrap struct { Wrap esDisMax `json:"dis_max,omitempty"` } esNestedDisMax struct { Wrap struct { IgnoreUnmapped bool `json:"ignore_unmapped"` Path string `json:"path,omitempty"` Query esDisMaxWrap `json:"query,omitempty"` } `json:"nested,omitempty"` } esMultiMatch struct { Wrap struct { Query string `json:"query"` Type string `json:"type"` // Operator string `json:"operator"` Fields []string `json:"fields"` MinimumShouldMatch string `json:"minimum_should_match"` } `json:"multi_match"` } VectorValue struct { Vector []float32 `json:"vector"` K int `json:"k"` Filter []interface{} `json:"filter,omitempty"` } esSearchParams struct { Query struct { Bool *struct { // query context Must []interface{} `json:"must,omitempty"` Should []interface{} `json:"should,omitempty"` // filter context Filter []interface{} `json:"filter,omitempty"` MustNot []interface{} `json:"must_not,omitempty"` } `json:"bool,omitempty"` } `json:"query"` Aggregations EsSearchNestedAggrTerms `json:"aggs,omitempty"` Source SourceFilter `json:"_source,omitempty"` MinScore float32 `json:"min_score,omitempty"` Highlight EsSearchHighlight `json:"highlight,omitempty"` } esSearchAggrTerm struct { Field string `json:"field,omitempty"` Size int `json:"size,omitempty"` } esSearchAggrComposite struct { Sources interface{} `json:"sources"` // it can be esSearchAggrTerm,.. (Histogram, Date histogram, GeoTile grid) Size int `json:"size,omitempty"` } SourceFilter struct { Excludes []string `json:"excludes"` //columns to exclude on the query response } esSearchNestedAggs struct { Path string `json:"path,omitempty"` } esSearchAggr struct { Nested esSearchNestedAggs `json:"nested,omitempty"` Terms esSearchAggrTerm `json:"terms,omitempty"` Aggregations EsSearchAggrTerms `json:"aggs,omitempty"` // Composite *esSearchAggrComposite `json:"composite"` } esSearchNestedAggrTerm struct { Terms esSearchAggrTerm `json:"terms,omitempty"` } esSearchNestedAggr struct { Nested esSearchNestedAggs `json:"nested,omitempty"` Aggregations map[string]esSearchNestedAggrTerm `json:"aggs,omitempty"` // Composite *esSearchAggrComposite `json:"composite"` } esSearchResponse struct { Took int `json:"took"` TimedOut bool `json:"timed_out"` Hits esSearchHits `json:"hits"` Aggregations esSearchAggregations `json:"aggregations"` } esSearchTotal struct { Value int `json:"value"` Relation string `json:"relation"` } esSearchHits struct { Total esSearchTotal `json:"total"` Hits []*esSearchHit `json:"hits"` } esSearchHit struct { Index string `json:"_index"` ID string `json:"_id"` Source json.RawMessage `json:"_source"` Highlight map[string]any `json:"highlight,omitempty"` } esSearchAggregations struct { Resource struct { DocCountErrorUpperBound int `json:"-"` SumOtherDocCount int `json:"-"` Buckets []struct { Key string `json:"key"` DocCount int `json:"doc_count"` ResourceName struct { DocCountErrorUpperBound int `json:"-"` SumOtherDocCount int `json:"-"` Buckets []struct { Key string `json:"key"` DocCount int `json:"doc_count"` } `json:"buckets"` } `json:"resourceName"` Namespaces struct { DocCountErrorUpperBound int `json:"-"` SumOtherDocCount int `json:"-"` Buckets []struct { Key string `json:"key"` DocCount int `json:"doc_count"` } `json:"buckets"` } `json:"namespaces"` Modules struct { DocCountErrorUpperBound int `json:"-"` SumOtherDocCount int `json:"-"` Buckets []struct { Key string `json:"key"` DocCount int `json:"doc_count"` } `json:"buckets"` } `json:"modules"` } `json:"buckets"` } `json:"resource"` Module struct { DocCount int `json:"doc_count"` Module struct { DocCountErrorUpperBound int `json:"doc_count_error_upper_bound"` SumOtherDocCount int `json:"sum_other_doc_count"` Buckets []struct { Key string `json:"key"` DocCount int `json:"doc_count"` } `json:"buckets"` } `json:"module"` } `json:"module"` Namespace struct { DocCount int `json:"doc_count"` Namespace struct { DocCountErrorUpperBound int `json:"doc_count_error_upper_bound"` SumOtherDocCount int `json:"sum_other_doc_count"` Buckets []struct { Key string `json:"key"` DocCount int `json:"doc_count"` } `json:"buckets"` } `json:"namespace"` } `json:"namespace"` } searchParams struct { title string query string moduleAggs []string namespaceAggs []string dumpRaw bool from int size int resourceType []string aggOnly bool mAggOnly bool allowedRoles map[interface{}]bool searchMode string } highlightSetting struct { PreTags []string `json:"pre_tags"` PostTags []string `json:"post_tags"` NumberOfFragments int `json:"number_of_fragments"` } ) var ( rtIndexMapping = map[string]string{ "compose:module": "compose-modules", "compose:record": "compose-records-*", "compose:namespace": "compose-namespaces", "system:user": "system-users", } ) func esSearch(ctx context.Context, log *zap.Logger, esc *elasticsearch.Client, p searchParams) (sr *esSearchResponse, page pagination, err error) { var ( buf bytes.Buffer roles []string userID uint64 _, claims, _ = jwtauth.FromContext(ctx) allowedRoleExist = false ) if _, has := claims["roles"]; has { if rr, is := claims["roles"].([]interface{}); is { for _, role := range rr { if _, ok := p.allowedRoles[role]; ok { allowedRoleExist = true } roles = append(roles, fmt.Sprintf("%s", role)) } } } if _, has := claims["sub"]; has { if sub, is := claims["sub"].(string); is { userID, _ = extractFromSubClaim(sub) } } noQ := len(p.query) == 0 noNSFilter := len(p.namespaceAggs) == 0 // noMFilter := len(p.moduleAggs) == 0 sqs := esSimpleQueryString{} sqs.Wrap.Query = strings.TrimSpace(strings.ToLower(p.query)) query := &esSearchParams{} index := esSearchParamsIndex{} if query.Query.Bool == nil { query.Query.Bool = &struct { Must []interface{} `json:"must,omitempty"` Should []interface{} `json:"should,omitempty"` Filter []interface{} `json:"filter,omitempty"` MustNot []interface{} `json:"must_not,omitempty"` }{} } // Decide what indexes we can use if userID == 0 { // Missing, invalid, expired access token (JWT) index.Prefix.Index.Value = "corteza-public-" } else { // Authenticated user index.Prefix.Index.Value = "corteza-private-" if !allowedRoleExist { // Skip all documents that do not have baring roles in to allow list query.Query.Bool.Filter = append(query.Query.Bool.Filter, map[string]map[string]interface{}{ "terms": {"security.allowedRoles": roles}, }) // Skip all documents that have baring roles in to deny list query.Query.Bool.MustNot = append(query.Query.Bool.MustNot, map[string]map[string]interface{}{ "terms": {"security.deniedRoles": roles}, }) } } // Aggregations V1.0 // if len(p.aggregations) > 0 { // query.Aggregations = make(map[string]esSearchAggr) // // for _, a := range p.aggregations { // query.Aggregations[a] = esSearchAggr{esSearchAggrTerm{Field: a + ".keyword"}} // } // } // Search string filter if !noQ { applyTraditionalSearch(query, index, sqs.Wrap.Query) } // add the highlight for the matched fields. query.Highlight = EsSearchHighlight{ RequireFieldMatch: false, Fields: map[string]highlightSetting{ "values.*": { PreTags: []string{"*"}, PostTags: []string{"*"}, NumberOfFragments: 0, }, }, } var ( mm = esMultiMatch{} mdd esNestedDisMax nsdd esNestedDisMax ) for _, mAggs := range p.moduleAggs { mm.Wrap.Query = mAggs mm.Wrap.Type = "cross_fields" mm.Wrap.MinimumShouldMatch = "100%" mm.Wrap.Fields = []string{"module.name"} // query.Query.Bool.Must = append(query.Query.Bool.Must, mm) // query.Query.DisMax.Queries = append(query.Query.DisMax.Queries, mm) // dd.Wrap.Queries = append(dd.Wrap.Queries, mm) mdd.Wrap.Path = "module" nsdd.Wrap.IgnoreUnmapped = true mdd.Wrap.Query.Wrap.Queries = append(mdd.Wrap.Query.Wrap.Queries, mm) } if len(mdd.Wrap.Query.Wrap.Queries) > 0 { query.Query.Bool.Must = append(query.Query.Bool.Must, mdd) } // no need now since we are adding below as filter // if p.aggOnly { for _, nAggs := range p.namespaceAggs { mm.Wrap.Query = nAggs mm.Wrap.Type = "cross_fields" mm.Wrap.MinimumShouldMatch = "100%" mm.Wrap.Fields = []string{"namespace.name"} // query.Query.Bool.Must = append(query.Query.Bool.Must, mm) // query.Query.DisMax.Queries = append(query.Query.DisMax.Queries, mm) // dd.Wrap.Queries = append(dd.Wrap.Queries, mm) nsdd.Wrap.Path = "namespace" nsdd.Wrap.IgnoreUnmapped = true nsdd.Wrap.Query.Wrap.Queries = append(nsdd.Wrap.Query.Wrap.Queries, mm) } // } if len(nsdd.Wrap.Query.Wrap.Queries) > 0 { query.Query.Bool.Must = append(query.Query.Bool.Must, nsdd) } // if !p.aggOnly && !noNSFilter { // nsf := make(map[string]interface{}) // nsf["terms"] = map[string][]string{ // "namespace.name.keyword": p.namespaceAggs, // } // query.Query.Bool.Filter = append(query.Query.Bool.Filter, nsf) // } // Aggregations V1.0 Improved // if len(p.aggregations) > 0 { // for _, a := range p.aggregations { // if len(a) > 0 { // sqs = esSimpleQueryString{} // sqs.Wrap.Query = a // query.Query.Bool.Must = append(query.Query.Bool.Must, sqs) // } // } // } // if noQ == 0 && len(p.moduleAggs) == 0 && len(p.namespaceAggs) == 0 { // query.Query.DisMax.Queries = append(query.Query.DisMax.Queries, index) // } query.Aggregations = make(map[string]esSearchNestedAggr) query.Aggregations["namespace"] = esSearchNestedAggr{ Nested: esSearchNestedAggs{Path: "namespace"}, Aggregations: map[string]esSearchNestedAggrTerm{ "namespace": { Terms: esSearchAggrTerm{ // Field: "namespace.name.keyword", Field: "namespace.handle", Size: 999, }, }, }, } query.Source = SourceFilter{ Excludes: []string{"vectorsValue"}, } if !noQ || !noNSFilter { query.Aggregations["module"] = esSearchNestedAggr{ Nested: esSearchNestedAggs{Path: "module"}, Aggregations: map[string]esSearchNestedAggrTerm{ "module": { Terms: esSearchAggrTerm{ // Field: "module.name.keyword", Field: "module.handle", Size: 999, }, }, }, } } // query.Aggregations["resource"] = esSearchAggr{ // Terms: esSearchAggrTerm{ // Field: "resourceType.keyword", // Size: 999, // }, // Aggregations: EsSearchAggrTerms{ // "resourceName": esSearchAggr{ // Terms: esSearchAggrTerm{ // Field: "name.keyword", // Size: 999, // }, // }, // "modules": esSearchAggr{ // Terms: esSearchAggrTerm{ // Field: "module.name.keyword", // Size: 999, // }, // }, // "namespaces": esSearchAggr{ // Terms: esSearchAggrTerm{ // Field: "namespace.name.keyword", // Size: 999, // }, // }, // }, // } // Aggregations V2.0 // if len(p.aggregations) > 0 { // query.Aggregations = (Aggregations{}).encodeTerms(p.aggregations) // } // spew.Dump("query: ", query) if err = json.NewEncoder(&buf).Encode(query); err != nil { err = fmt.Errorf("could not encode query: %q", err) return } spew.Dump("buf: ", buf) log.Debug("searching ", zap.String("for", p.title), // zap.String("open search body", buf.String()), ) // Why set size to 999? default value for size is 10, // so we needed to set value till we add (@todo) pagination to search result if p.from > 0 && p.from >= p.size { p.size = p.from + 100 } if p.size == 0 { p.size = 999 } page.From = p.from page.Size = p.size // In case of requesting specific resource types, use only the appropriate indexes indexes := []string{} if len(p.resourceType) > 0 { for _, rt := range p.resourceType { if _, ok := rtIndexMapping[rt]; !ok { err = fmt.Errorf("invalid resource type: %s", rt) return } indexes = append(indexes, fmt.Sprintf("%s%s", index.Prefix.Index.Value, rtIndexMapping[rt])) } } sReqArgs := []func(*esapi.SearchRequest){ esc.Search.WithContext(ctx), esc.Search.WithBody(&buf), esc.Search.WithTrackTotalHits(true), // esc.Search.WithScroll(), esc.Search.WithSize(p.size), esc.Search.WithFrom(p.from), esc.Search.WithIndex(indexes...), // esc.Search.WithExplain(true), // debug } if p.dumpRaw { sReqArgs = append( sReqArgs, esc.Search.WithSourceExcludes("security"), esc.Search.WithPretty(), ) } // Perform the search request. res, err := esc.Search(sReqArgs...) if err != nil { return } if err = validElasticResponse(res, err); err != nil { err = fmt.Errorf("invalid search response: %w", err) return } defer res.Body.Close() if p.dumpRaw { // Copy body buf and then restore it bodyBytes, _ := ioutil.ReadAll(res.Body) res.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes)) os.Stdout.Write(bodyBytes) } sr = &esSearchResponse{} if err = json.NewDecoder(res.Body).Decode(sr); err != nil { return } // Print the response status, number of results, and request duration. log.Debug("search completed", zap.String("for", p.title), zap.String("query", sqs.Wrap.Query), zap.String("indexPrefix", index.Prefix.Index.Value), zap.Int("from", p.from), zap.Int("size", p.size), zap.String("status", res.Status()), zap.Int("took", sr.Took), zap.Bool("timedOut", sr.TimedOut), zap.Int("hits", sr.Hits.Total.Value), zap.String("hitsRelation", sr.Hits.Total.Relation), zap.Int("namespaceAggs", len(sr.Aggregations.Namespace.Namespace.Buckets)), zap.Int("moduleAggs", len(sr.Aggregations.Module.Module.Buckets)), ) return } // @todo move this to es service func validElasticResponse(res *esapi.Response, err error) error { if err != nil { return fmt.Errorf("failed to get response from search backend: %w", err) } if res.IsError() { defer res.Body.Close() var rsp struct { Error struct { Type string Reason string } } if err := json.NewDecoder(res.Body).Decode(&rsp); err != nil { return fmt.Errorf("could not parse response body: %w", err) } else { return fmt.Errorf("search backend responded with an error: %s (type: %s, status: %s)", rsp.Error.Reason, rsp.Error.Type, res.Status()) } } return nil } func extractFromSubClaim(sub string) (userID uint64, rr []uint64) { parts := strings.Split(sub, " ") rr = make([]uint64, len(parts)-1) for p := range parts { id, _ := strconv.ParseUint(parts[p], 10, 64) if p == 0 { userID = id } else { rr[p-1] = id } } return } func applyTraditionalSearch(query *esSearchParams, index interface{}, searchQuery string) { query.Query.Bool.Must = []interface{}{index} query.Query.Bool.Must = append(query.Query.Bool.Must, buildWildcardQuery(searchQuery)) } func buildWildcardQuery(searchQuery string) map[string]interface{} { return map[string]interface{}{ "wildcard": map[string]interface{}{ "catch_all": map[string]interface{}{ "value": fmt.Sprintf("*%s*", searchQuery), }, }, } } func buildKNNQuery(vector interface{}) map[string]interface{} { return map[string]interface{}{ "knn": map[string]interface{}{ "vectorsValue": map[string]interface{}{ "vector": vector, "k": 10, }, }, } } func buildMultiMatchQuery(searchQuery string) map[string]interface{} { return map[string]interface{}{ "multi_match": map[string]interface{}{ "query": searchQuery, "fields": []string{"catch_all"}, "boost": 1.0, }, } }