test(cpworker): test checksums, parse_packet bounds, zmq flush and control-socket errors by spec

cpworker/tests/integration/control_socket_test.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"path/filepath"
@@ -87,36 +88,69 @@ func (c *controlConn) send(gap time.Duration, parts ...string) error {
 	return nil
 }
 
-// expectOK reads one reply line and checks that its status is OK.
-func (c *controlConn) expectOK() error {
+// readReply reads one reply line and decodes it.
+func (c *controlConn) readReply() (map[string]any, error) {
 	c.conn.SetReadDeadline(time.Now().Add(controlIOTimeout))
 	line, err := c.r.ReadBytes('\n')
 	if err != nil {
-		return fmt.Errorf("read reply: %w", err)
+		return nil, fmt.Errorf("read reply: %w", err)
 	}
 	var resp map[string]any
 	if err := json.Unmarshal(line, &resp); err != nil {
-		return fmt.Errorf("decode reply %q: %w", line, err)
+		return nil, fmt.Errorf("decode reply %q: %w", line, err)
+	}
+	return resp, nil
+}
+
+// expectOK reads one reply line and checks that its status is OK.
+func (c *controlConn) expectOK() error {
+	resp, err := c.readReply()
+	if err != nil {
+		return err
 	}
 	if resp["status"] != "OK" {
-		return fmt.Errorf("reply status is not OK: %s", line)
+		return fmt.Errorf("reply status is not OK: %v", resp)
 	}
 	return nil
 }
 
-// handshakeAndPing runs the v1 handshake and one ping on a new connection.
-func handshakeAndPing(sock string) error {
+// expectClosed checks that the manager closes the connection without replying.
+func (c *controlConn) expectClosed() error {
+	c.conn.SetReadDeadline(time.Now().Add(controlIOTimeout))
+	line, err := c.r.ReadBytes('\n')
+	if err == io.EOF && len(line) == 0 {
+		return nil
+	}
+	if err != nil {
+		return fmt.Errorf("want the connection closed, got: %w", err)
+	}
+	return fmt.Errorf("want the connection closed, got reply %q", line)
+}
+
+// dialHandshaken opens a connection and runs the v1 handshake on it.
+func dialHandshaken(sock string) (*controlConn, error) {
 	c, err := dialControl(sock)
 	if err != nil {
-		return fmt.Errorf("dial: %w", err)
+		return nil, fmt.Errorf("dial: %w", err)
 	}
-	defer c.conn.Close()
 	if err := c.send(0, `{"version":"v1"}`+"\n"); err != nil {
-		return fmt.Errorf("send handshake: %w", err)
+		c.conn.Close()
+		return nil, fmt.Errorf("send handshake: %w", err)
 	}
 	if err := c.expectOK(); err != nil {
-		return fmt.Errorf("handshake: %w", err)
+		c.conn.Close()
+		return nil, fmt.Errorf("handshake: %w", err)
+	}
+	return c, nil
+}
+
+// handshakeAndPing runs the v1 handshake and one ping on a new connection.
+func handshakeAndPing(sock string) error {
+	c, err := dialHandshaken(sock)
+	if err != nil {
+		return err
 	}
+	defer c.conn.Close()
 	if err := c.send(0, `{"command":"ping"}`+"\n"); err != nil {
 		return fmt.Errorf("send ping: %w", err)
 	}
@@ -134,6 +168,8 @@ func handshakeAndPing(sock string) error {
 //   - request_split_across_writes: a request that arrives in several writes is answered (#261)
 //   - silent_client_does_not_block_others: a client that connects and never sends the
 //     handshake does not stop the manager from serving other clients
+//   - rejected_requests: a bad handshake, a request that is not JSON or has no string
+//     "command" closes that connection; an unknown command is answered with ERROR
 func TestControlSocket(t *testing.T) {
 	if os.Getenv("CPWORKER_BIN") == "" {
 		t.Skip("CPWORKER_BIN not set")
@@ -208,4 +244,63 @@ func TestControlSocket(t *testing.T) {
 			t.Errorf("client behind a silent connection: %v", err)
 		}
 	})
+
+	t.Run("rejected_requests", func(t *testing.T) {
+		sock := startControlWorker(t, "rejected_requests")
+
+		for _, hs := range []string{`not json`, `{"version":"v2"}`, `{"version":1}`, `{}`} {
+			c, err := dialControl(sock)
+			if err != nil {
+				t.Fatalf("dial: %v", err)
+			}
+			if err := c.send(0, hs+"\n"); err != nil {
+				t.Fatalf("send handshake %s: %v", hs, err)
+			}
+			if err := c.expectClosed(); err != nil {
+				t.Errorf("handshake %s: %v", hs, err)
+			}
+			c.conn.Close()
+		}
+
+		// An unknown command is answered, and the connection stays usable.
+		c, err := dialHandshaken(sock)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := c.send(0, `{"command":"no_such_command"}`+"\n"); err != nil {
+			t.Fatalf("send unknown command: %v", err)
+		}
+		resp, err := c.readReply()
+		if err != nil {
+			t.Fatalf("unknown command: %v", err)
+		}
+		if resp["status"] != "ERROR" || resp["message"] != "unknown command" {
+			t.Errorf("unknown command: got %v, want status ERROR, message \"unknown command\"", resp)
+		}
+		if err := c.send(0, `{"command":"ping"}`+"\n"); err != nil {
+			t.Fatalf("send ping: %v", err)
+		}
+		if err := c.expectOK(); err != nil {
+			t.Errorf("ping after unknown command: %v", err)
+		}
+		c.conn.Close()
+
+		for _, req := range []string{`not json`, `{"command":42}`, `{}`} {
+			c, err := dialHandshaken(sock)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := c.send(0, req+"\n"); err != nil {
+				t.Fatalf("send request %s: %v", req, err)
+			}
+			if err := c.expectClosed(); err != nil {
+				t.Errorf("request %s: %v", req, err)
+			}
+			c.conn.Close()
+		}
+
+		if err := handshakeAndPing(sock); err != nil {
+			t.Errorf("client after the rejected requests: %v", err)
+		}
+	})
 }

cpworker/tests/unit/output_zmq.c

@@ -545,6 +545,75 @@ void test_rate_limit_charges_and_counts_sliced_frame_length(void)
     TEST_ASSERT_EQUAL_UINT64(1000, stats.ratelimit_drop_bytes.bytes);
 }
 
+/* ---- docs/ZMQ-WIRE-FORMAT.md §8: batch flush rules ---- */
+
+// A batch is flushed only when the next record would take it past ZMQ_MAX_BATCH_BUF_SIZE. A
+// 2451-byte frame makes a 2473-byte record (length, packet header, frame, MPLS label), and 424
+// of them after the batch header fill the buffer exactly.
+void test_batch_flushed_when_next_record_exceeds_buffer(void)
+{
+    char errbuf[ERROR_BUFFER_SIZE] = {0};
+    zmq_output_t *output = new_output(0, (char *)VALID_UUID, errbuf);
+    TEST_ASSERT_NOT_NULL_MESSAGE(output, errbuf);
+
+    static uint8_t frame[2451];
+    memset(frame, 0x5a, sizeof(frame));
+    frame[12] = 0x08;
+    frame[13] = 0x00;
+
+    for (int i = 0; i < 424; i++)
+        TEST_ASSERT_EQUAL_INT(0, send_frame(output, frame, sizeof(frame)));
+    TEST_ASSERT_EQUAL_UINT32(ZMQ_MAX_BATCH_BUF_SIZE, output->pkts_buf.batch_bufpos);
+    TEST_ASSERT_EQUAL_UINT16(424, output->pkts_buf.batch_hdr.pkts_num);
+
+    TEST_ASSERT_EQUAL_INT(0, send_frame(output, frame, sizeof(frame)));
+    TEST_ASSERT_EQUAL_UINT16(1, output->pkts_buf.batch_hdr.pkts_num);
+    TEST_ASSERT_EQUAL_INT(ZMQ_MAX_BATCH_BUF_SIZE, recv_batch());
+    TEST_ASSERT_EQUAL_UINT16(424, batch_pkts_num());
+
+    zmq_output_destroy(&output->base);
+}
+
+// A batch is flushed when a packet arrives more than ZMQ_PKTS_FLUSH_MAX_DUR_SEC after its first
+// packet; that packet starts the next batch.
+void test_batch_flushed_when_packet_is_over_a_second_later(void)
+{
+    char errbuf[ERROR_BUFFER_SIZE] = {0};
+    zmq_output_t *output = new_output(0, (char *)VALID_UUID, errbuf);
+    TEST_ASSERT_NOT_NULL_MESSAGE(output, errbuf);
+
+    send_frame_at(output, 100);
+    send_frame_at(output, 101);
+    TEST_ASSERT_EQUAL_UINT16(2, output->pkts_buf.batch_hdr.pkts_num);
+
+    send_frame_at(output, 102);
+    TEST_ASSERT_EQUAL_UINT16(1, output->pkts_buf.batch_hdr.pkts_num);
+    TEST_ASSERT_GREATER_THAN_INT(0, recv_batch());
+    TEST_ASSERT_EQUAL_UINT16(2, batch_pkts_num());
+
+    zmq_output_destroy(&output->base);
+}
+
+// With no traffic, the periodic zmq_heartbeat() call flushes a batch once its first packet is
+// more than ZMQ_PKTS_FLUSH_MAX_DUR_SEC old, even when heartbeats are disabled.
+void test_idle_batch_flushed_after_a_second(void)
+{
+    char errbuf[ERROR_BUFFER_SIZE] = {0};
+    zmq_output_t *output = new_output(0, (char *)VALID_UUID, errbuf);
+    TEST_ASSERT_NOT_NULL_MESSAGE(output, errbuf);
+
+    send_frame_at(output, 100);
+    zmq_heartbeat(&output->base, 101);
+    TEST_ASSERT_EQUAL_UINT16(1, output->pkts_buf.batch_hdr.pkts_num);
+
+    zmq_heartbeat(&output->base, 102);
+    TEST_ASSERT_EQUAL_UINT16(0, output->pkts_buf.batch_hdr.pkts_num);
+    TEST_ASSERT_GREATER_THAN_INT(0, recv_batch());
+    TEST_ASSERT_EQUAL_UINT16(1, batch_pkts_num());
+
+    zmq_output_destroy(&output->base);
+}
+
 int main(void)
 {
     UNITY_BEGIN();
@@ -573,5 +642,9 @@ int main(void)
     RUN_TEST(test_counters_use_sliced_frame_length_without_headers);
     RUN_TEST(test_rate_limit_charges_and_counts_sliced_frame_length);
 
+    RUN_TEST(test_batch_flushed_when_next_record_exceeds_buffer);
+    RUN_TEST(test_batch_flushed_when_packet_is_over_a_second_later);
+    RUN_TEST(test_idle_batch_flushed_after_a_second);
+
     return UNITY_END();
 }