reorganized monkeyleash module

monkeyleash/monkey/config.go

@@ -0,0 +1,93 @@
+package monkey
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"os"
+	"strings"
+
+	"bou.ke/monkey"
+)
+
+type ConfigCaps struct {
+	Allowlists map[string][]string `json:"allowlists"`
+	Hashes     map[string][]string `json:"hashes"`
+}
+
+type Config struct {
+	configCaps     ConfigCaps
+	capFuncMapping map[string]string
+	isHashMode     bool
+}
+
+func LoadConfig(capsFile, interestingFile string) (*Config, error) {
+
+	config := &Config{}
+
+	config.configCaps = loadConfigCaps(capsFile)
+
+	var err error
+	config.capFuncMapping, err = loadCapFuncMapping(interestingFile)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load CapFuncMapping: %v", err)
+	}
+
+	config.isHashMode = os.Getenv("GOCAP_HASH_MODE") == "true"
+
+	return config, nil
+}
+
+func loadConfigCaps(capsFile string) ConfigCaps {
+	var configCaps ConfigCaps
+
+	file, err := os.Open(capsFile)
+	if err != nil {
+		fmt.Printf("Error opening caps file: %v\n", err)
+		return configCaps
+	}
+	defer file.Close()
+
+	decoder := json.NewDecoder(file)
+	err = decoder.Decode(&configCaps)
+	if err != nil {
+		fmt.Printf("Error decoding JSON: %v\n", err)
+	}
+
+	return configCaps
+}
+
+func loadCapFuncMapping(filename string) (map[string]string, error) {
+	file, err := os.Open(filename)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	funcCapMap := make(map[string]string)
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if strings.HasPrefix(line, "func ") {
+			parts := strings.Fields(line)
+			if len(parts) == 3 {
+				funcCapMap[parts[1]] = parts[2]
+			}
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return nil, err
+	}
+
+	return funcCapMap, nil
+}
+
+func writeConfigFile(filename string, configCaps ConfigCaps) {
+	jsonData, err := json.MarshalIndent(configCaps, "", "    ")
+	if err == nil {
+		monkey.Unpatch(os.WriteFile)
+		os.WriteFile(filename, jsonData, 0644)
+		monkey.Patch(os.WriteFile, WriteFileReplaceHash)
+	}
+}

monkeyleash/monkey/guard.go

