// Copyright The CW Authors


#include "java_common.h"

static __inline struct apm_span_context *build_sc(struct ebpf_proc_info proc_info, char *host) {
	u32 key = 0;
	struct apm_span_context *cw_sc = bpf_map_lookup_elem(&apm_span_context_heap, &key);

	if (cw_sc == NULL) {
		return NULL;
	}

	struct apm_span_context *cw_psc = cw_get_parent_tracking_span();
	if (cw_psc) {
		copy_byte_arrays(cw_psc->trace_id, cw_sc->trace_id, APM_TRACE_ID_SIZE);
	}
	// new spanid
	generate_random_bytes(cw_sc->span_id, APM_SPAN_ID_SIZE);

	// set host_id/appid
	u32 k0 = 0;
	struct trace_conf_t *trace_conf = trace_conf_map__lookup(&k0);
	if (trace_conf) {
//		for (int i = 0; i < 8; ++i) {
//			cw_bpf_debug("[Client] host_id:%02x", trace_conf->host_id[i]);
//		}
		copy_byte_arrays(trace_conf->host_id, cw_sc->host_id, APM_HOST_ID_SIZE);
//		copy_byte_arrays(trace_conf->app_id, cw_sc->app_id, APM_APP_ID_SIZE);
	}

//	__u64 pid_tgid = bpf_get_current_pid_tgid();
//	__u32 tgid = pid_tgid >> 32;
//	struct ebpf_proc_info *proc_info =
//			bpf_map_lookup_elem(&proc_info_map, &tgid);
//	if (proc_info) {
	copy_byte_arrays(proc_info.instance_id, cw_sc->instance_id, APM_APP_ID_SIZE);
	copy_byte_arrays(proc_info.app_id, cw_sc->app_id, APM_APP_ID_SIZE);
//	}

	// set assumed_app_id
	set_assumed_app_id_arrays(host, cw_sc->assumed_app_id, APM_ASSUMED_APP_ID_STRING_SIZE);
	cw_save_current_tracking_span(cw_sc);
	return cw_sc;
//	char header[CW_STREAM_HEADER_LEN];
//	span_context_to_cw_string_stream(cw_span_context, header, '1');
//	cw_bpf_debug("Successfully HEADER  %s\n", header);
//	return header;
}

