// Copyright The CW Authors


#include "java_common.h"
static __inline int updataSocket2(struct sock_t *map_data_res, char * payload,int len,void * jbytechar_ptr,void * len_from_rbp_ptr) {

	long res = bpf_probe_write_user((void *) jbytechar_ptr,  &map_data_res->payload, sizeof(map_data_res->payload));
	bpf_printk("sizeof(map_data_res->payload) %d\n", sizeof(map_data_res->payload));
	bpf_printk("sizeof(payload) %d\n", sizeof(payload));
	bpf_printk("&payload 0x%lx\n", &payload);

	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", jbytechar_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", jbytechar_ptr, res);
	}

	unsigned long new_val;
	new_val = len;
	res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val, sizeof(new_val));
	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", len_from_rbp_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);
	}

	bpf_printk("len %d\n", len);

//	bpf_printk("payload %s\n", payload);
//	for (int i = 270; i < len; ++i) {
//		bpf_printk("data[%d]=%c", i,payload[i]);
//		if(payload[i]=='\0'){
//			break;
//		}
//	}
	return 0;
}

static __inline int updataSocket(struct sock_t *map_data_res,void * jbytechar_ptr,void * len_from_rbp_ptr) {

	long res = bpf_probe_write_user((void *) jbytechar_ptr, &map_data_res->payload, sizeof(map_data_res->payload));
	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", jbytechar_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", jbytechar_ptr, res);
	}

	unsigned long new_val;
	new_val = map_data_res->size;
	res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val, sizeof(new_val));
	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", len_from_rbp_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);
	}

	bpf_printk("Successfully %d\n", map_data_res->size);

	bpf_printk("Successfully %s\n", map_data_res->payload);
//	for (int i = 270; i < 290; ++i) {
//		bpf_printk("data[%d]=%c", i,map_data_res->payload[i]);
//		if(map_data_res->payload[i]=='\0'){
//			break;
//		}
//	}
	return 0;
}

static __inline struct sock_t* buildHeader(struct sock_t *map_data) {

	int key = 0;
	struct sock_t *map_data_res = bpf_map_lookup_elem(&socket_res, &key);

	if (!map_data_res) {
		bpf_printk("Failed to lookup socket_heap");
		return NULL;
	}

//	__builtin_memset(map_data_res, 0, sizeof(struct sock_t));
//	char header[5] = "cwt\r\n";

//	char header[CW_STREAM_HEADER_LEN] = "cwtrace: 00:00:1015481350055581:5450531005555981:5610250100539899:304775019cd3218a304775019cd3218a:1001025098564810:140acc88cde8773f\r\n";
	struct apm_span_context *cw_span_context = bpf_map_lookup_elem(&apm_span_context_heap, &key);

	if (cw_span_context == 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_span_context->trace_id, APM_TRACE_ID_SIZE);
		// new spanid
		generate_random_bytes(cw_span_context->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_span_context->host_id, APM_HOST_ID_SIZE);
		copy_byte_arrays(trace_conf->app_id, cw_span_context->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) {
		for (int i = 0; i < 8; ++i) {
//			cw_bpf_debug("[Client] instance_id:%02x", proc_info->instance_id[i]);
		}
		copy_byte_arrays(proc_info->instance_id, cw_span_context->instance_id, APM_APP_ID_SIZE);
	}

	// set assumed_app_id
	set_assumed_app_id_arrays("111", cw_span_context->assumed_app_id, APM_ASSUMED_APP_ID_STRING_SIZE);
	cw_save_current_tracking_span(cw_span_context);

	char header[CW_STREAM_HEADER_LEN];
	span_context_to_cw_string_stream(cw_span_context, header, '1');
	bpf_printk("Successfully HEADER  %s\n", header);

//	bpf_probe_read(map_data_res->payload, map_data->header_offset_idx, map_data->payload);
//	bpf_printk("Successfully %s\n", data);
//	return 1;
#pragma unroll
	for (int i = 0; i < MAX_LEN; i++) {
		if (i < map_data->header_offset_idx) {
			map_data_res->payload[i] = map_data->payload[i];
		} else {
			if (i == map_data->header_offset_idx) {
#pragma unroll
				for (int k = 0; k < CW_STREAM_HEADER_LEN; k++) {
					int tmp_len = i + k;
					if (tmp_len < MAX_LEN) {
						map_data_res->payload[tmp_len] = header[k];
					}
				}
			}
			int tmp_len = i + CW_STREAM_HEADER_LEN;
			if (tmp_len < MAX_LEN) {
				map_data_res->payload[tmp_len] = map_data->payload[i];
				if (map_data->payload[i] == '\0')
					break;
//				bpf_printk("map_data->payload: %c->i=%d, offindex %d", map_data_res->payload[i], i, map_data->header_offset_idx);
			}
		}
	}
	map_data_res->size = map_data->size + CW_STREAM_HEADER_LEN;

	return map_data_res;
}

static __inline int insertHeader(struct sock_t *map_data,void * jbytechar_ptr,void * len_from_rbp_ptr) {

	struct sock_t *map_data_res= buildHeader(map_data);
	if (map_data_res == NULL) {
		bpf_printk("Failed to lookup socket_heap");
		return -1;
	}

	/*修改部分*/
	return updataSocket2(map_data_res, map_data_res->payload, map_data_res->size, jbytechar_ptr, len_from_rbp_ptr);
}


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

	bpf_printk("--------");


//	void *len_from_rsp_ptr = (void *) PT_REGS_RSP(ctx) - 65640 +8 ;
	void *len_from_rbp_ptr = (void *) (PT_REGS_RBP(ctx)) - 0x10058;
