package main

//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -type Config proxy proxy.c

import (
 "fmt"
 "io"
 "log"
 "net"
 "os"
 "syscall"
 "time"
 "unsafe"

 "github.com/cilium/ebpf"
 "github.com/cilium/ebpf/link"
 "github.com/cilium/ebpf/rlimit"
)

const (
 CGROUP_PATH = "/sys/fs/cgroup" // Root cgroup 路径
 PROXY_PORT      = 18000 // 代理监听的端口
 SO_ORIGINAL_DST = 80 // 获取原目标地址的套接字选项
)

// SockAddrIn 是用于保存通过 SO_ORIGINAL_DST 检索的 IPv4 sockaddr_in 结构。
type SockAddrIn struct {
 SinFamily uint16
 SinPort   [2]byte
 SinAddr   [4]byte
 // 填充以匹配 sockaddr_in 大小
 Pad [8]byte
}

// getsockopt 辅助函数
func getsockopt(s int, level int, optname int, optval unsafe.Pointer, optlen *uint32) (err error) {
 _, _, e := syscall.Syscall6(syscall.SYS_GETSOCKOPT, uintptr(s), uintptr(level), uintptr(optname), uintptr(optval), uintptr(unsafe.Pointer(optlen)), 0)
 if e != 0 {
  return e
 }
 return
}

// HTTP 代理请求处理程序
func handleConnection(conn net.Conn) {
 defer conn.Close()

 // 在 Go 中需要通过 RawConn 对底层 socket 文件描述符执行低级操作。
 // 从而让我们可以用 getsockopt 来检索由 SO_ORIGINAL_DST 选项设置的原始目的地址，
 // 该信息不能通过 Go 高级网络 API 直接访问。
 rawConn, err := conn.(*net.TCPConn).SyscallConn()
 if err != nil {
  log.Printf("Failed to get raw connection: %v", err)
  return
 }

 var originalDst SockAddrIn
 // 如果 Control 不为 nil，则在创建网络连接后，在绑定到操作系统之前调用。
 rawConn.Control(func(fd uintptr) {
  optlen := uint32(unsafe.Sizeof(originalDst))
  // 通过 SO_ORIGINAL_DST 选项进行系统调用来检索原始目的地址。
  err = getsockopt(int(fd), syscall.SOL_IP, SO_ORIGINAL_DST, unsafe.Pointer(&originalDst), &optlen)
  if err != nil {
   log.Printf("getsockopt SO_ORIGINAL_DST failed: %v", err)
  }
 })

 targetAddr := net.IPv4(originalDst.SinAddr[0], originalDst.SinAddr[1], originalDst.SinAddr[2], originalDst.SinAddr[3]).String()
 targetPort := (uint16(originalDst.SinPort[0]) << 8) | uint16(originalDst.SinPort[1])

 fmt.Printf("Original destination: %s:%d\n", targetAddr, targetPort)

 // 检查原始目的地址是否可以从代理访问
 targetConn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", targetAddr, targetPort), 5*time.Second)
 if err != nil {
  log.Printf("Failed to connect to original destination: %v", err)
  return
 }
 defer targetConn.Close()

 fmt.Printf("Proxying connection from %s to %s\n", conn.RemoteAddr(), targetConn.RemoteAddr())

  // 下面的代码创建了两个数据传输通道：
  // - 从客户端到目标服务器（由单独例程处理）。
  // - 从目标服务器到客户端（由主程序处理）。
 go func() {
  _, err = io.Copy(targetConn, conn)
  if err != nil {
   log.Printf("Failed copying data to target: %v", err)
  }
 }()
 _, err = io.Copy(conn, targetConn)
 if err != nil {
  log.Printf("Failed copying data from target: %v", err)
 }
}

