src.elv.sh@v0.21.0-dev.0.20240515223629-06979efb9a2a/pkg/mods/unix/rlimit.go (about)

     1  //go:build unix
     2  
     3  package unix
     4  
     5  import (
     6  	"fmt"
     7  	"sync"
     8  
     9  	"golang.org/x/sys/unix"
    10  	"src.elv.sh/pkg/eval/errs"
    11  	"src.elv.sh/pkg/eval/vals"
    12  )
    13  
    14  type rlimitsVar struct{}
    15  
    16  var (
    17  	getRlimit = unix.Getrlimit
    18  	setRlimit = unix.Setrlimit
    19  )
    20  
    21  var (
    22  	rlimitMutex sync.Mutex
    23  	rlimits     map[int]*unix.Rlimit
    24  )
    25  
    26  func (rlimitsVar) Get() any {
    27  	rlimitMutex.Lock()
    28  	defer rlimitMutex.Unlock()
    29  
    30  	initRlimits()
    31  	rlimitsMap := vals.EmptyMap
    32  	for res, lim := range rlimits {
    33  		limMap := vals.EmptyMap
    34  		if lim.Cur != unix.RLIM_INFINITY {
    35  			limMap = limMap.Assoc("cur", convertRlimT(lim.Cur))
    36  		}
    37  		if lim.Max != unix.RLIM_INFINITY {
    38  			limMap = limMap.Assoc("max", convertRlimT(lim.Max))
    39  		}
    40  		rlimitsMap = rlimitsMap.Assoc(rlimitKeys[res], limMap)
    41  	}
    42  	return rlimitsMap
    43  }
    44  
    45  func (rlimitsVar) Set(v any) error {
    46  	newRlimits, err := parseRlimitsMap(v)
    47  	if err != nil {
    48  		return err
    49  	}
    50  
    51  	rlimitMutex.Lock()
    52  	defer rlimitMutex.Unlock()
    53  
    54  	initRlimits()
    55  	for res := range rlimits {
    56  		if *rlimits[res] != *newRlimits[res] {
    57  			err := setRlimit(res, newRlimits[res])
    58  			if err != nil {
    59  				return fmt.Errorf("setrlimit %s: %w", rlimitKeys[res], err)
    60  			}
    61  			rlimits[res] = newRlimits[res]
    62  		}
    63  	}
    64  	return nil
    65  }
    66  
    67  func initRlimits() {
    68  	if rlimits != nil {
    69  		return
    70  	}
    71  	rlimits = make(map[int]*unix.Rlimit)
    72  	for res := range rlimitKeys {
    73  		var lim unix.Rlimit
    74  		err := getRlimit(res, &lim)
    75  		if err == nil {
    76  			rlimits[res] = &lim
    77  		} else {
    78  			// Since getrlimit should only ever return an error when the
    79  			// resource is not supported, this should normally never happen. But
    80  			// be defensive nonetheless.
    81  			logger.Println("initialize rlimits", res, rlimitKeys[res], err)
    82  			// Remove this key, so that rlimitKeys is always consistent with the
    83  			// value of rlimits (and thus $unix:rlimits).
    84  			delete(rlimitKeys, res)
    85  		}
    86  	}
    87  }
    88  
    89  func parseRlimitsMap(val any) (map[int]*unix.Rlimit, error) {
    90  	if err := checkRlimitsMapKeys(val); err != nil {
    91  		return nil, err
    92  	}
    93  	limits := make(map[int]*unix.Rlimit, len(rlimitKeys))
    94  	for res, key := range rlimitKeys {
    95  		limitVal, err := vals.Index(val, key)
    96  		if err != nil {
    97  			return nil, err
    98  		}
    99  		limits[res], err = parseRlimitMap(limitVal)
   100  		if err != nil {
   101  			return nil, err
   102  		}
   103  	}
   104  	return limits, nil
   105  }
   106  
   107  func checkRlimitsMapKeys(val any) error {
   108  	wantedKeys := make(map[string]struct{}, len(rlimitKeys))
   109  	for _, key := range rlimitKeys {
   110  		wantedKeys[key] = struct{}{}
   111  	}
   112  	var errKey error
   113  	err := vals.IterateKeys(val, func(k any) bool {
   114  		ks, ok := k.(string)
   115  		if !ok {
   116  			errKey = errs.BadValue{What: "key of $unix:rlimits",
   117  				Valid: "string", Actual: vals.Kind(k)}
   118  			return false
   119  		}
   120  		if _, valid := wantedKeys[ks]; !valid {
   121  			errKey = errs.BadValue{What: "key of $unix:rlimits",
   122  				Valid: "valid resource key", Actual: vals.ReprPlain(k)}
   123  			return false
   124  		}
   125  		delete(wantedKeys, ks)
   126  		return true
   127  	})
   128  	if err != nil {
   129  		return errs.BadValue{What: "$unix:rlimits",
   130  			Valid: "map", Actual: vals.Kind(val)}
   131  	}
   132  	if errKey != nil {
   133  		return errKey
   134  	}
   135  	if len(wantedKeys) > 0 {
   136  		return errs.BadValue{What: "$unix:rlimits",
   137  			Valid: "map containing all resource keys", Actual: vals.ReprPlain(val)}
   138  	}
   139  	return nil
   140  }
   141  
   142  func parseRlimitMap(val any) (*unix.Rlimit, error) {
   143  	if err := checkRlimitMapKeys(val); err != nil {
   144  		return nil, err
   145  	}
   146  	cur, err := indexRlimitMap(val, "cur")
   147  	if err != nil {
   148  		return nil, err
   149  	}
   150  	max, err := indexRlimitMap(val, "max")
   151  	if err != nil {
   152  		return nil, err
   153  	}
   154  	return &unix.Rlimit{Cur: cur, Max: max}, nil
   155  }
   156  
   157  func checkRlimitMapKeys(val any) error {
   158  	var errKey error
   159  	err := vals.IterateKeys(val, func(k any) bool {
   160  		if k != "cur" && k != "max" {
   161  			errKey = errs.BadValue{What: "key of rlimit value",
   162  				Valid: "cur or max", Actual: vals.ReprPlain(k)}
   163  			return false
   164  		}
   165  		return true
   166  	})
   167  	if err != nil {
   168  		return errs.BadValue{What: "rlimit value",
   169  			Valid: "map", Actual: vals.Kind(val)}
   170  	}
   171  	return errKey
   172  }
   173  
   174  func indexRlimitMap(m any, key string) (rlimT, error) {
   175  	val, err := vals.Index(m, key)
   176  	if err != nil {
   177  		return unix.RLIM_INFINITY, nil
   178  	}
   179  	if r, ok := parseRlimT(val); ok {
   180  		return r, nil
   181  	}
   182  	return 0, errs.BadValue{What: key + " in rlimit value",
   183  		Valid: rlimTValid, Actual: vals.ReprPlain(val)}
   184  }