@@ -0,0 +1,139 @@
+package monkey
+
+import (
+	"encoding/hex"
+	"fmt"
+	"net"
+	"os"
+	"runtime"
+	"strings"
+
+	"bou.ke/monkey"
+	"github.com/cespare/xxhash"
+)
+
+var (
+	config *Config
+)
+
+const initialBufferSize = 1024
+const maxStackDepth = 10
+
+func init() {
+	var err error
+
+	config, err = LoadConfig("./go.cap", "../../goleash/interesting.cm")
+	if err != nil {
+		panic(fmt.Errorf("failed to load configuration: %v", err))
+	}
+
+	RegisterHooks()
+}
+
+func RegisterHooks() {
+	if config.isHashMode {
+		monkey.Patch(os.WriteFile, WriteFileReplaceHash)
+		monkey.Patch(os.ReadFile, ReadFileReplaceHash)
+		monkey.Patch(os.Create, CreateReplaceHash)
+		monkey.Patch((*os.File).Chmod, FileChmodReplaceHash)
+		monkey.Patch(os.Chmod, ChmodReplaceHash)
+		monkey.Patch(net.LookupHost, LookupHostReplaceHash)
+
+	} else {
+		monkey.Patch(os.WriteFile, WriteFileReplace)
+		monkey.Patch(os.ReadFile, ReadFileReplace)
+		monkey.Patch(os.Create, CreateReplace)
+		monkey.Patch((*os.File).Chmod, FileChmodReplace)
+		monkey.Patch(os.Chmod, ChmodReplace)
+		monkey.Patch(net.LookupHost, LookupHostReplace)
+	}
+}
+
+func getStackTrace(funcName string, skip int) string {
+
+	var sb strings.Builder
+	sb.Grow(initialBufferSize)
+	sb.WriteString(funcName)
+	sb.WriteString("|")
+
+	pcs := make([]uintptr, maxStackDepth)
+	n := runtime.Callers(skip, pcs)
+	frames := runtime.CallersFrames(pcs[:n])
+
+	for frame, more := frames.Next(); more; frame, more = frames.Next() {
+		sb.WriteString(frame.Function)
+		sb.WriteString("|")
+	}
+
+	return sb.String()
+}
+
+func generateHash(stackTrace string) string {
+	hash := xxhash.Sum64String(stackTrace)
+	// hash := sha256.Sum256([]byte(stackTrace))
+
+	return hex.EncodeToString([]byte{
+		byte(hash >> 56),
+		byte(hash >> 48),
+		byte(hash >> 40),
+		byte(hash >> 32),
+		byte(hash >> 24),
+		byte(hash >> 16),
+		byte(hash >> 8),
+		byte(hash),
+	})
+	// return hex.EncodeToString(hash[:])
+}
+
+func UpdateStackHashes(funcName string) {
+	stackTrace := getStackTrace(funcName, 4)
+	stackHash := generateHash(stackTrace)
+	invokedCapability := config.capFuncMapping[funcName]
+
+	hashes := config.configCaps.Hashes[invokedCapability]
+	if !contains(hashes, stackHash) {
+		config.configCaps.Hashes[invokedCapability] = append(hashes, stackHash)
+		writeConfigFile("./go.cap", config.configCaps)
+	}
+}
+
+func getCallerDependency(stackTrace string) string {
+	callerDepName := ""
+	for depName := range config.configCaps.Allowlists {
+		if strings.Contains(stackTrace, depName) {
+			callerDepName = depName
+			break
+		}
+	}
+
+	return callerDepName
+}
+
+func CheckCapability(funcName string) error {
+	stackTrace := getStackTrace(funcName, 4)
+	stackHash := generateHash(stackTrace)
+	callerDepName := getCallerDependency(stackTrace)
+
+	invokedCapability, ok := config.capFuncMapping[funcName]
+	if !ok {
+		return nil
+	}
+
+	allowedHashes := config.configCaps.Hashes[invokedCapability]
+	allowedCapabilities := config.configCaps.Allowlists[callerDepName]
+
+	if len(allowedHashes) == 0 || !contains(allowedHashes, stackHash) || !contains(allowedCapabilities, invokedCapability) {
+		return fmt.Errorf("capability '%s' denied for caller '%s'", invokedCapability, callerDepName)
+	}
+
+	return nil
+}
+
+func contains(slice []string, item string) bool {
+	set := make(map[string]struct{}, len(slice))
+	for _, s := range slice {
+		set[s] = struct{}{}
+	}
+	_, exists := set[item]
+	return exists
+}

monkeyleash/monkey/hooks.go

@@ -0,0 +1,81 @@
+package monkey
+
+import (
+	"fmt"
+	"net"
+	"os"
+
+	"bou.ke/monkey"
+)
+
+//go:noinline
+func WriteFileReplace(name string, data []byte, perm os.FileMode) error {
+	fmt.Println("os.WriteFile hooked.")
+	if err := CheckCapability("os.WriteFile"); err != nil {
+		panic(fmt.Sprintf("\n%v", err))
+	}
+	monkey.Unpatch(os.WriteFile)
+	ret1 := os.WriteFile(name, data, perm)
+	monkey.Patch(os.WriteFile, WriteFileReplace)
+	return ret1
+}
+
+//go:noinline
+func ReadFileReplace(name string) ([]byte, error) {
+	fmt.Println("os.ReadFile hooked.")
+	if err := CheckCapability("os.ReadFile"); err != nil {
+		panic(fmt.Sprintf("\n%v", err))
+	}
+	monkey.Unpatch(os.ReadFile)
+	ret1, ret2 := os.ReadFile(name)
+	monkey.Patch(os.ReadFile, ReadFileReplace)
+	return ret1, ret2
+}
+
+//go:noinline
+func CreateReplace(name string) (*os.File, error) {
+	fmt.Println("os.Create hooked.")
+	if err := CheckCapability("os.Create"); err != nil {
+		panic(fmt.Sprintf("\n%v", err))
+	}
+	monkey.Unpatch(os.Create)
+	ret1, ret2 := os.Create(name)
+	monkey.Patch(os.Create, CreateReplace)
+	return ret1, ret2
+}
+
+//go:noinline
+func ChmodReplace(name string, mode os.FileMode) error {
+	fmt.Println("os.Chmod hooked.")
+	if err := CheckCapability("os.Chmod"); err != nil {
+		panic(fmt.Sprintf("\n%v", err))
+	}
+	monkey.Unpatch(os.Chmod)
+	ret := os.Chmod(name, mode)
+	monkey.Patch(os.Chmod, ChmodReplace)
+	return ret
+}
+
+//go:noinline
+func FileChmodReplace(f *os.File, mode os.FileMode) error {
+	fmt.Println("(*os.File).Chmod hooked.")
+	if err := CheckCapability("(*os.File).Chmod"); err != nil {
+		panic(fmt.Sprintf("\n%v", err))
+	}
+	monkey.Unpatch((*os.File).Chmod)
+	ret := f.Chmod(mode)
+	monkey.Patch((*os.File).Chmod, FileChmodReplace)
+	return ret
+}
+
+//go:noinline
+func LookupHostReplace(host string) ([]string, error) {
+	fmt.Println("net.LookupHost hooked.")
+	if err := CheckCapability("net.LookupHost"); err != nil {
+		panic(fmt.Sprintf("\n%v", err))
+	}
+	monkey.Unpatch(net.LookupHost)
+	ret1, ret2 := net.LookupHost(host)
+	monkey.Patch(net.LookupHost, LookupHostReplace)
+	return ret1, ret2
+}