func main() {
 // 删除 <5.11 内核的资源限制。
 if err := rlimit.RemoveMemlock(); err != nil { 
  log.Print("Removing memlock:", err)
 }

 // 加载编译好的 eBPF ELF 并将其加载到内核中
 // 注意：也可以 pin eBPF 程序
 var objs proxyObjects
 if err := loadProxyObjects(&objs, nil); err != nil {
   log.Print("Error loading eBPF objects:", err)
 }
 defer objs.Close()

 // 将 eBPF 程序附加到 root cgroup
 connect4Link, err := link.AttachCgroup(link.CgroupOptions{
  Path:    CGROUP_PATH, 
  Attach:  ebpf.AttachCGroupInet4Connect,
  Program: objs.CgConnect4,
 })
 if err != nil {
   log.Print("Attaching CgConnect4 program to Cgroup:", err)
 }
 defer connect4Link.Close() 

 sockopsLink, err := link.AttachCgroup(link.CgroupOptions{
  Path:    CGROUP_PATH,
  Attach:  ebpf.AttachCGroupSockOps,
  Program: objs.CgSockOps,
 })
 if err != nil {
   log.Print("Attaching CgSockOps program to Cgroup:", err)
 }
 defer sockopsLink.Close() 

 sockoptLink, err := link.AttachCgroup(link.CgroupOptions{
  Path:    CGROUP_PATH,
  Attach:  ebpf.AttachCGroupGetsockopt,
  Program: objs.CgSockOpt,
 })
 if err != nil {
   log.Print("Attaching CgSockOpt program to Cgroup:", err)
 }
 defer sockoptLink.Close() 

 // 在 localhost 上启动代理服务器
 // 本例中只演示了 IPv4，但同样方法也可用于 IPv6
 proxyAddr := fmt.Sprintf("127.0.0.1:%d", PROXY_PORT)
 listener, err := net.Listen("tcp", proxyAddr)
 if err != nil {
  log.Fatalf("Failed to start proxy server: %v", err)
 }
 defer listener.Close()

 // 更新 proxyMaps 映射表以包含代理服务器配置信息，因为我们需要知道代理服务器的进程标识符，
 // 以便过滤掉由代理服务器自身生成的 eBPF 事件，从而避免其循环代理自身的数据包。
 var key uint32 = 0
 config := proxyConfig{
  ProxyPort: PROXY_PORT,
  ProxyPid: uint64(os.Getpid()),
 }
 err = objs.proxyMaps.MapConfig.Update(&key, &config, ebpf.UpdateAny)
 if err != nil {
  log.Fatalf("Failed to update proxyMaps map: %v", err)
 }

 log.Printf("Proxy server with PID %d listening on %s", os.Getpid(), proxyAddr)
 for {
  conn, err := listener.Accept()
  if err != nil {
   log.Printf("Failed to accept connection: %v", err)
   continue
  }

  go handleConnection(conn)
 }
}

//go:build ignore
#include <stddef.h>
#include <linux/bpf.h>
#include <linux/netfilter_ipv4.h>
#include <linux/in.h>
#include <sys/socket.h>

#include "bpf-builtin.h"
#include "bpf-utils.h"

