// sk_msg.c - L4 socket-level programs for HTTP header injection
// Includes: sockops, cgroup/skb, sk_msg, sk_skb/stream_verdict
// These operate at the socket layer, not L7 protocol parsing.

// sockhash map for stream_verdict program
struct {
	__uint(type, BPF_MAP_TYPE_SOCKHASH);
	__uint(max_entries, 65535);
	__type(key, struct bpf_sock_tuple);
	__type(value, __u32); // socket cookie
} sockhash SEC(".maps");

// 用于在cgroup/skb和sk_skb程序间传递信息的map
struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(max_entries, 65535);
	__type(key, __u32); // socket cookie
	__type(value, __u32); // flag: 1=需要修改HTTP请求
} http_modify_flags SEC(".maps");

// sk_msg map for sk_msg program
struct {
	__uint(type, BPF_MAP_TYPE_SOCKMAP);
	__uint(max_entries, 65535);
	__type(key, __u32); // socket cookie
	__type(value, __u32); // placeholder
} sk_msg_map SEC(".maps");

// cwother 开关：key=0, value=1 启用，0 或不存在则跳过
struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__type(key, __u32);
	__type(value, __u8);
	__uint(max_entries, 1);
} l4_cwother_enabled SEC(".maps");

//SEC("sk_msg")
//int sk_msg_handler(struct sk_msg_md *msg)
//{
//	const char header[] = "X-Debug: 1234\r\n";
//	const int hdr_len = sizeof(header) - 1; // 不含结尾 '\0'
//
//	// 在最开头插入 hdr_len 字节
//	if (bpf_msg_push_data(msg, 0, hdr_len, 0) < 0) {
//		bpf_printk("push_data failed\n");
//		return SK_DROP;
//	}
//
//	// 把新数据段线性化，确保能写
//	if (bpf_msg_pull_data(msg, 0, hdr_len, 0) < 0) {
//		bpf_printk("pull_data failed\n");
//		return SK_DROP;
//	}
//
//	char *data = (char *)(long)msg->data;
//	char *data_end = (char *)(long)msg->data_end;
//
//	// 写入 header 字符串
//#pragma clang loop unroll(full)
//	for (int i = 0; i < hdr_len; i++) {
//		if (data + i + 1 > data_end)
//			break;
//		*(char *)(data + i) = header[i];
//	}
//
//	bpf_printk("Inserted header len=%d\n", hdr_len);
//	return SK_PASS;
//}

// sockops程序：将socket添加到sockhash map中
SEC("sockops")
int sockops_cb(struct bpf_sock_ops *skops) {
	// 添加调试信息
//	bpf_printk("sockops: received event, op=%d, family=%d\n", skops->op, skops->family);

	// 只处理TCP连接
	bpf_printk("skops->family - %d",skops->family);
	if (skops->family != AF_INET && skops->family != AF_INET6) {
		return 0;
	}



	// 只处理客户端主动建立的连接（出站连接）
	// 这样stream_verdict就能看到客户端发送的HTTP请求
	if (skops->op != BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB) {
		return 0;
	}
	__u32 saddr = bpf_ntohl(skops->local_ip4);
	__u32 daddr = bpf_ntohl(skops->remote_ip4);
	__u16 sport = bpf_ntohs(skops->local_port);
	__u16 dport = bpf_ntohs(skops->remote_port);
//	bpf_printk("sockops op=%d src=%08x:%u\n", skops->op, saddr, sport);
//	bpf_printk("        dst=%08x:%u\n", daddr, dport);

	// 创建socket tuple
	struct bpf_sock_tuple tuple = {};
	tuple.ipv4.saddr = skops->local_ip4;
	tuple.ipv4.daddr = skops->remote_ip4;
	tuple.ipv4.sport = skops->local_port;
	tuple.ipv4.dport = skops->remote_port;

	// 将socket添加到sockhash map
	__u32 cookie = bpf_get_socket_cookie(skops);
	bpf_sock_hash_update(skops, &sockhash, &tuple, BPF_ANY);

	// 同时将socket添加到sockmap（用于sk_msg程序）
	bpf_sock_map_update(skops, &sk_msg_map, &cookie, BPF_ANY);

//	bpf_printk("sockops: added socket to sockhash and sockmap, cookie=%u\n", cookie);

	return 0;
}

