package elasticsearch import ( "bytes" "context" "encoding/json" "fmt" "io" "strings" "sync" "time" "github.com/elastic/go-elasticsearch/v8" "tm/pkg/logger" ) type client struct { es *elasticsearch.Client config Config logger logger.Logger queue chan AuditDocument done chan struct{} wg sync.WaitGroup } // NewClient creates an Elasticsearch client for audit log indexing and search. func NewClient(config Config, log logger.Logger) (Client, error) { if !config.Enabled || strings.TrimSpace(config.URL) == "" { return NewNoopClient(), nil } if config.RequestTimeout == 0 { config.RequestTimeout = 10 * time.Second } if config.BulkFlushPeriod == 0 { config.BulkFlushPeriod = 5 * time.Second } if config.QueueSize == 0 { config.QueueSize = 1000 } if config.IndexPrefix == "" { config.IndexPrefix = "tm-audit-logs" } esCfg := elasticsearch.Config{ Addresses: []string{config.URL}, Username: config.Username, Password: config.Password, } es, err := elasticsearch.NewClient(esCfg) if err != nil { return nil, fmt.Errorf("create elasticsearch client: %w", err) } c := &client{ es: es, config: config, logger: log, queue: make(chan AuditDocument, config.QueueSize), done: make(chan struct{}), } ctx, cancel := context.WithTimeout(context.Background(), config.RequestTimeout) defer cancel() if err := c.ping(ctx); err != nil { return nil, err } c.wg.Add(1) go c.runBulkWorker() log.Info("Successfully connected to Elasticsearch", map[string]interface{}{ "url": config.URL, "index_prefix": config.IndexPrefix, }) return c, nil } func (c *client) ping(ctx context.Context) error { res, err := c.es.Ping(c.es.Ping.WithContext(ctx)) if err != nil { return fmt.Errorf("ping elasticsearch: %w", err) } defer res.Body.Close() if res.IsError() { return fmt.Errorf("elasticsearch ping failed: %s", readErrorBody(res.Body)) } return nil } func (c *client) Index(_ context.Context, doc AuditDocument) error { select { case c.queue <- doc: return nil default: c.logger.Warn("Audit log queue full, dropping entry", map[string]interface{}{ "action": doc.Action, }) return nil } } func (c *client) Search(ctx context.Context, filter SearchFilter) (*SearchResult, error) { indexPattern := c.config.IndexPrefix + "-*" query := buildSearchQuery(filter) payload, err := json.Marshal(query) if err != nil { return nil, fmt.Errorf("marshal search query: %w", err) } res, err := c.es.Search( c.es.Search.WithContext(ctx), c.es.Search.WithIndex(indexPattern), c.es.Search.WithBody(bytes.NewReader(payload)), c.es.Search.WithTrackTotalHits(true), ) if err != nil { return nil, fmt.Errorf("search audit logs: %w", err) } defer res.Body.Close() if res.IsError() { return nil, fmt.Errorf("elasticsearch search error: %s", readErrorBody(res.Body)) } return decodeSearchResponse(res.Body) } func (c *client) Close() error { close(c.done) c.wg.Wait() return nil } func (c *client) runBulkWorker() { defer c.wg.Done() ticker := time.NewTicker(c.config.BulkFlushPeriod) defer ticker.Stop() batch := make([]AuditDocument, 0, 64) flush := func() { if len(batch) == 0 { return } if err := c.flushBatch(batch); err != nil { c.logger.Error("Failed to flush audit logs to Elasticsearch", map[string]interface{}{ "error": err.Error(), "count": len(batch), }) } batch = batch[:0] } for { select { case <-c.done: for { select { case doc := <-c.queue: batch = append(batch, doc) default: flush() return } } case doc := <-c.queue: batch = append(batch, doc) if len(batch) >= 64 { flush() } case <-ticker.C: flush() } } } func (c *client) flushBatch(batch []AuditDocument) error { if len(batch) == 0 { return nil } var buf bytes.Buffer for _, doc := range batch { indexName := indexNameFor(c.config.IndexPrefix, doc.At) meta, err := marshalBulkIndexAction(indexName) if err != nil { return fmt.Errorf("marshal bulk index action: %w", err) } buf.Write(meta) buf.WriteByte('\n') body, err := json.Marshal(doc) if err != nil { return fmt.Errorf("marshal audit document: %w", err) } buf.Write(body) buf.WriteByte('\n') } ctx, cancel := context.WithTimeout(context.Background(), c.config.RequestTimeout) defer cancel() res, err := c.es.Bulk(bytes.NewReader(buf.Bytes()), c.es.Bulk.WithContext(ctx)) if err != nil { return fmt.Errorf("bulk index audit logs: %w", err) } defer res.Body.Close() if res.IsError() { return fmt.Errorf("elasticsearch bulk error: %s", readErrorBody(res.Body)) } return nil } type bulkIndexAction struct { Index struct { Index string `json:"_index"` } `json:"index"` } func marshalBulkIndexAction(indexName string) ([]byte, error) { action := bulkIndexAction{} action.Index.Index = indexName return json.Marshal(action) } // readErrorBody best-effort reads an Elasticsearch error response for diagnostics. // Read errors are intentionally ignored because the HTTP call already failed. func readErrorBody(body io.Reader) string { responseBody, _ := io.ReadAll(body) return string(responseBody) } func indexNameFor(prefix string, at int64) string { t := time.Unix(at, 0).UTC() return fmt.Sprintf("%s-%04d.%02d", prefix, t.Year(), int(t.Month())) } func buildSearchQuery(filter SearchFilter) map[string]interface{} { must := make([]map[string]interface{}, 0) filterClauses := make([]map[string]interface{}, 0) addTerm := func(field, value string) { if value == "" { return } filterClauses = append(filterClauses, map[string]interface{}{ "term": map[string]interface{}{field: value}, }) } addTerm("actor_id", filter.ActorID) addTerm("actor_type", filter.ActorType) addTerm("action", filter.Action) addTerm("target_type", filter.TargetType) addTerm("target_id", filter.TargetID) if filter.From > 0 || filter.To > 0 { rangeQuery := map[string]interface{}{} if filter.From > 0 { rangeQuery["gte"] = filter.From } if filter.To > 0 { rangeQuery["lte"] = filter.To } filterClauses = append(filterClauses, map[string]interface{}{ "range": map[string]interface{}{"at": rangeQuery}, }) } boolQuery := map[string]interface{}{} if len(must) > 0 { boolQuery["must"] = must } if len(filterClauses) > 0 { boolQuery["filter"] = filterClauses } limit := filter.Limit if limit <= 0 { limit = 20 } offset := filter.Offset if offset < 0 { offset = 0 } return map[string]interface{}{ "query": map[string]interface{}{ "bool": boolQuery, }, "sort": []map[string]interface{}{ {"at": map[string]interface{}{"order": "desc"}}, }, "from": offset, "size": limit, } } func decodeSearchResponse(body io.Reader) (*SearchResult, error) { var parsed struct { Hits struct { Total struct { Value int64 `json:"value"` } `json:"total"` Hits []struct { Source AuditDocument `json:"_source"` } `json:"hits"` } `json:"hits"` } if err := json.NewDecoder(body).Decode(&parsed); err != nil { return nil, fmt.Errorf("decode search response: %w", err) } items := make([]AuditDocument, 0, len(parsed.Hits.Hits)) for _, hit := range parsed.Hits.Hits { items = append(items, hit.Source) } return &SearchResult{ Items: items, TotalCount: parsed.Hits.Total.Value, }, nil }