maomaowm/mmsg/mmsg.c

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

int main(int argc, char *argv[]) {
	if (argc < 2) {
		fprintf(stderr, "Usage: mmsg <command> [args...]\n");
		fprintf(stderr, "  get <type> ...      one-shot request\n");
		fprintf(stderr, "  watch <type> ...    persistent stream\n");
		return EXIT_FAILURE;
	}

	const char *socket_path = getenv("MANGO_INSTANCE_SIGNATURE");
	if (!socket_path) {
		fprintf(stderr, "Error: MANGO_INSTANCE_SIGNATURE not set.\n");
		return EXIT_FAILURE;
	}

	int sock = socket(AF_UNIX, SOCK_STREAM, 0);
	if (sock < 0) {
		perror("socket");
		return EXIT_FAILURE;
	}

	struct sockaddr_un addr = {.sun_family = AF_UNIX};
	strncpy(addr.sun_path, socket_path, sizeof(addr.sun_path) - 1);

	if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
		perror("connect");
		close(sock);
		return EXIT_FAILURE;
	}

	/* 拼接命令 */
	char cmd[1024] = {0};
	int offset = 0;
	for (int i = 1; i < argc; i++) {
		offset += snprintf(cmd + offset, sizeof(cmd) - offset, "%s%s", argv[i],
						   (i == argc - 1) ? "" : " ");
	}

	send(sock, cmd, strlen(cmd), 0);

	bool is_watch = (strncmp(argv[1], "watch", 5) == 0);

	if (is_watch) {
		/* watch 模式：持续接收 JSON 并输出，直到连接断开 */
		char buf[4096];
		ssize_t n;
		while ((n = recv(sock, buf, sizeof(buf) - 1, 0)) > 0) {
			buf[n] = '\0';
			printf("%s", buf);
			fflush(stdout);
		}
		close(sock);
		return EXIT_SUCCESS;
	} else {
		/* 一次性命令：读取所有数据，然后关闭输出 */
		char buf[4096];
		ssize_t n;
		while ((n = recv(sock, buf, sizeof(buf) - 1, 0)) > 0) {
			buf[n] = '\0';
			printf("%s", buf);
			fflush(stdout);
		}
		close(sock);
		return EXIT_SUCCESS;
	}
}