// 专门处理HTTP请求的程序 - 使用cgroup/skb
SEC("cgroup/skb")
int http_request_handler(struct __sk_buff *skb) {
	// 检查数据包长度
//	if (skb->len < 4) {
//		return 1; // 允许通过
//	}
//
//	// 读取前几个字节来检查是否是HTTP请求
//	char buf[100] = {};
//	if (bpf_skb_load_bytes(skb, 0, buf, 100) < 0) {
//		return 1; // 允许通过
//	}
//
//
//	// 跳过IP头部(20字节)和TCP头部(20字节)，直接读取应用层数据
//	// IP头部长度 = (buf[0] & 0x0F) * 4 = 5 * 4 = 20字节
//	// TCP头部长度 = (buf[32] >> 4) * 4，但通常也是20字节
//	int ip_header_len = (buf[0] & 0x0F) * 4;
//	int tcp_header_len = 20; // 通常TCP头部是20字节
//	int app_data_offset = ip_header_len + tcp_header_len;
//
//	// 检查是否有足够的应用层数据
//	if (skb->len < app_data_offset + 4) {
//		return 1; // 允许通过
//	}
//
//	// 读取应用层数据的前几个字节
//	char app_data[100] = {};
//	if (bpf_skb_load_bytes(skb, app_data_offset, app_data, 100) < 0) {
//		return 1; // 允许通过
//	}
//
//	// 检查是否是HTTP请求
//	if ((app_data[0] == 'G' && app_data[1] == 'E' && app_data[2] == 'T' && app_data[3] == ' ') ||
//	    (app_data[0] == 'P' && app_data[1] == 'O' && app_data[2] == 'S' && app_data[3] == 'T') ||
//	    (app_data[0] == 'P' && app_data[1] == 'U' && app_data[2] == 'T' && app_data[3] == ' ') ||
//	    (app_data[0] == 'D' && app_data[1] == 'E' && app_data[2] == 'L' && app_data[3] == 'E') ||
//	    (app_data[0] == 'H' && app_data[1] == 'E' && app_data[2] == 'A' && app_data[3] == 'D') ||
//	    (app_data[0] == 'O' && app_data[1] == 'P' && app_data[2] == 'T' && app_data[3] == 'I')) {
//
////		bpf_printk("=== HTTP REQUEST FOUND ===\n");
////		bpf_printk("IP header len: %d, TCP header len: %d\n", ip_header_len, tcp_header_len);
////		bpf_printk("App data offset: %d\n", app_data_offset);
////		bpf_printk("Content: %s\n", app_data);
////		bpf_printk("=== END HTTP REQUEST ===\n");
//
//		// 标记这个socket需要修改HTTP请求
//		// 通过http_modify_flags map传递信息给stream_verdict程序
//		__u32 cookie = bpf_get_socket_cookie(skb);
//		__u32 flag = 1; // 标记需要修改
//		if (bpf_map_update_elem(&http_modify_flags, &cookie, &flag, BPF_ANY) == 0) {
//			bpf_printk("Marked socket %u for HTTP request modification\n", cookie);
//		}
//	}
//
	return 1; // 允许数据包通过
}