SEC("uprobe/Java_java_net_SocketOutputStream_socketWrite0")
int uprobe_Java_java_net_SocketOutputStream_socketWrite0(struct pt_regs *ctx) {
	// 捕获第六个参数 data_count
	int data_count = PT_REGS_PARM6(ctx);
	if (data_count < MIN_LEN) {
		return 0;
	}

	cw_bpf_debug("[java client] data_count:%d\n", data_count);

	void *len_from_rbp_ptr = (void *) PT_LEN_REGS(ctx);

	int len_val = 0;
	cw_bpf_debug("[java client] address: len_from_rbp_ptr<0x%lx>\n", len_from_rbp_ptr);

	bpf_probe_read(&len_val, sizeof(len_val), (void *) len_from_rbp_ptr);
	cw_bpf_debug("[java client] [len_val] before addr<0x%lx>, %d \n", len_from_rbp_ptr, len_val);
	if (len_val != data_count) {
		cw_bpf_debug("len check failed: len_val=%d data_count=%d", len_val, data_count);
		return 0;
	}

	int key = 0;
	struct sock_t *map_data = bpf_map_lookup_elem(&socket_heap, &key);
	if (map_data == NULL) {
		cw_bpf_debug("[java client] Failed to lookup socket_heap\n");
		return 1;
	}

	__builtin_memset(map_data->header_stream, 0, sizeof(map_data->header_stream));
	__builtin_memset(map_data->host, 0, sizeof(map_data->host));
	// 获取jbytearray
	__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 0;
	}

	// cw_bpf_debug("[java client] code_type:%d\n", map_data->proc_info->code_type);
	void *jbytearray_ptr = NULL;

	// 1003 是 javaAOT 的
	if (proc_info->code_type == CodeTypeJavaAot) {
		jbytearray_ptr = (void *) (PT_REGS_SP(ctx)) + 0x8 + 0x10060 + 0x20;
	} else {
		jbytearray_ptr = (void *) PT_REGS_PARM4(ctx);
	}

	if(!jbytearray_ptr){
		return 0;
	}
	// cw_bpf_debug("[java client] jbytechar_ptr_from_rcx <0x%lx>", jbytearray_ptr);
	
	unsigned long jbytechar_head_ptr;
	// 读取一级指针
	long ret = bpf_probe_read_user(&jbytechar_head_ptr, sizeof(unsigned long), (void *) jbytearray_ptr);
	// cw_bpf_debug("[java client] [jbytechar_head_ptr] <0x%lx>", jbytechar_head_ptr);
	
	if (ret != 0) {
		cw_bpf_debug("[java client] Failed to read first level ptr: %d\n", ret);
		return 0;
	}

	// 检查一级指针是否有效
	if (!jbytechar_head_ptr) {
		cw_bpf_debug("[java client] First level pointer is null.\n");
		return 0;
	}


	// 读取用户空间数据到 map_data->payload
	void *jbytechar_ptr = (void *) (jbytechar_head_ptr + 16);
	// cw_bpf_debug("[java client] [jbytechar_ptr] <0x%lx>", jbytechar_ptr);
	long err = bpf_probe_read_user_str(map_data->payload, sizeof(map_data->payload), jbytechar_ptr);

	if (err < 0) {
		cw_bpf_debug("[java client] bpf_probe_read_user failed with return code: %d\n", err);
		return 0;
	}
	map_data->size = data_count;
	map_data->payload_char_p = (void *) jbytechar_ptr;
	map_data->payload_len_p = (void *) len_from_rbp_ptr;

	// http协议过滤
	if (!is_http_request2(map_data->payload, map_data->size)) {
		return -1;
	}

	// 查找 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;
		}
	}

	if (map_data->header_offset_idx == 0) {
		cw_bpf_debug("[java client] header_offset_idx error:\n");
		return 1;
	}

	map_data->end_str_len = data_count - map_data->header_offset_idx + 1;
	u32 copy_size = L7_IOVEC_BUF_SIZE;
	if (map_data->end_str_len > copy_size) {
		u32 chunk = (u32) map_data->end_str_len / copy_size;
		if (chunk > MAX_CHUNK) {
			cw_bpf_debug("[java client] Exceeding the max_chunk:%d need_chunk:%d\n", MAX_CHUNK, chunk);
			return 1;
		}
	}
	// cw_bpf_debug("[java client] header_offset_idx:%d\n", map_data->header_offset_idx);

	// tail payload保存
	u32 tail_key = 1;
	char *tail_payload = bpf_map_lookup_elem(&large_array_map, &tail_key);
	if (!tail_payload) {
		// cw_bpf_debug("[java client ]header ->end %s\n", tail_payload);
		return 1;
	}
	__builtin_memset(tail_payload, 0, sizeof(char *));

	bpf_probe_read_user_str(tail_payload, JAVA_MAX_BUFFER_SIZE,
	                        (void *) (jbytechar_ptr + map_data->header_offset_idx));
	// host查询
	bpf_tail_call(ctx, &NAME(progs_jmp_up_map), PROG_DATA_JAVA_FIND_HOST_UP_IDX);
	return 0;

/*设置初始值*/
	/*
	 * debug
	char *full_payload= bpf_map_lookup_elem(&large_array_map, &key);
	if (!full_payload)
		return -1;
	// 读取完整http
	bpf_probe_read_user_str(full_payload, JAVA_MAX_BUFFER_SIZE, jbytechar_ptr);
	cw_bpf_debug("[java client]origin full_payload %s\n", full_payload);
*/

	// 尾部payload
