add tests for prlimit

Signed-off-by: lifubang <[email protected]>

Author: lifubang

src/syscall/syscall_linux_test.go

@@ -17,6 +17,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"testing"
 	"unsafe"
@@ -654,3 +655,63 @@ func TestAllThreadsSyscallBlockedSyscall(t *testing.T) {
 	wr.Close()
 	wg.Wait()
 }
+
+//go:linkname syscall_prlimit syscall.prlimit
+func syscall_prlimit(pid, resource int, newlimit, old *syscall.Rlimit) error
+
+//go:linkname syscall_origRlimitNofile syscall.origRlimitNofile
+var syscall_origRlimitNofile atomic.Pointer[syscall.Rlimit]
+
+func TestPrlimitSelf(t *testing.T) {
+	syscall_origRlimitNofile.Store(&syscall.Rlimit{
+		Cur: 1024,
+		Max: 65536,
+	})
+
+	// Get the current nofile limit
+	var lim syscall.Rlimit
+	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim); err != nil {
+		t.Fatalf("Failed to get the current nofile limit: %v", err)
+	}
+	// Set current process's nofile limit through prlimit
+	if err := syscall_prlimit(0, syscall.RLIMIT_NOFILE, &lim, nil); err != nil {
+		t.Fatalf("Prlimit self failed: %v", err)
+	}
+
+	rlimLater := syscall_origRlimitNofile.Load()
+	if rlimLater != nil {
+		t.Fatalf("origRlimitNofile expected to be nil, but got %v", rlimLater)
+	}
+}
+
+func TestPrlimitOtherProcess(t *testing.T) {
+	syscall_origRlimitNofile.Store(&syscall.Rlimit{
+		Cur: 1024,
+		Max: 65536,
+	})
+	rlimOrig := syscall_origRlimitNofile.Load()
+
+	// Get the current nofile limit
+	var lim syscall.Rlimit
+	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim); err != nil {
+		t.Fatalf("Failed to get the current nofile limit: %v", err)
+	}
+
+	// Start a child process firstly,
+	// so we can use Prlimit to set it's nofile limit.
+	cmd := exec.Command("sleep", "infinity")
+	cmd.Start()
+
+	// Set child process's nofile rlimit through prlimit
+	if err := syscall_prlimit(cmd.Process.Pid, syscall.RLIMIT_NOFILE, &lim, nil); err != nil {
+		t.Fatalf("Prlimit(%d) failed: %v", cmd.Process.Pid, err)
+	}
+
+	cmd.Process.Kill()
+	cmd.Process.Wait()
+
+	rlimLater := syscall_origRlimitNofile.Load()
+	if rlimLater != rlimOrig {
+		t.Fatalf("origRlimitNofile expected to be %v, but got %v", rlimOrig, rlimLater)
+	}
+}