mirror of
https://github.com/safedep/pmg.git
synced 2026-08-03 07:24:09 +02:00
implement decision logic
This commit is contained in:
Sahilb315
+161
-41
@@ -1,4 +1,4 @@
|
||||
//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -type event bpf connect.c -- -I/usr/include/aarch64-linux-gnu
|
||||
//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -type event -type target bpf connect.c -- -I/usr/include/aarch64-linux-gnu
|
||||
|
||||
package main
|
||||
|
||||
@@ -7,10 +7,13 @@ import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
|
||||
"github.com/cilium/ebpf"
|
||||
@@ -19,46 +22,65 @@ import (
|
||||
"github.com/cilium/ebpf/rlimit"
|
||||
)
|
||||
|
||||
// Mirrors the ACTION_* constants in connect.c.
|
||||
var actionNames = map[uint8]string{
|
||||
0: "REDIRECT",
|
||||
1: "skip/proto",
|
||||
2: "skip/loopback",
|
||||
3: "skip/exempt",
|
||||
4: "skip/dport",
|
||||
5: "skip/no-target",
|
||||
}
|
||||
|
||||
func main() {
|
||||
err := rlimit.RemoveMemlock()
|
||||
if err != nil {
|
||||
fmt.Printf("Failed to remove mem lock: %v\n", err)
|
||||
return
|
||||
proxyFlag := flag.String("proxy", "", "redirect target as host:port, e.g. 127.0.0.1:8443")
|
||||
exemptFlag := flag.String("exempt-uid", "", "comma separated uids that are never redirected")
|
||||
tcpOnlyFlag := flag.Bool("tcp-only", false, "only print TCP events, hiding DNS and route probe noise")
|
||||
flag.Parse()
|
||||
|
||||
if err := run(*proxyFlag, *exemptFlag, *tcpOnlyFlag); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "pmgwatch: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func run(proxyAddr, exemptUIDs string, tcpOnly bool) error {
|
||||
if err := rlimit.RemoveMemlock(); err != nil {
|
||||
return fmt.Errorf("remove memlock: %w", err)
|
||||
}
|
||||
|
||||
var objs bpfObjects
|
||||
err = loadBpfObjects(&objs, nil)
|
||||
if err != nil {
|
||||
fmt.Printf("Failed to load bpf objects: %v\n", err)
|
||||
return
|
||||
if err := loadBpfObjects(&objs, nil); err != nil {
|
||||
return fmt.Errorf("load bpf objects: %w", err)
|
||||
}
|
||||
|
||||
defer objs.Close()
|
||||
|
||||
rootCgroup := "/sys/fs/cgroup"
|
||||
if err := configureTarget(&objs, proxyAddr); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := configureExempt(&objs, exemptUIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
l, err := link.AttachCgroup(link.CgroupOptions{
|
||||
Path: rootCgroup,
|
||||
Path: "/sys/fs/cgroup",
|
||||
Attach: ebpf.AttachCGroupInet4Connect,
|
||||
Program: objs.Connect4,
|
||||
})
|
||||
if err != nil {
|
||||
fmt.Printf("Failed to attach cgroup: %v\n", err)
|
||||
return
|
||||
return fmt.Errorf("attach cgroup: %w", err)
|
||||
}
|
||||
defer l.Close()
|
||||
|
||||
rd, err := ringbuf.NewReader(objs.Events)
|
||||
if err != nil {
|
||||
fmt.Printf("Failed to create ringbuf reader: %v\n", err)
|
||||
return
|
||||
return fmt.Errorf("open ringbuf: %w", err)
|
||||
}
|
||||
defer rd.Close()
|
||||
|
||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||||
defer stop() // Cleans up resources allocated by the signal package
|
||||
|
||||
fmt.Println("Application started. Press Ctrl+C to exit.")
|
||||
defer stop()
|
||||
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
@@ -66,41 +88,139 @@ func main() {
|
||||
rd.Close()
|
||||
}()
|
||||
|
||||
fmt.Println("Shadow mode: decisions are recorded, nothing is redirected. Ctrl+C to exit.")
|
||||
fmt.Printf("%-15s %-17s %-7s %-8s %-22s %s\n", "ACTION", "COMMAND", "UID", "PID", "DESTINATION", "PROTO")
|
||||
|
||||
return drain(rd, tcpOnly)
|
||||
}
|
||||
|
||||
func drain(rd *ringbuf.Reader, tcpOnly bool) error {
|
||||
var e bpfEvent
|
||||
|
||||
for {
|
||||
rec, err := rd.Read()
|
||||
if err != nil {
|
||||
if errors.Is(err, ringbuf.ErrClosed) {
|
||||
return
|
||||
return nil
|
||||
}
|
||||
fmt.Printf("failed to read raw event: %v\n", err)
|
||||
|
||||
fmt.Fprintf(os.Stderr, "read event: %v\n", err)
|
||||
continue
|
||||
}
|
||||
|
||||
err = binary.Read(bytes.NewReader(rec.RawSample), binary.LittleEndian, &e)
|
||||
if err != nil {
|
||||
fmt.Printf("failed to read event: %v\n", err)
|
||||
if err := binary.Read(bytes.NewReader(rec.RawSample), binary.LittleEndian, &e); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "decode event: %v\n", err)
|
||||
continue
|
||||
}
|
||||
|
||||
ip := make(net.IP, 4)
|
||||
binary.LittleEndian.PutUint32(ip, e.Daddr)
|
||||
|
||||
name := e.Comm[:]
|
||||
if i := bytes.IndexByte(name, 0); i != -1 {
|
||||
name = name[:i] // keep bytes before the first NUL
|
||||
}
|
||||
comm := string(name)
|
||||
|
||||
proto := ""
|
||||
switch e.Proto {
|
||||
case syscall.IPPROTO_TCP:
|
||||
proto = "TCP"
|
||||
case syscall.IPPROTO_UDP:
|
||||
proto = "UDP"
|
||||
if tcpOnly && e.Proto != syscall.IPPROTO_TCP {
|
||||
continue
|
||||
}
|
||||
|
||||
fmt.Println()
|
||||
fmt.Printf("PID: %v\nUID: %v\ndAddr: %v\ndPort: %v\nProto: %v\nCommand: %v\n", e.Pid, e.Uid, ip, e.Dport, proto, comm)
|
||||
printEvent(&e)
|
||||
}
|
||||
}
|
||||
|
||||
func printEvent(e *bpfEvent) {
|
||||
ip := make(net.IP, 4)
|
||||
binary.LittleEndian.PutUint32(ip, e.Daddr)
|
||||
|
||||
action, ok := actionNames[e.Action]
|
||||
if !ok {
|
||||
action = fmt.Sprintf("action/%d", e.Action)
|
||||
}
|
||||
|
||||
fmt.Printf("%-15s %-17s %-7d %-8d %-22s %s\n",
|
||||
action, comm(e.Comm[:]), e.Uid, e.Pid,
|
||||
net.JoinHostPort(ip.String(), strconv.Itoa(int(e.Dport))), proto(e.Proto))
|
||||
}
|
||||
|
||||
// comm is a fixed 16 byte buffer, NUL terminated when the name is shorter.
|
||||
func comm(raw []byte) string {
|
||||
if i := bytes.IndexByte(raw, 0); i != -1 {
|
||||
raw = raw[:i]
|
||||
}
|
||||
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
func proto(p uint8) string {
|
||||
switch p {
|
||||
case syscall.IPPROTO_TCP:
|
||||
return "TCP"
|
||||
case syscall.IPPROTO_UDP:
|
||||
return "UDP"
|
||||
default:
|
||||
return strconv.Itoa(int(p))
|
||||
}
|
||||
}
|
||||
|
||||
// configureTarget writes the redirect destination the hook will use. Leaving it
|
||||
// unset is valid: every otherwise eligible connection then reports
|
||||
// skip/no-target, which is how shadow mode runs before a proxy exists.
|
||||
func configureTarget(objs *bpfObjects, addr string) error {
|
||||
if addr == "" {
|
||||
fmt.Println("No redirect target set, eligible connections will report skip/no-target")
|
||||
return nil
|
||||
}
|
||||
|
||||
host, portStr, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse proxy address %q: %w", addr, err)
|
||||
}
|
||||
|
||||
ip := net.ParseIP(host)
|
||||
if ip == nil || ip.To4() == nil {
|
||||
return fmt.Errorf("proxy address %q must be an IPv4 address", addr)
|
||||
}
|
||||
|
||||
port, err := strconv.ParseUint(portStr, 10, 16)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse proxy port %q: %w", portStr, err)
|
||||
}
|
||||
|
||||
// The kernel holds the address in network byte order, which is what reading
|
||||
// the four octets little endian reproduces.
|
||||
target := bpfTarget{
|
||||
Ip: binary.LittleEndian.Uint32(ip.To4()),
|
||||
Port: uint16(port),
|
||||
}
|
||||
|
||||
if err := objs.TargetMap.Put(uint32(0), target); err != nil {
|
||||
return fmt.Errorf("write redirect target: %w", err)
|
||||
}
|
||||
|
||||
fmt.Printf("Redirect target: %s\n", addr)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func configureExempt(objs *bpfObjects, list string) error {
|
||||
if list == "" {
|
||||
fmt.Println("No exempt uids set, nothing bypasses the ladder")
|
||||
return nil
|
||||
}
|
||||
|
||||
var applied []string
|
||||
for _, field := range strings.Split(list, ",") {
|
||||
field = strings.TrimSpace(field)
|
||||
if field == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
uid, err := strconv.ParseUint(field, 10, 32)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse exempt uid %q: %w", field, err)
|
||||
}
|
||||
|
||||
if err := objs.ExemptMap.Put(uint32(uid), uint8(1)); err != nil {
|
||||
return fmt.Errorf("write exempt uid %d: %w", uid, err)
|
||||
}
|
||||
|
||||
applied = append(applied, field)
|
||||
}
|
||||
|
||||
fmt.Printf("Exempt uids: %s\n", strings.Join(applied, ", "))
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user