//	u32 key2 = 1;
//	char *tail_payload = bpf_map_lookup_elem(&large_array_map, &key2);
//	if (!tail_payload) {
//		cw_bpf_debug("[java client ]header ->end %s\n", tail_payload);
//		return -1;
//	}
//
//	u32 copy_size = L7_IOVEC_BUF_SIZE;
//	long end_str_len = bpf_probe_read_user_str(tail_payload, JAVA_MAX_BUFFER_SIZE,
//	                                           (void *) (jbytechar_ptr + map_data->header_offset_idx));
//	cw_bpf_debug("[java client] end_str_len %d\n", end_str_len);
//	map_data->end_str_len = end_str_len;

//	u32 copy_size = L7_IOVEC_BUF_SIZE;
//	if (end_str_len > copy_size) {
//		u32 chunk = (u32) end_str_len / copy_size;
//		if (chunk > MAX_CHUNK) {
//			cw_bpf_debug("[java client] Exceeding the max_chunk:%d need_chunk:%d\n", MAX_CHUNK, chunk);
//			return -1;
//		}
//	}



//	char host[30];
//	bpf_probe_read(&host, sizeof host, map_data->payload + map_data->host_offset_idx);
//	host[29] = '\0';
//	cw_bpf_debug("[java client] http host %s\n", host);
//	char tmp_header[CW_STREAM_HEADER_LEN];
//	struct apm_span_context *cw_sc = build_sc(proc_info, host);
//	if (cw_sc == NULL){
//		return 0;
//	}
//	span_context_to_cw_string_stream(cw_sc, map_data->header_stream, '1');
//	return 1;
//	cw_bpf_debug("header:[%s],sizeof %d ",map_data->header_stream,sizeof(map_data->header_stream));


//	bpf_probe_read(map_data->header_stream, sizeof(map_data->header_stream), (void *) (tmp_header));
//	cw_bpf_debug("stream_header1->header------%s",map_data->header_stream);

//	if (insert_pos >= 0 && insert_pos <= JAVA_MAX_BUFFER_SIZE - sizeof(map_data->header_stream)) {
//		__builtin_memcpy(full_payload + insert_pos, map_data->header_stream, sizeof(map_data->header_stream));
//	}



	// 修改长度
//	long new_val = map_data->size + HEADER_LEN;
//
//	long res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val, sizeof(new_val));
//	if (res < 0) {
//		cw_bpf_debug("[java client] Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);
////		cw_bpf_debug("Failed to write len to user address: %p\n", len_from_rbp_ptr);
//		return -1;
//	}

//	u32 insert_pos = map_data->header_offset_idx;
//	u32 tail_offset_idx = insert_pos + sizeof(map_data->header_stream) - 1;
//
//	// 单次
//	if (end_str_len <= copy_size) {
//		// 写头
//		bpf_probe_write_user((void *) jbytechar_ptr + insert_pos, &map_data->header_stream, sizeof(map_data->header_stream));
//		// 写尾
//		bpf_probe_write_user((void *) jbytechar_ptr + tail_offset_idx, tail_payload, L7_IOVEC_BUF_SIZE);
//	}
//
//	// 多批
//	if (end_str_len > copy_size) {
////		bpf_tail_call(ctx, &jmp_table1, 1);
//		bpf_tail_call(ctx, &NAME(progs_jmp_up_map), PROG_DATA_JAVA_UPDATE_HEADER_UP_IDX);
//	}

	bpf_tail_call(ctx, &NAME(progs_jmp_up_map), PROG_DATA_JAVA_UPDATE_HEADER_UP_IDX);

	return 0;
}

//bpf_prog_up_build_header
PROGUP(java_find_host)(struct pt_regs *ctx) {
	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 1;
	}

	// host查询
	long host_native = 0x203a74736f480a0dLL; // 小端序下的 "\r\nHost: " ('0d' '0a' '48' '6f' '73' '74' '3a' '20')