// sk_msg program for handling socket messages
SEC("sk_msg")
int sk_msg_handler(struct sk_msg_md *msg)
{
	long ret;
	if (bpf_msg_pull_data(msg, 0, 16, 0) < 0)
		return SK_PASS;

	void *p = (void *) (long) msg->data;
	void *e = (void *) (long) msg->data_end;
	if (!is_http_request_in_sk((void *) (__u64) msg->data)) {
		return SK_PASS;
	}

	__u32 tgid = (__u32) (bpf_get_current_pid_tgid() >> 32);
	struct ebpf_proc_info *proc_info = bpf_map_lookup_elem(&proc_info_map, &tgid);
	if (!proc_info) {
		return SK_PASS;
	}

	if (!mk_header_in_sk_msg(proc_info->code_type)) {
//		bpf_printk("not allowd");
		return SK_PASS;
	}

	__u32 max_buf = msg->size < MAX_L7_IOVEC_BUF_SIZE ? msg->size : MAX_L7_IOVEC_BUF_SIZE;

	// 拉满整条消息查Header
	if (bpf_msg_pull_data(msg, 0, max_buf, 0) < 0){
//		bpf_printk("HTTP request return");
		return SK_PASS;
	}
	bpf_printk("HTTP request return2");
//	bpf_printk("HTTP request return2 ：%d",  (__u64)msg->data_end - (__u64)msg->data);

//	return SK_PASS;


	__u32 k0 = 0;
	struct sock_t *map_data = bpf_map_lookup_elem(&socket_heap, &k0);
	if (!map_data) {
		cw_bpf_debug("[java client] Failed to lookup socket_heap\n");
		return SK_PASS;
	}

	if (bpf_probe_read_kernel(map_data->payload, sizeof(map_data->payload), (void *) (__u64) msg->data) != 0) {
		return SK_PASS;
	}
	bpf_printk("HTTP request: %s",map_data->payload);

	// header 查询
	long header_start_native = 0x0a0d312e312f5054LL; // 小端序下的 "TP/1.1\r\n" (0x54 0x50 0x2f 0x31 0x2e 0x31 0x0d 0x0a)
#pragma clang loop unroll(full)
	for (__u32 i = 0; i < 200 - 8; i++) {
		long data = *(long long *) (map_data->payload + i);
		if (data == header_start_native) {
			map_data->header_offset_idx = i + 8;
			break;
		}
	}
	bpf_printk("header_offset_idx %d",map_data->header_offset_idx);
	if (map_data->header_offset_idx == 0) {
		bpf_printk("[java client] header_offset_idx error:\n");
		return SK_PASS;
	}

//	bpf_printk("HTTP header_offset_idx: %d\n", map_data->header_offset_idx);

	if (map_data->header_offset_idx > msg->size){
		return SK_PASS;
	}

	// 先拉直 [0, offset) 让插入点可达
	if (bpf_msg_pull_data(msg, 0, map_data->header_offset_idx, 0) < 0){
		bpf_printk("HTTP request return");
		return SK_PASS;
	}

	// TODO real Header
//	const char insert[] = "AA: a\r\n";   // 要插入的 7 字节
//	const char insert[CW_STREAM_HEADER_LEN] = "cwtrace: 00:00:1015481350055581:5450531005555981:5610250100539899:304775019cd3218a304775019cd3218a:1001025098564810:140acc88cde8773f\r\n";

//	int ins_len = sizeof(insert) - 1;  // 不含 '\0'


//	__u32 rip_net = msg->remote_ip4;           // network order
//	__u32 raw_port = msg->remote_port;         // 32-bit field
//	unsigned char assumed_app_id[16];
//	ipport_to_8digits_from_u32_high(rip_net, raw_port, assumed_app_id);

//	for (int i = 0; i < APM_ASSUMED_APP_ID_SIZE; i++) {
//	    bpf_printk("assumed_app_id-assumed_app_id[%d] = %02x", i, assumed_app_id[i]);
//	}

	struct apm_span_context *cw_sc = build_sc_by_ipport(*proc_info, msg->remote_ip4, msg->remote_port);
	if (cw_sc == NULL) {
		return SK_PASS;
	}


	span_context_to_cw_string_stream(cw_sc, map_data->header_stream, code_type_to_code(proc_info->code_type));

	// cwtrace 写完后，在 header_stream 后面拼接 cwother
	__u32 ins_len = CW_STREAM_HEADER_LEN - 1; // cwtrace 内容长度（不含 '\0'）= 134

	__u32 _k0 = 0;
	__u8 *cwother_flag = bpf_map_lookup_elem(&l4_cwother_enabled, &_k0);
	bpf_printk("%d",proc_info->sysvc_len);

	if (cwother_flag && *cwother_flag && proc_info->sysvc[0] != '\0') {
		// 计算 sysvc 实际长度
		__u64 sys_val_len = 0;
#pragma clang loop unroll(full)
		for (int i = 0; i < CW_SYS_HEADER_VAL_LENGTH; i++) {
			if (proc_info->sysvc[i] == '\0') {
				sys_val_len = i;
				break;
			}
			if (i == CW_SYS_HEADER_VAL_LENGTH - 1)
				sys_val_len = CW_SYS_HEADER_VAL_LENGTH;
		}

		// 自动获取
		// sys_val_len = proc_info->sysvc_len;
// sys_val_len &= (CW_SYS_HEADER_VAL_LENGTH - 1); // 假设是2的幂

		if (sys_val_len > 0) {
			__u32 off = ins_len;
			__builtin_memcpy(&map_data->header_stream[off], "cwother: ", CW_SYS_HEADER_KEY_LENGTH + 2);
			off += CW_SYS_HEADER_KEY_LENGTH + 2;
			__builtin_memcpy(&map_data->header_stream[off], proc_info->sysvc, CW_SYS_HEADER_VAL_LENGTH);
			off += sys_val_len;
			map_data->header_stream[off]     = '\r';
			map_data->header_stream[off + 1] = '\n';
			off += 2;
			ins_len = off;
//			bpf_printk("sk_msg cwother: ins_len=%u sys_val_len=%u\n", ins_len, sys_val_len);
		}
	}

	__u32 rip = __builtin_bswap32(msg->remote_ip4); // host order
	__u16 rport = __builtin_bswap32(msg->remote_port);
	__u16 lport = (__u16)(msg->local_port); // host order
//	bpf_printk("sk_msg: rport=%u lport=%u ins_len=%u\n", rport, lport, ins_len);

	__u32 header_offset = map_data->header_offset_idx;
	if (header_offset > msg->size){
		return SK_PASS;
	}

	// 2) 执行插入
	ret = bpf_msg_push_data(msg, header_offset, ins_len, 0);
	if (ret) {
		bpf_printk("push off=%u len=%u ret=%ld\n", header_offset, ins_len, ret);
		return SK_PASS;
	}

//	bpf_printk("bpf_msg_push_data suc.");

	// push_data 后；拉到写满Header处的长度
//	__u32 need = msg->size < MAX_L7_IOVEC_BUF_SIZE ? msg->size : MAX_L7_IOVEC_BUF_SIZE;

	ret = bpf_msg_pull_data(msg, 0, header_offset + ins_len, 0);
	if (ret) {
		bpf_printk("pull 0..%u fail ret=%ld\n", header_offset + ins_len, ret);
		return SK_PASS;
	}
//	bpf_printk("bpf_msg_pull_data suc.");

	// 3. 写入插入的内容
	char *data     = (char *)(long)msg->data;
	char *data_end = (char *)(long)msg->data_end;
	__u64 len =(__u32)(data_end - data);

//	__u32 header_off = offset;

	// 运行时长度 n：先限制到常量 MAX_INS，再裁剪到剩余空间
//	__u32 n = ins_len;                  // 运行时的插入/覆盖长度

	if (header_offset > len)
		return SK_PASS;

#pragma clang loop unroll(full)
	for (int i = 0; i < MAX_L7_IOVEC_BUF_SIZE; i++) {
		/*test1*/
		if (i >= header_offset) {
			for (int j = 0; j < CW_COMBINED_STREAM_MAX_LEN; j++) { // 编译期上界 221
				if (j >= ins_len) break;  // 运行时退出
				if (data + i + 1 > data_end) break;
				*(char *) (data + i) = map_data->header_stream[j];
				i++;
			}
			break;
		}
		if (i > header_offset) break;
		/*test1 end*/


		/*test2
		 * if (i < header_offset){
			bpf_printk("continue-%d",i);
			continue;
		}
		if (data + i + 1 > data_end) break;
		if (i == header_offset + ins_len){
			bpf_printk("break-%d",i);
			break;
		}
		for (int j = 0; j < ins_len; j++) {
			if (data + i + 1 > data_end) break;
			*(char *)(data + i) = insert[j];
			bpf_printk("set-%d %d", i, j);
			i++;
		}
		break;*/

	}

	// save span context
//	bpf_printk("sk_msg save: rip=0x%x rport=%u lport=%u\n", rip, rport, lport);
	cw_save_current_span_context_by_ipport(rip, rport, lport, cw_sc);

	return SK_PASS;

//	bpf_printk("sk_msg");
//	const char header[] = "X-Debug: 1234\r\n";
//	const int hdr_len = sizeof(header) - 1; // 不含结尾 '\0'

	// 读取HTTP请求数据
//	__u64 data_start = (__u64)msg->data;
//	__u64 data_end = (__u64)msg->data_end;
//	__u64 data_len = (__u64)msg->data_end - (__u64)msg->data;

	// 查找第一行结束位置（\r\n）
//	GET /aaaaaa HTTP/1.1
//	int first_line_end = 22;
//
//
//	if (first_line_end == -1) {
//		bpf_printk("No first line found\n");
//		return SK_PASS;
//	}
//
////	bpf_printk("First line ends at offset: %d\n", first_line_end);
//	// 先测试基本的读取功能
//	char test_buf[200] = {};
//	if (bpf_probe_read_kernel(test_buf, 200, (void *)(__u64)msg->data) == 0) {
//		bpf_printk("HTTP request: %s\n", test_buf);
//	}
//	if (test_buf[0] != 'G' && test_buf[1] != 'E' && test_buf[1] != 'T') {
//		return SK_PASS;
//	}
//	bpf_printk("data_len %d",msg->size);
//
//	// 尝试插入自定义头部
//	bpf_printk("Attempting to insert header at offset %d\n", first_line_end);
//	/*test*/
////	const char insert[] = "AA: a\r\n";   // 要插入的 7 字节
////	int ins_len = sizeof(insert) - 1;  // 不含 '\0'
////	int offset = 22;  // 偏移位置，例：在 "GET /" 后面插入
//
//	// 1) 先拉直 [0, offset) 让插入点可达（注意第三参是长度）
//	 ret = bpf_msg_pull_data(msg, 0, offset, 0);
//	if (ret) {
//		bpf_printk("pull pre off=%u len=%u ret=%ld\n", offset, offset, ret);
//		return SK_PASS;
//	}
//
//	// 2) 执行插入
//	ret = bpf_msg_push_data(msg, offset, ins_len, 0);
//	if (ret) {
//		bpf_printk("push off=%u len=%u ret=%ld\n", offset, ins_len, ret);
//		return SK_PASS;
//	}
//
//	// 插入后
////	__u32 need = msg->size;          // 拉满整条消息
//// 可选：上限保护，避免一次拉太大导致失败
////	if (need > 4096) need = 4096;    // 例如先最多拉 4KB
//
////	ret = bpf_msg_pull_data(msg, 0, need, 0);
////	if (ret) {
////		bpf_printk("pull 0..%u fail ret=%ld\n", need, ret);
////		return SK_PASS;
////	}
//
//	// 3. 写入插入的内容
////	char *data     = (char *)(long)msg->data;
////	char *data_end = (char *)(long)msg->data_end;
//
//#pragma clang loop unroll(full)
//	for (int i = 0; i < ins_len; i++) {
//		if (data + offset + i + 1 > data_end) break;
//		*(char *)(data + offset + i) = insert[i];
//	}
//
////	bpf_printk("push_data succeeded\n");
//
//	/*test2*/
//	// 参数：msg, start_offset, end_offset, flags
////	if (bpf_msg_pull_data(msg, 0, 6, 0) < 0) {
////		bpf_printk("pull_data failed\n");
////		return SK_PASS;
////	}
////
////	bpf_printk("pull_data succeeded, writing header\n");
////
////	// 测试直接修改数据
////	char *data_ptr = (char *)(long)msg->data;
////	char *data_end_ptr = (char *)(long)msg->data_end;
////
////	// 方法1：直接修改指针指向的数据
////	if (data_ptr + 1 <= data_end_ptr) {
//////		*data_ptr = 'E';
////	}
////	int offset = 5; // 第6个字符，0-based
////	if (data_ptr + offset + 1 <= data_end_ptr) {
////		*(data_ptr + offset) = 'b';
////	}
//	/*test2*/
//
////	data_ptr[6] = 'b';
////	bpf_printk("Modified data[6] to 'b' using direct pointer\n");
//
//
//	// 验证修改结果
//	char verify_buf[200] = {};
//	if (bpf_probe_read_kernel(verify_buf, 200, (void *)data) == 0) {
//		bpf_printk("Modified HTTP request: %s\n", verify_buf);
//		bpf_printk("Modified HTTP request: %d\n", (long) msg->data_end - (long) msg->data);
//	}
//
//	__u64 id = bpf_get_current_pid_tgid();
//	struct read_args *args = bpf_map_lookup_elem(&active_reads, &id);
//	if (!args) {
//		return SK_PASS;
//	}
//
//	struct l7_request_key k = {};
//	k.pid = id >> 32;
//	k.fd = args->fd;
//
//	bpf_printk("Response ------- sk msg: k.pid:%d, k.fd:%d", k.pid, k.fd);
//
//
//	return SK_PASS;
}


