[fix] bound what one anonymous push request can consume

hertzbeat-push/src/main/java/org/apache/hertzbeat/push/service/impl/PushGatewayServiceImpl.java

@@ -19,6 +19,7 @@
 
 package org.apache.hertzbeat.push.service.impl;
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.time.Instant;
 import java.util.LinkedList;
@@ -37,6 +38,7 @@
 import org.apache.hertzbeat.common.util.SnowFlakeIdGenerator;
 import org.apache.hertzbeat.push.dao.PushMonitorDao;
 import org.apache.hertzbeat.push.service.PushGatewayService;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Service;
 
 /**
@@ -52,29 +54,78 @@ public class PushGatewayServiceImpl implements PushGatewayService {
     private final PushMonitorDao pushMonitorDao;
 
     private final Map<String, Long> jobInstanceMap;
-
-    public PushGatewayServiceImpl(CommonDataQueue commonDataQueue, PushMonitorDao pushMonitorDao) {
+
+    /**
+     * Cap on push monitors created automatically from unknown job/instance pairs.
+     *
+     * <p>The route is unauthenticated by design, and every new pair used to persist a
+     * monitor row and add a `jobInstanceMap` entry that is never removed, so a caller
+     * iterating over made up names could grow the database and the heap without bound.
+     * Above the cap an unknown pair is refused while the pairs already known keep working,
+     * which is why eviction is not used here: evicting a live entry would make the next
+     * push for that pair create a second monitor for the same job and instance.
+     */
+    private final int maxAutoCreatedMonitors;
+
+    /**
+     * Cap on how many bytes a single push body may carry.
+     *
+     * <p>`OnlineParser.parseMetrics` builds its result in memory with no limit of its own,
+     * and the servlet container does not bound a non form request body, so one unbounded
+     * request was enough to exhaust the heap.
+     */
+    private final long maxBodyBytes;
+
+    /**
+     * Cap on how many samples a single push body may carry, applied after parsing so a body
+     * that is small on the wire cannot still flood the collection queue.
+     */
+    private final int maxSamples;
+
+    public PushGatewayServiceImpl(CommonDataQueue commonDataQueue, PushMonitorDao pushMonitorDao,
+                                  @Value("${hertzbeat.push.max-auto-created-monitors:10000}") int maxAutoCreatedMonitors,
+                                  @Value("${hertzbeat.push.max-body-bytes:5242880}") long maxBodyBytes,
+                                  @Value("${hertzbeat.push.max-samples:10000}") int maxSamples) {
         this.commonDataQueue = commonDataQueue;
         this.pushMonitorDao = pushMonitorDao;
+        this.maxAutoCreatedMonitors = maxAutoCreatedMonitors;
+        this.maxBodyBytes = maxBodyBytes;
+        this.maxSamples = maxSamples;
         jobInstanceMap = new ConcurrentHashMap<>();
-        pushMonitorDao.findMonitorsByType((byte) 1).forEach(monitor -> 
+        pushMonitorDao.findMonitorsByType((byte) 1).forEach(monitor ->
                 jobInstanceMap.put(monitor.getApp() + "_" + monitor.getName(), monitor.getId()));
     }
 
     @Override
     public boolean pushPrometheusMetrics(InputStream inputStream, String job, String instance) {
         try {
             long curTime = Instant.now().toEpochMilli();
-            Map<String, MetricFamily> metricFamilyMap = OnlineParser.parseMetrics(inputStream);
+            Map<String, MetricFamily> metricFamilyMap =
+                    OnlineParser.parseMetrics(new BoundedInputStream(inputStream, maxBodyBytes));
             if (metricFamilyMap == null) {
                 log.error("parse prometheus metrics is null, job: {}, instance: {}", job, instance);
                 return false;
             }
+            int samples = metricFamilyMap.values().stream()
+                    .mapToInt(family -> family.getMetricList().size())
+                    .sum();
+            if (samples > maxSamples) {
+                log.warn("reject prometheus push carrying {} samples, limit is {}, job: {}, instance: {}",
+                        samples, maxSamples, job, instance);
+                return false;
+            }
             long id = 0L;
             if (job != null && instance != null) {
                 // auto create monitor when job and instance not null
                 // job is app, instance is the name
-                id = jobInstanceMap.computeIfAbsent(job + "_" + instance, key -> {
+                String key = job + "_" + instance;
+                if (!jobInstanceMap.containsKey(key) && jobInstanceMap.size() >= maxAutoCreatedMonitors) {
+                    log.warn("reject prometheus push for unknown job: {}, instance: {}, "
+                                    + "already tracking {} push monitors, limit is {}",
+                            job, instance, jobInstanceMap.size(), maxAutoCreatedMonitors);
+                    return false;
+                }
+                id = jobInstanceMap.computeIfAbsent(key, ignored -> {
                     log.info("auto create monitor by prometheus push, job: {}, instance: {}", job, instance);
                     long monitorId = SnowFlakeIdGenerator.generateId();
                     Monitor monitor = Monitor.builder()
@@ -130,4 +181,53 @@ public boolean pushPrometheusMetrics(InputStream inputStream, String job, String
             return false;
         }
     }
+
+    /**
+     * Fails the read once the body has delivered more than {@code limit} bytes, instead of
+     * letting the parser accumulate an unbounded body in memory. Reading stops at the
+     * failure, so the bytes beyond the limit are never buffered.
+     */
+    static final class BoundedInputStream extends InputStream {
+
+        private final InputStream delegate;
+
+        private final long limit;
+
+        private long read;
+
+        BoundedInputStream(InputStream delegate, long limit) {
+            this.delegate = delegate;
+            this.limit = limit;
+        }
+
+        @Override
+        public int read() throws IOException {
+            int value = delegate.read();
+            if (value != -1) {
+                count(1);
+            }
+            return value;
+        }
+
+        @Override
+        public int read(byte[] buffer, int offset, int length) throws IOException {
+            int count = delegate.read(buffer, offset, length);
+            if (count > 0) {
+                count(count);
+            }
+            return count;
+        }
+
+        private void count(int increment) throws IOException {
+            read += increment;
+            if (read > limit) {
+                throw new IOException("push body exceeds the " + limit + " byte limit");
+            }
+        }
+
+        @Override
+        public void close() throws IOException {
+            delegate.close();
+        }
+    }
 }

hertzbeat-push/src/test/java/org/apache/hertzbeat/push/service/impl/PushGatewayServiceImplTest.java

@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hertzbeat.push.service.impl;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.lenient;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import org.apache.hertzbeat.common.entity.manager.Monitor;
+import org.apache.hertzbeat.common.queue.CommonDataQueue;
+import org.apache.hertzbeat.push.dao.PushMonitorDao;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+/**
+ * Test case for {@link PushGatewayServiceImpl}.
+ *
+ * <p>`/api/push/prometheus/**` is unauthenticated by design, so the resource a single
+ * anonymous request may consume has to be bounded: the body it may carry, the samples it
+ * may enqueue, and the number of push monitors it may bring into existence.
+ */
+@ExtendWith(MockitoExtension.class)
+class PushGatewayServiceImplTest {
+
+    private static final String BODY = "sample_metric{label=\"a\"} 1\n";
+
+    @Mock
+    private CommonDataQueue commonDataQueue;
+
+    @Mock
+    private PushMonitorDao pushMonitorDao;
+
+    @BeforeEach
+    void setUp() {
+        // the stream test below builds no service, so this default must not be strict
+        lenient().when(pushMonitorDao.findMonitorsByType((byte) 1)).thenReturn(List.of());
+    }
+
+    private PushGatewayServiceImpl service(int maxMonitors, long maxBodyBytes, int maxSamples) {
+        return new PushGatewayServiceImpl(commonDataQueue, pushMonitorDao, maxMonitors, maxBodyBytes, maxSamples);
+    }
+
+    private static ByteArrayInputStream body(String content) {
+        return new ByteArrayInputStream(content.getBytes(StandardCharsets.UTF_8));
+    }
+
+    @Test
+    void testPushIsAcceptedWithinTheLimits() {
+        PushGatewayServiceImpl service = service(10, 1024, 100);
+
+        assertTrue(service.pushPrometheusMetrics(body(BODY), "job1", "instance1"));
+
+        verify(pushMonitorDao).save(any(Monitor.class));
+    }
+
+    @Test
+    void testBodyBeyondTheByteLimitIsRejected() {
+        PushGatewayServiceImpl service = service(10, 16, 100);
+
+        assertFalse(service.pushPrometheusMetrics(body(BODY.repeat(100)), "job1", "instance1"));
+
+        verify(pushMonitorDao, never()).save(any(Monitor.class));
+    }
+
+    @Test
+    void testBodyBeyondTheSampleLimitIsRejected() {
+        PushGatewayServiceImpl service = service(10, 1024 * 1024, 2);
+        StringBuilder many = new StringBuilder();
+        for (int index = 0; index < 10; index++) {
+            many.append("sample_metric{label=\"value").append(index).append("\"} 1\n");
+        }
+
+        assertFalse(service.pushPrometheusMetrics(body(many.toString()), "job1", "instance1"));
+
+        verify(pushMonitorDao, never()).save(any(Monitor.class));
+    }
+
+    /**
+     * An unknown job/instance pair persists a monitor row and adds a map entry that is
+     * never removed, so without a cap an anonymous caller iterating over made up names
+     * grows the database and the heap without bound.
+     */
+    @Test
+    void testAutoCreationStopsAtTheMonitorLimit() {
+        PushGatewayServiceImpl service = service(2, 1024, 100);
+
+        assertTrue(service.pushPrometheusMetrics(body(BODY), "job1", "instance1"));
+        assertTrue(service.pushPrometheusMetrics(body(BODY), "job2", "instance2"));
+        assertFalse(service.pushPrometheusMetrics(body(BODY), "job3", "instance3"));
+
+        verify(pushMonitorDao, times(2)).save(any(Monitor.class));
+    }
+
+    /**
+     * The cap must not turn into eviction: a pair already known has to keep resolving to
+     * the monitor it created, otherwise a later push would create a second monitor for the
+     * same job and instance.
+     */
+    @Test
+    void testKnownPairsKeepWorkingAtTheLimit() {
+        PushGatewayServiceImpl service = service(1, 1024, 100);
+
+        assertTrue(service.pushPrometheusMetrics(body(BODY), "job1", "instance1"));
+        assertFalse(service.pushPrometheusMetrics(body(BODY), "other", "instance"));
+        assertTrue(service.pushPrometheusMetrics(body(BODY), "job1", "instance1"));
+
+        verify(pushMonitorDao, times(1)).save(any(Monitor.class));
+    }
+
+    @Test
+    void testMonitorsLoadedAtStartupCountTowardsTheLimit() {
+        lenient().when(pushMonitorDao.findMonitorsByType((byte) 1)).thenReturn(List.of(
+                Monitor.builder().id(1L).app("job1").name("instance1").build()));
+        PushGatewayServiceImpl service = service(1, 1024, 100);
+
+        assertFalse(service.pushPrometheusMetrics(body(BODY), "job2", "instance2"));
+        assertTrue(service.pushPrometheusMetrics(body(BODY), "job1", "instance1"));
+
+        verify(pushMonitorDao, never()).save(any(Monitor.class));
+    }
+
+    @Test
+    void testBoundedStreamStopsAtTheLimit() throws Exception {
+        PushGatewayServiceImpl.BoundedInputStream stream =
+                new PushGatewayServiceImpl.BoundedInputStream(body("abcdef"), 3);
+
+        assertEquals('a', stream.read());
+        assertEquals('b', stream.read());
+        assertEquals('c', stream.read());
+        assertThrows(IOException.class, stream::read);
+    }
+}