#pragma clang loop unroll(full)
	for (u32 i = 0; i < 900 - 8; i++) {
		long data = *(long long *) (map_data->payload + i);
		if (data == host_native) {
			map_data->host_offset_idx = i + 8;
			break;
		}
	}
	if (map_data->host_offset_idx == 0) {
		cw_bpf_debug("[java client] host_offset_idx error:\n");
		return 1;
	}


	bpf_probe_read_user(&map_data->host, sizeof map_data->host, map_data->payload + map_data->host_offset_idx);
	// Build Header
	bpf_tail_call(ctx, &NAME(progs_jmp_up_map), PROG_DATA_JAVA_BUILD_HEADER_UP_IDX);
	return 0;
}

// bpf_prog_up_build_header
PROGUP(java_build_header)(struct pt_regs *ctx) {
	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 1;
	}

	__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 0;
	}

	for (u32 i = 0; i < 30; i++) {
		if (map_data->host[i] == '\r') {
			map_data->host[i] = '\0';
			break;
		}
	}
	cw_bpf_debug("[java client] http host %s\n", map_data->host);

	struct apm_span_context *cw_sc = build_sc(*proc_info, map_data->host);
	if (cw_sc == NULL) {
		return 0;
	}
	span_context_to_cw_string_stream(cw_sc, map_data->header_stream, '1');
	cw_bpf_debug("[java client] header:[%s],sizeof %d ", map_data->header_stream, sizeof(map_data->header_stream));
	// 更新Header
	bpf_tail_call(ctx, &NAME(progs_jmp_up_map), PROG_DATA_JAVA_UPDATE_HEADER_UP_IDX);
	return 0;
}

// bpf_prog_up_update_header
PROGUP(java_update_header)(struct pt_regs *ctx) {
	u32 k0 = 0;
	u32 copy_size = L7_IOVEC_BUF_SIZE;

	u32 k1 = 1;
	char *tail_payload = bpf_map_lookup_elem(&large_array_map, &k1);
	if (!tail_payload) {
		cw_bpf_debug("[java client] header ->end %s\n", tail_payload);
		return -1;
	}

	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 -1;
	}

	u32 insert_pos = map_data->header_offset_idx;
	u32 tail_offset_idx = insert_pos + sizeof(map_data->header_stream) - 1;
	u32 end_str_len = map_data->end_str_len;
	void *jbytechar_ptr = map_data->payload_char_p;
	void *len_from_rbp_ptr = map_data->payload_len_p;

	if (!jbytechar_ptr) {
		return -1;
	}

	long new_val = map_data->size + HEADER_LEN;

	long res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val, sizeof(new_val));
	if (res < 0) {
		cw_bpf_debug("[java client] Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);
		return -1;
	}

	// 单次
	if (end_str_len <= copy_size) {
		// 写头
		bpf_probe_write_user((void *) jbytechar_ptr + insert_pos, &map_data->header_stream,
		                     sizeof(map_data->header_stream));
		// 写尾
		bpf_probe_write_user((void *) jbytechar_ptr + tail_offset_idx, tail_payload, L7_IOVEC_BUF_SIZE);
		return 0;
	}


	// 多批
	u32 chunk = end_str_len / copy_size;
	bpf_probe_write_user((void *) jbytechar_ptr + insert_pos, &map_data->header_stream,
	                     sizeof(map_data->header_stream));
	for (u32 i = 0; i < MAX_CHUNK; i++) {
		if (i <= chunk) {
			tail_offset_idx = insert_pos + sizeof(map_data->header_stream) - 1 + i * copy_size;
			bpf_probe_write_user((void *) jbytechar_ptr + tail_offset_idx, tail_payload + i * copy_size,
			                     L7_IOVEC_BUF_SIZE);
		}
	}

	return 0;
}