monkeyleash/monkey/hooks_hash.go

@@ -0,0 +1,69 @@
+package monkey
+
+import (
+	"fmt"
+	"net"
+	"os"
+
+	"bou.ke/monkey"
+)
+
+//go:noinline
+func WriteFileReplaceHash(name string, data []byte, perm os.FileMode) error {
+	fmt.Println("os.WriteFile hooked for hash.")
+	UpdateStackHashes("os.WriteFile")
+	monkey.Unpatch(os.WriteFile)
+	ret1 := os.WriteFile(name, data, perm)
+	monkey.Patch(os.WriteFile, WriteFileReplaceHash)
+	return ret1
+}
+
+//go:noinline
+func ReadFileReplaceHash(name string) ([]byte, error) {
+	fmt.Println("os.ReadFile hooked for hash.")
+	UpdateStackHashes("os.ReadFile")
+	monkey.Unpatch(os.ReadFile)
+	ret1, ret2 := os.ReadFile(name)
+	monkey.Patch(os.ReadFile, ReadFileReplaceHash)
+	return ret1, ret2
+}
+
+//go:noinline
+func CreateReplaceHash(name string) (*os.File, error) {
+	fmt.Println("os.Create hooked for hash.")
+	UpdateStackHashes("os.Create")
+	monkey.Unpatch(os.Create)
+	ret1, ret2 := os.Create(name)
+	monkey.Patch(os.Create, CreateReplaceHash)
+	return ret1, ret2
+}
+
+//go:noinline
+func ChmodReplaceHash(name string, mode os.FileMode) error {
+	fmt.Println("os.Chmod hooked for hash.")
+	UpdateStackHashes("os.Chmod")
+	monkey.Unpatch(os.Chmod)
+	ret := os.Chmod(name, mode)
+	monkey.Patch(os.Chmod, ChmodReplaceHash)
+	return ret
+}
+
+//go:noinline
+func FileChmodReplaceHash(f *os.File, mode os.FileMode) error {
+	fmt.Println("(*os.File).Chmod hooked for hash.")
+	UpdateStackHashes("(*os.File).Chmod")
+	monkey.Unpatch((*os.File).Chmod)
+	ret := f.Chmod(mode)
+	monkey.Patch((*os.File).Chmod, FileChmodReplaceHash)
+	return ret
+}
+
+//go:noinline
+func LookupHostReplaceHash(host string) ([]string, error) {
+	fmt.Println("net.LookupHost hooked for hash.")
+	UpdateStackHashes("net.LookupHost")
+	monkey.Unpatch(net.LookupHost)
+	ret1, ret2 := net.LookupHost(host)
+	monkey.Patch(net.LookupHost, LookupHostReplaceHash)
+	return ret1, ret2
+}