#undef bpf_printk
#define bpf_printk(fmt, ...)                            \
({                                                      \
        static const char ____fmt[] = fmt;              \
        bpf_trace_printk(____fmt, sizeof(____fmt),      \
                         ##__VA_ARGS__);                \
})

#define MAX_CONNECTIONS 20000

struct Config {
  __u16 proxy_port;
  __u64 proxy_pid;
};

struct Socket {
  __u32 src_addr;
  __u16 src_port;
  __u32 dst_addr;
  __u16 dst_port;
};

struct {
  int (*type)[BPF_MAP_TYPE_ARRAY];
  int (*max_entries)[1];
  __u32 *key;
  struct Config *value;
} map_config SEC(".maps");

struct {
  int (*type)[BPF_MAP_TYPE_HASH];
  int (*max_entries)[MAX_CONNECTIONS];
  __u64 *key;
  struct Socket *value;
} map_socks SEC(".maps");

struct {
  int (*type)[BPF_MAP_TYPE_HASH];
  int (*max_entries)[MAX_CONNECTIONS];
  __u16 *key;
  __u64 *value;
} map_ports SEC(".maps");

// 当进程（在附加了该钩子的 cgroup 中）调用 connect() 系统调用时触发该钩子
// 钩子将连接重定向到透明代理，但将原始目的地址和端口存储在 map_socks 中
SEC("cgroup/connect4")
int cg_connect4(struct bpf_sock_addr *ctx) {
  // 只转发 IPv4 TCP 连接
  if (ctx->user_family != AF_INET) return 1;
  if (ctx->protocol != IPPROTO_TCP) return 1;

  // 防止代理代理自己
  __u32 key = 0;
  struct Config *conf = bpf_map_lookup_elem(&map_config, &key);
  if (!conf) return 1;
  if ((bpf_get_current_pid_tgid() >> 32) == conf->proxy_pid) return 1;

  // 该字段包含传递给 connect() 系统调用的 IPv4 地址，即连接到此套接字的目标地址及端口。
  __u32 dst_addr = ntohl(ctx->user_ip4);
  // 该字段包含传递给 connect() 系统调用的端口号
  __u16 dst_port = ntohl(ctx->user_port) >> 16;
  // 目的套接字的唯一标识符
  __u64 cookie = bpf_get_socket_cookie(ctx);

  // 在 cookie key 下存储目标套接字
  struct Socket sock;
  __builtin_memset(&sock, 0, sizeof(sock));
  sock.dst_addr = dst_addr;
  sock.dst_port = dst_port;
  bpf_map_update_elem(&map_socks, &cookie, &sock, 0);

  // 将连接重定向到代理
  ctx->user_ip4 = htonl(0x7f000001); // 127.0.0.1 == proxy IP
  ctx->user_port = htonl(conf->proxy_port << 16); // Proxy port

  bpf_printk("Redirecting client connection to proxy\n");

  return 1;
}

// 当特定 cgroup 上有套接字操作（重传超时，建立连接等）时，该程序被调用。
// 在与代理成功建立连接后，记录客户端源地址和端口
SEC("sockops")
int cg_sock_ops(struct bpf_sock_ops *ctx) {
  // 只转发 IPv4 连接
  if (ctx->family != AF_INET) return 0;

  // 已建立连接的活跃套接字
  if (ctx->op == BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB) {
    __u64 cookie = bpf_get_socket_cookie(ctx);

    // 在 map 中查找对应 cookie
    // 如果 socket 存在，存储源端口和 socket 映射
    struct Socket *sock = bpf_map_lookup_elem(&map_socks, &cookie);
    if (sock) {
      __u16 src_port = ctx->local_port;
      bpf_map_update_elem(&map_ports, &src_port, &cookie, 0);
    }
  }

  bpf_printk("sockops hook successful\n");

  return 0;
}

// 代理通过 getsockopt SO_ORIGINAL_DST 查询原始目的地信息时触发。
// 该程序利用客户端源端口从 map_ports 中获取套接字 cookie，
// 再从 map_socks 中获取原始目标信息，
// 然后与原始目标建立连接并转发客户端请求。
SEC("cgroup/getsockopt")
int cg_sock_opt(struct bpf_sockopt *ctx) {
  // SO_ORIGINAL_DST 套接字选项主要用于网络地址转换（NAT）和透明代理。
  // 在典型 NAT 或透明代理设置中，传入的数据包从原始目的地重定向到代理服务器。
  // 代理服务器在接收到数据包后，通常需要知道原始目的地址，以便以合适的方式处理流量。
  // 这就是 SO_ORIGINAL_DST 发挥作用的地方。
  if (ctx->optname != SO_ORIGINAL_DST) return 1;
  // 只转发 IPv4 TCP 连接
  if (ctx->sk->family != AF_INET) return 1;
  if (ctx->sk->protocol != IPPROTO_TCP) return 1;

  // 获取客户端源端口。
  // 实际上应该是 sk->dst_port，因为通过设置了 SO_ORIGINAL_DST 套接字选项的 getsockopt() 系统调用会获取客户端原始目的端口，所以这是在查询客户端的最终目的端口。
  __u16 src_port = ntohs(ctx->sk->dst_port);

  // 基于客户端 src_port 获取 socket cookie
  __u64 *cookie = bpf_map_lookup_elem(&map_ports, &src_port);
  if (!cookie) return 1;

  // 通过该 cookie（套接字标识符），从 map_socks 中获取原始套接字（客户端连接至目标端的套接字）
  struct Socket *sock = bpf_map_lookup_elem(&map_socks, cookie);
  if (!sock) return 1;

  struct sockaddr_in *sa = ctx->optval;
  if ((void*)(sa + 1) > ctx->optval_end) return 1;

  // 与原目标地址建立连接
  ctx->optlen = sizeof(*sa);
  sa->sin_family = ctx->sk->family; // 地址族
  sa->sin_addr.s_addr = htonl(sock->dst_addr); // 目标地址
  sa->sin_port = htons(sock->dst_port); // 目标端口
  ctx->retval = 0;

  bpf_printk("Redirecting connection to original destination\n");

  return 1;
}

char __LICENSE[] SEC("license") = "GPL";