github.com/LanceLRQ/deer-common@v0.0.9-0.20210319081233-e8222ac018a8/sandbox/forkexec/setrlimit_unix.go (about)

     1  // +build darwin linux
     2  
     3  package forkexec
     4  
     5  import (
     6  	"math"
     7  	"runtime"
     8  	"syscall"
     9  )
    10  
    11  const DarwinSafeStackSize = 65500
    12  
    13  // 定义ITimer的常量
    14  const (
    15  	ITIMER_REAL    = 0
    16  	ITIMER_VIRTUAL = 1
    17  	ITIMER_PROF    = 2
    18  )
    19  
    20  type RLimit struct {
    21  	Which  int
    22  	Enable bool
    23  	RLim   syscall.Rlimit
    24  }
    25  
    26  type ITimerVal struct {
    27  	ItInterval TimeVal
    28  	ItValue    TimeVal
    29  }
    30  
    31  type TimeVal struct {
    32  	TvSec  uint64
    33  	TvUsec uint64
    34  }
    35  
    36  type ExecRLimit struct {
    37  	TimeLimit     int // 时间限制 (ms)
    38  	RealTimeLimit int // 真实时间限制 (ms, 触发SIGALRM)
    39  	MemoryLimit   int // 内存限制 (KB)
    40  	FileSizeLimit int // 文件读写限制 (B)
    41  	StackLimit    int // 栈大小限制 (KB，0表示用内存限制的值，-1表示不限制，建议设置为2倍。Mac下有坑，不要去设置。)
    42  }
    43  
    44  type RlimitOptions struct {
    45  	Rlimits     []RLimit
    46  	ITimerValue ITimerVal
    47  }
    48  
    49  // 解析ExecRLimit结构体并获取setrlimit操作需要的信息
    50  func GetRlimitOptions(sysRlimit *ExecRLimit) *RlimitOptions {
    51  	// Make stack limit
    52  	stackLimit := uint64(sysRlimit.StackLimit)
    53  	if stackLimit <= 0 {
    54  		stackLimit = uint64(sysRlimit.MemoryLimit * 2)
    55  	}
    56  
    57  	return &RlimitOptions{
    58  		Rlimits: []RLimit{
    59  			// Set time limit: RLIMIT_CPU
    60  			{
    61  				Which:  syscall.RLIMIT_CPU,
    62  				Enable: sysRlimit.TimeLimit > 0,
    63  				RLim: syscall.Rlimit{
    64  					Cur: uint64(math.Ceil(float64(sysRlimit.TimeLimit) / 1000.0)),
    65  					Max: uint64(math.Ceil(float64(sysRlimit.TimeLimit) / 1000.0)),
    66  				},
    67  			},
    68  			// Set memory limit: RLIMIT_DATA
    69  			{
    70  				Which:  syscall.RLIMIT_DATA,
    71  				Enable: sysRlimit.MemoryLimit > 0,
    72  				RLim: syscall.Rlimit{
    73  					Cur: uint64(sysRlimit.MemoryLimit * 1024),
    74  					Max: uint64(sysRlimit.MemoryLimit * 1024 * 2),
    75  				},
    76  			},
    77  			// Set memory limit: RLIMIT_AS
    78  			{
    79  				Which:  syscall.RLIMIT_AS,
    80  				Enable: sysRlimit.MemoryLimit > 0,
    81  				RLim: syscall.Rlimit{
    82  					Cur: uint64(sysRlimit.MemoryLimit * 1024 * 2),
    83  					Max: uint64(sysRlimit.MemoryLimit*1024*2 + 1024),
    84  				},
    85  			},
    86  			// Set stack limit (坑：macos不要去搞这个!)
    87  			{
    88  				Which:  syscall.RLIMIT_STACK,
    89  				Enable: stackLimit > 0 && sysRlimit.StackLimit >= 0 && runtime.GOOS != "darwin",
    90  				RLim: syscall.Rlimit{
    91  					Cur: stackLimit * 1024,
    92  					Max: stackLimit*1024 + 1024,
    93  				},
    94  			},
    95  			// Set file size limit: RLIMIT_FSIZE
    96  			{
    97  				Which:  syscall.RLIMIT_FSIZE,
    98  				Enable: sysRlimit.FileSizeLimit > 0,
    99  				RLim: syscall.Rlimit{
   100  					Cur: uint64(sysRlimit.FileSizeLimit),
   101  					Max: uint64(sysRlimit.FileSizeLimit),
   102  				},
   103  			},
   104  		},
   105  		ITimerValue: ITimerVal{
   106  			ItInterval: TimeVal{
   107  				TvSec:  uint64(math.Floor(float64(sysRlimit.RealTimeLimit) / 1000.0)),
   108  				TvUsec: uint64(sysRlimit.RealTimeLimit % 1000 * 1000),
   109  			},
   110  			ItValue: TimeVal{
   111  				TvSec:  uint64(math.Floor(float64(sysRlimit.RealTimeLimit) / 1000.0)),
   112  				TvUsec: uint64(sysRlimit.RealTimeLimit % 1000 * 1000),
   113  			},
   114  		},
   115  	}
   116  }