SEC("sk_skb/stream_verdict")
int http_hdr_inject(struct __sk_buff *skb) {
	// 检查数据包长度
	if (skb->len < 4) {
		return SK_PASS;
	}
	char buf[100] = {};
	if (bpf_skb_load_bytes(skb, 0, buf, 100) < 0) {
		return SK_PASS;
	}
//	bpf_printk("stream_verdict: ----%s",buf);

	return SK_PASS;
	// 获取socket cookie
	__u32 cookie = bpf_get_socket_cookie(skb);

	// 检查这个socket是否被标记为需要修改HTTP请求
	__u32 *flag = bpf_map_lookup_elem(&http_modify_flags, &cookie);
	if (flag){
		bpf_printk("%d",*flag);
	}
	if (flag && *flag == 1) {
		// 读取HTTP请求数据
		char buf[100] = {};
		if (bpf_skb_load_bytes(skb, 0, buf, 100) < 0) {
			return SK_PASS;
		}
//		bpf_printk("%s",buf);


		// 检查是否是HTTP请求
		if ((buf[0] == 'G' && buf[1] == 'E' && buf[2] == 'T' && buf[3] == ' ') ||
		    (buf[0] == 'P' && buf[1] == 'O' && buf[2] == 'S' && buf[3] == 'T') ||
		    (buf[0] == 'P' && buf[1] == 'U' && buf[2] == 'T' && buf[3] == ' ') ||
		    (buf[0] == 'D' && buf[1] == 'E' && buf[2] == 'L' && buf[3] == 'E') ||
		    (buf[0] == 'H' && buf[1] == 'E' && buf[2] == 'A' && buf[3] == 'D') ||
		    (buf[0] == 'O' && buf[1] == 'P' && buf[2] == 'T' && buf[3] == 'I')) {

//			bpf_printk("=== MODIFYING HTTP REQUEST ===\n");
//			bpf_printk("Socket cookie: %u\n", cookie);
//			bpf_printk("Original content: %s\n", buf);

			// 修改HTTP请求头
			if (skb->len > 20) {
				// 在HTTP头部后添加自定义头部
				char custom_header[] = "X-Custom-Header: Modified\r\n";
				if (bpf_skb_adjust_room(skb, sizeof(custom_header) - 1, BPF_ADJ_ROOM_NET, 0) == 0) {
					if (bpf_skb_store_bytes(skb, 20, custom_header, sizeof(custom_header) - 1, 0) == 0) {
						bpf_printk("Successfully added custom header to HTTP request\n");
					}
				}
			}

			// 清除标记
			__u32 clear_flag = 0;
			bpf_map_update_elem(&http_modify_flags, &cookie, &clear_flag, BPF_ANY);
			bpf_printk("=== END MODIFYING HTTP REQUEST ===\n");
		}
	}

	return SK_PASS;
}