//	void *jbytearray_from_rbp_p_p_p = (void *) (PT_REGS_RBP(ctx)) - 0x10050;

//	long res;
	void *len_ptr = 0;
	bpf_printk("address: len_from_rbp_ptr<0x%lx>\n", len_from_rbp_ptr);
//	bpf_printk("len_from_rsp_ptr address: 0x%lx\n", len_from_rsp_ptr);

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

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

	// 获取jbytearray
	void *jbytearray_ptr = (void *) PT_REGS_PARM4(ctx);
	bpf_printk("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);
	bpf_printk("[jbytechar_head_ptr] <0x%lx>", jbytechar_head_ptr);

	if (ret != 0) {
		bpf_printk("Failed to read first level ptr: %d\n", ret);
		return 0;
	}

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

	// 定义 Map 键值和 Map 数据结构
	int key = 0;
	struct sock_t *map_data = bpf_map_lookup_elem(&socket_heap, &key);
	if (!map_data) {
		bpf_printk("Failed to lookup socket_heap\n");
		return 1;
	}
//	__builtin_memset(map_data, 0, sizeof(struct sock_t));
	// 读取用户空间数据到 map_data->payload
	void *jbytechar_ptr = (void *) (jbytechar_head_ptr + 16);
	bpf_printk("[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) {
		bpf_printk("bpf_probe_read_user failed with return code: %d\n", err);
		return 0;
	}
	map_data->size = data_count;


	// 查找 Header 开始位置
	for (int i = 0; i < MAX_LEN - 9; i++) {
		if (map_data->payload[i] == '1' &&
		    map_data->payload[i + 1] == '.' &&
		    map_data->payload[i + 2] == '1' &&
		    map_data->payload[i + 3] == '\r' &&
		    map_data->payload[i + 4] == '\n') {
			map_data->header_offset_idx = i + 5;
			break;
		}
	}

	if (map_data->header_offset_idx == 0) {
		return -1;
	}

	// http协议过滤
	if (!is_http_request2(map_data->payload, map_data->size)) {
		return -1;
	}
//	char syscall_infer_buf[100];
///*实现2*/
/*实现2
//	int key = 0;
	char *data;
	data = bpf_map_lookup_elem(&large_array_map, &key);
	if (!data)
		return -1;
	// 读取完整http
	long aaa = bpf_probe_read_user_str(data, MAX_BUFFER_SIZE, jbytechar_ptr);
	bpf_printk("data %s\n", data);

	char *data2;
	int key2 = 1;
	data2 = bpf_map_lookup_elem(&large_array_map, &key2);
	if (!data2) {
		bpf_printk("data2 %s\n", data2);
		return -1;
	}

	long end_str_len = bpf_probe_read_user_str(data2, MAX_BUFFER_SIZE, (void *) (jbytechar_ptr + map_data->header_offset_idx));
	bpf_printk("end_str_len %d\n", end_str_len);

	// 挪动数据以腾出空间插入子串
	int insert_pos = map_data->header_offset_idx;

	char subs[CW_STREAM_HEADER_LEN+1] = "cwtrace: 00:00:1015481350055581:5450531005555981:5610250100539899:304775019cd3218a304775019cd3218a:1001025098564810:140acc88cde8773f\r\n\0";
	if (insert_pos >= 0 && insert_pos <= MAX_BUFFER_SIZE - sizeof(subs)) {
		__builtin_memcpy(data + insert_pos, subs, sizeof(subs));
	}

	int insert_pos2 = insert_pos + sizeof(subs) - 1;
	bpf_printk("sizeof(data2) %d\n", sizeof(data2));


	int copy_size = 375;
	int max_chunk = 40;
//	if (end_str_len > copy_size) {
//		int chunk = end_str_len / copy_size ; // 除以常量10
//#pragma unroll
//		for (int i = 0; i < max_chunk; i++) {
//			if (i <= chunk) {
//				insert_pos2 = insert_pos + sizeof(subs) - 1 + i*copy_size;
//				bpf_printk("insert_pos2:%d\n", insert_pos2);
//				bpf_printk("chunk i:%i %d\n", i, chunk);
//				if (insert_pos2 >= 0 && insert_pos2 <= MAX_BUFFER_SIZE - copy_size) {
//					__builtin_memcpy(data + insert_pos2, data2 + i*copy_size, copy_size);
//				}
//			}
//		}
//	}

	if (end_str_len <= copy_size) {
		if (insert_pos2 >= 0 && insert_pos2 <= MAX_BUFFER_SIZE - copy_size){
			bpf_printk("nochunk %d\n", end_str_len);
			__builtin_memcpy(data + insert_pos2, data2, copy_size);
		}
	}
	bpf_printk("Successfully %s\n", data);
	for (int i = 270; i < 290; ++i) {
		bpf_printk("data[%d]=%c", i,data[i]);
		if(data[i]=='\0'){
			break;
		}
	}

	res = bpf_probe_write_user((void *) jbytechar_ptr, &data, sizeof(data[MAX_BUFFER_SIZE]));
	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", jbytechar_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", jbytechar_ptr, res);
	}

	unsigned long new_val2;
	new_val2 = map_data->size + CW_STREAM_HEADER_LEN;
	bpf_printk("Successfully %d\n", new_val2);

	res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val2, sizeof(new_val2));
	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", len_from_rbp_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);
	}


	return 2;

	实现2-end */
///*实现2-end*/

/*实现1*/

	insertHeader(map_data,jbytechar_ptr,len_from_rbp_ptr);


//	updataSocket(map_data_res,jbytechar_ptr,len_from_rbp_ptr);

	return 0;
}