SEC("uprobe/javaaot_asmnop")
int javaaot_asmnop(struct pt_regs *ctx) {
	bpf_printk("enter the ret javaaot_asmnop\n");
	unsigned long jhttpdata_ptr;
#if defined(__x86_64__)
	jhttpdata_ptr = (ctx)->rsi;
	unsigned long len_from_rbp_ptr = (ctx)->rsp-0x8;
#elif defined(__aarch64__)
	jhttpdata_ptr = (ctx)->regs[1];
	unsigned long len_from_rbp_ptr = (ctx)->regs[16]-0x8;
#endif
	bpf_printk("[len_from_rbp_ptr] <0x%lx>", len_from_rbp_ptr);

	// 捕获第六个参数 data_count
	int data_count = PT_REGS_PARM3(ctx);
	if (data_count < MIN_LEN) {
		return 0;
	}

	cw_bpf_debug("[java client] data_count:%d\n", data_count);

	void *len_ptr = 0;
	cw_bpf_debug("[java client] address: len_from_rbp_ptr<0x%lx>\n", len_from_rbp_ptr);

	bpf_probe_read(&len_ptr, sizeof(len_ptr), (void *) len_from_rbp_ptr);
	cw_bpf_debug("[java client] [len_ptr] before addr<0x%lx>, %d \n", len_from_rbp_ptr, len_ptr);

	if ((long) len_ptr != data_count) {
		cw_bpf_debug("[java client] [len_ptr] check error.");
		return 0;
	}

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

	// 获取jbytearray
	__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 0;
	}

	// 读取用户空间数据到 map_data->payload
	void *jbytechar_ptr = (void *) jhttpdata_ptr;
	// cw_bpf_debug("[java client] [jbytechar_ptr] <0x%lx>", jbytechar_ptr);
	long err = bpf_probe_read_user_str(map_data->payload, sizeof(map_data->payload), jbytechar_ptr);

	if (err < 0) {
		cw_bpf_debug("[java client] bpf_probe_read_user failed with return code: %d\n", err);
		return 0;
	}
	map_data->size = data_count;
	map_data->payload_char_p = (void *) jbytechar_ptr;
	map_data->payload_len_p = (void *) len_from_rbp_ptr;

	// http协议过滤
	if (!is_http_request2(map_data->payload, map_data->size)) {
		return -1;
	}

	// 查找 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;
		}
	}

	if (map_data->header_offset_idx == 0) {
		cw_bpf_debug("[java client] header_offset_idx error:\n");
		return 1;
	}

	map_data->end_str_len = data_count - map_data->header_offset_idx + 1;
	u32 copy_size = L7_IOVEC_BUF_SIZE;
	if (map_data->end_str_len > copy_size) {
		u32 chunk = (u32) map_data->end_str_len / copy_size;
		if (chunk > MAX_CHUNK) {
			cw_bpf_debug("[java client] Exceeding the max_chunk:%d need_chunk:%d\n", MAX_CHUNK, chunk);
			return 1;
		}
	}
	// cw_bpf_debug("[java client] header_offset_idx:%d\n", map_data->header_offset_idx);

	// tail payload保存
	u32 tail_key = 1;
	char *tail_payload = bpf_map_lookup_elem(&large_array_map, &tail_key);
	if (!tail_payload) {
		// cw_bpf_debug("[java client ]header ->end %s\n", tail_payload);
		return 1;
	}
	__builtin_memset(tail_payload, 0, sizeof(char *));

	bpf_probe_read_user_str(tail_payload, JAVA_MAX_BUFFER_SIZE,
	                        (void *) (jbytechar_ptr + map_data->header_offset_idx));
	// host查询
	bpf_tail_call(ctx, &NAME(progs_jmp_up_map), PROG_DATA_JAVA_FIND_HOST_UP_IDX);
	return 0;

	bpf_tail_call(ctx, &NAME(progs_jmp_up_map), PROG_DATA_JAVA_UPDATE_HEADER_UP_IDX);

	return 0;
}