github.com/criyle/go-sandbox@v0.10.3/pkg/rlimit/rlimit.go (about)

     1  // Package rlimit provides data structure for resource limits by setrlimit syscall on linux.
     2  package rlimit
     3  
     4  import (
     5  	"fmt"
     6  	"strings"
     7  	"syscall"
     8  
     9  	"github.com/criyle/go-sandbox/runner"
    10  )
    11  
    12  // RLimits defines the rlimit applied by setrlimit syscall to traced process
    13  type RLimits struct {
    14  	CPU          uint64 // in s
    15  	CPUHard      uint64 // in s
    16  	Data         uint64 // in bytes
    17  	FileSize     uint64 // in bytes
    18  	Stack        uint64 // in bytes
    19  	AddressSpace uint64 // in bytes
    20  	OpenFile     uint64 // count
    21  	DisableCore  bool   // set core to 0
    22  }
    23  
    24  // RLimit is the resource limits defined by Linux setrlimit
    25  type RLimit struct {
    26  	// Res is the resource type (e.g. syscall.RLIMIT_CPU)
    27  	Res int
    28  	// Rlim is the limit applied to that resource
    29  	Rlim syscall.Rlimit
    30  }
    31  
    32  func getRlimit(cur, max uint64) syscall.Rlimit {
    33  	return syscall.Rlimit{Cur: cur, Max: max}
    34  }
    35  
    36  // PrepareRLimit creates rlimit structures for tracee
    37  // TimeLimit in s, SizeLimit in byte
    38  func (r *RLimits) PrepareRLimit() []RLimit {
    39  	var ret []RLimit
    40  	if r.CPU > 0 {
    41  		cpuHard := r.CPUHard
    42  		if cpuHard < r.CPU {
    43  			cpuHard = r.CPU
    44  		}
    45  
    46  		ret = append(ret, RLimit{
    47  			Res:  syscall.RLIMIT_CPU,
    48  			Rlim: getRlimit(r.CPU, cpuHard),
    49  		})
    50  	}
    51  	if r.Data > 0 {
    52  		ret = append(ret, RLimit{
    53  			Res:  syscall.RLIMIT_DATA,
    54  			Rlim: getRlimit(r.Data, r.Data),
    55  		})
    56  	}
    57  	if r.FileSize > 0 {
    58  		ret = append(ret, RLimit{
    59  			Res:  syscall.RLIMIT_FSIZE,
    60  			Rlim: getRlimit(r.FileSize, r.FileSize),
    61  		})
    62  	}
    63  	if r.Stack > 0 {
    64  		ret = append(ret, RLimit{
    65  			Res:  syscall.RLIMIT_STACK,
    66  			Rlim: getRlimit(r.Stack, r.Stack),
    67  		})
    68  	}
    69  	if r.AddressSpace > 0 {
    70  		ret = append(ret, RLimit{
    71  			Res:  syscall.RLIMIT_AS,
    72  			Rlim: getRlimit(r.AddressSpace, r.AddressSpace),
    73  		})
    74  	}
    75  	if r.OpenFile > 0 {
    76  		ret = append(ret, RLimit{
    77  			Res:  syscall.RLIMIT_NOFILE,
    78  			Rlim: getRlimit(r.OpenFile, r.OpenFile),
    79  		})
    80  	}
    81  	if r.DisableCore {
    82  		ret = append(ret, RLimit{
    83  			Res:  syscall.RLIMIT_CORE,
    84  			Rlim: getRlimit(0, 0),
    85  		})
    86  	}
    87  	return ret
    88  }
    89  
    90  func (r RLimit) String() string {
    91  	t := ""
    92  	switch r.Res {
    93  	case syscall.RLIMIT_CPU:
    94  		return fmt.Sprintf("CPU[%d s:%d s]", r.Rlim.Cur, r.Rlim.Max)
    95  	case syscall.RLIMIT_NOFILE:
    96  		return fmt.Sprintf("OpenFile[%d:%d]", r.Rlim.Cur, r.Rlim.Max)
    97  	case syscall.RLIMIT_DATA:
    98  		t = "Data"
    99  	case syscall.RLIMIT_FSIZE:
   100  		t = "File"
   101  	case syscall.RLIMIT_STACK:
   102  		t = "Stack"
   103  	case syscall.RLIMIT_AS:
   104  		t = "AddressSpace"
   105  	case syscall.RLIMIT_CORE:
   106  		t = "Core"
   107  	}
   108  	return fmt.Sprintf("%s[%v:%v]", t, runner.Size(r.Rlim.Cur), runner.Size(r.Rlim.Max))
   109  }
   110  
   111  func (r RLimits) String() string {
   112  	var sb strings.Builder
   113  	sb.WriteString("RLimits[")
   114  	for i, rl := range r.PrepareRLimit() {
   115  		if i > 0 {
   116  			sb.WriteByte(',')
   117  		}
   118  		sb.WriteString(rl.String())
   119  	}
   120  	sb.WriteString("]")
   121  	return sb.String()
   122  }