github.com/ncodes/nomad@v0.5.7-0.20170403112158-97adf4a74fb3/client/driver/env/env.go (about)

     1  package env
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"os"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/ncodes/nomad/helper"
    11  	hargs "github.com/ncodes/nomad/helper/args"
    12  	"github.com/ncodes/nomad/nomad/structs"
    13  )
    14  
    15  // A set of environment variables that are exported by each driver.
    16  const (
    17  	// AllocDir is the environment variable with the path to the alloc directory
    18  	// that is shared across tasks within a task group.
    19  	AllocDir = "NOMAD_ALLOC_DIR"
    20  
    21  	// TaskLocalDir is the environment variable with the path to the tasks local
    22  	// directory where it can store data that is persisted to the alloc is
    23  	// removed.
    24  	TaskLocalDir = "NOMAD_TASK_DIR"
    25  
    26  	// SecretsDir is the environment variable with the path to the tasks secret
    27  	// directory where it can store sensitive data.
    28  	SecretsDir = "NOMAD_SECRETS_DIR"
    29  
    30  	// MemLimit is the environment variable with the tasks memory limit in MBs.
    31  	MemLimit = "NOMAD_MEMORY_LIMIT"
    32  
    33  	// CpuLimit is the environment variable with the tasks CPU limit in MHz.
    34  	CpuLimit = "NOMAD_CPU_LIMIT"
    35  
    36  	// AllocID is the environment variable for passing the allocation ID.
    37  	AllocID = "NOMAD_ALLOC_ID"
    38  
    39  	// AllocName is the environment variable for passing the allocation name.
    40  	AllocName = "NOMAD_ALLOC_NAME"
    41  
    42  	// TaskName is the environment variable for passing the task name.
    43  	TaskName = "NOMAD_TASK_NAME"
    44  
    45  	// JobName is the environment variable for passing the job name.
    46  	JobName = "NOMAD_JOB_NAME"
    47  
    48  	// AllocIndex is the environment variable for passing the allocation index.
    49  	AllocIndex = "NOMAD_ALLOC_INDEX"
    50  
    51  	// AddrPrefix is the prefix for passing both dynamic and static port
    52  	// allocations to tasks.
    53  	// E.g $NOMAD_ADDR_http=127.0.0.1:80
    54  	AddrPrefix = "NOMAD_ADDR_"
    55  
    56  	// IpPrefix is the prefix for passing the IP of a port allocation to a task.
    57  	IpPrefix = "NOMAD_IP_"
    58  
    59  	// PortPrefix is the prefix for passing the port allocation to a task.
    60  	PortPrefix = "NOMAD_PORT_"
    61  
    62  	// HostPortPrefix is the prefix for passing the host port when a portmap is
    63  	// specified.
    64  	HostPortPrefix = "NOMAD_HOST_PORT_"
    65  
    66  	// MetaPrefix is the prefix for passing task meta data.
    67  	MetaPrefix = "NOMAD_META_"
    68  
    69  	// VaultToken is the environment variable for passing the Vault token
    70  	VaultToken = "VAULT_TOKEN"
    71  )
    72  
    73  // The node values that can be interpreted.
    74  const (
    75  	nodeIdKey    = "node.unique.id"
    76  	nodeDcKey    = "node.datacenter"
    77  	nodeNameKey  = "node.unique.name"
    78  	nodeClassKey = "node.class"
    79  
    80  	// Prefixes used for lookups.
    81  	nodeAttributePrefix = "attr."
    82  	nodeMetaPrefix      = "meta."
    83  )
    84  
    85  // TaskEnvironment is used to expose information to a task via environment
    86  // variables and provide interpolation of Nomad variables.
    87  type TaskEnvironment struct {
    88  	Env              map[string]string
    89  	TaskMeta         map[string]string
    90  	AllocDir         string
    91  	TaskDir          string
    92  	SecretsDir       string
    93  	CpuLimit         int
    94  	MemLimit         int
    95  	TaskName         string
    96  	AllocIndex       int
    97  	AllocId          string
    98  	AllocName        string
    99  	Node             *structs.Node
   100  	Networks         []*structs.NetworkResource
   101  	PortMap          map[string]int
   102  	VaultToken       string
   103  	InjectVaultToken bool
   104  	JobName          string
   105  	Alloc            *structs.Allocation
   106  
   107  	// taskEnv is the variables that will be set in the tasks environment
   108  	TaskEnv map[string]string
   109  
   110  	// nodeValues is the values that are allowed for interprolation from the
   111  	// node.
   112  	NodeValues map[string]string
   113  }
   114  
   115  func NewTaskEnvironment(node *structs.Node) *TaskEnvironment {
   116  	return &TaskEnvironment{Node: node, AllocIndex: -1}
   117  }
   118  
   119  // ParseAndReplace takes the user supplied args replaces any instance of an
   120  // environment variable or nomad variable in the args with the actual value.
   121  func (t *TaskEnvironment) ParseAndReplace(args []string) []string {
   122  	replaced := make([]string, len(args))
   123  	for i, arg := range args {
   124  		replaced[i] = hargs.ReplaceEnv(arg, t.TaskEnv, t.NodeValues)
   125  	}
   126  
   127  	return replaced
   128  }
   129  
   130  // ReplaceEnv takes an arg and replaces all occurrences of environment variables
   131  // and nomad variables.  If the variable is found in the passed map it is
   132  // replaced, otherwise the original string is returned.
   133  func (t *TaskEnvironment) ReplaceEnv(arg string) string {
   134  	return hargs.ReplaceEnv(arg, t.TaskEnv, t.NodeValues)
   135  }
   136  
   137  // Build must be called after all the tasks environment values have been set.
   138  func (t *TaskEnvironment) Build() *TaskEnvironment {
   139  	t.NodeValues = make(map[string]string)
   140  	t.TaskEnv = make(map[string]string)
   141  
   142  	// Build the meta
   143  	for k, v := range t.TaskMeta {
   144  		t.TaskEnv[fmt.Sprintf("%s%s", MetaPrefix, strings.ToUpper(k))] = v
   145  		t.TaskEnv[fmt.Sprintf("%s%s", MetaPrefix, k)] = v
   146  	}
   147  
   148  	// Build the ports
   149  	for _, network := range t.Networks {
   150  		for label, value := range network.MapLabelToValues(nil) {
   151  			t.TaskEnv[fmt.Sprintf("%s%s", IpPrefix, label)] = network.IP
   152  			t.TaskEnv[fmt.Sprintf("%s%s", HostPortPrefix, label)] = strconv.Itoa(value)
   153  			if forwardedPort, ok := t.PortMap[label]; ok {
   154  				value = forwardedPort
   155  			}
   156  			t.TaskEnv[fmt.Sprintf("%s%s", PortPrefix, label)] = strconv.Itoa(value)
   157  			IPPort := net.JoinHostPort(network.IP, strconv.Itoa(value))
   158  			t.TaskEnv[fmt.Sprintf("%s%s", AddrPrefix, label)] = IPPort
   159  
   160  		}
   161  	}
   162  
   163  	// Build the directories
   164  	if t.AllocDir != "" {
   165  		t.TaskEnv[AllocDir] = t.AllocDir
   166  	}
   167  	if t.TaskDir != "" {
   168  		t.TaskEnv[TaskLocalDir] = t.TaskDir
   169  	}
   170  	if t.SecretsDir != "" {
   171  		t.TaskEnv[SecretsDir] = t.SecretsDir
   172  	}
   173  
   174  	// Build the resource limits
   175  	if t.MemLimit != 0 {
   176  		t.TaskEnv[MemLimit] = strconv.Itoa(t.MemLimit)
   177  	}
   178  	if t.CpuLimit != 0 {
   179  		t.TaskEnv[CpuLimit] = strconv.Itoa(t.CpuLimit)
   180  	}
   181  
   182  	// Build the tasks ids
   183  	if t.AllocId != "" {
   184  		t.TaskEnv[AllocID] = t.AllocId
   185  	}
   186  	if t.AllocName != "" {
   187  		t.TaskEnv[AllocName] = t.AllocName
   188  	}
   189  	if t.AllocIndex != -1 {
   190  		t.TaskEnv[AllocIndex] = strconv.Itoa(t.AllocIndex)
   191  	}
   192  	if t.TaskName != "" {
   193  		t.TaskEnv[TaskName] = t.TaskName
   194  	}
   195  	if t.JobName != "" {
   196  		t.TaskEnv[JobName] = t.JobName
   197  	}
   198  
   199  	// Build the addr of the other tasks
   200  	if t.Alloc != nil {
   201  		for taskName, resources := range t.Alloc.TaskResources {
   202  			if taskName == t.TaskName {
   203  				continue
   204  			}
   205  			for _, nw := range resources.Networks {
   206  				ports := make([]structs.Port, 0, len(nw.ReservedPorts)+len(nw.DynamicPorts))
   207  				for _, port := range nw.ReservedPorts {
   208  					ports = append(ports, port)
   209  				}
   210  				for _, port := range nw.DynamicPorts {
   211  					ports = append(ports, port)
   212  				}
   213  				for _, p := range ports {
   214  					key := fmt.Sprintf("%s%s_%s", AddrPrefix, taskName, p.Label)
   215  					t.TaskEnv[key] = fmt.Sprintf("%s:%d", nw.IP, p.Value)
   216  					key = fmt.Sprintf("%s%s_%s", IpPrefix, taskName, p.Label)
   217  					t.TaskEnv[key] = nw.IP
   218  					key = fmt.Sprintf("%s%s_%s", PortPrefix, taskName, p.Label)
   219  					t.TaskEnv[key] = strconv.Itoa(p.Value)
   220  				}
   221  			}
   222  		}
   223  	}
   224  
   225  	// Build the node
   226  	if t.Node != nil {
   227  		// Set up the node values.
   228  		t.NodeValues[nodeIdKey] = t.Node.ID
   229  		t.NodeValues[nodeDcKey] = t.Node.Datacenter
   230  		t.NodeValues[nodeNameKey] = t.Node.Name
   231  		t.NodeValues[nodeClassKey] = t.Node.NodeClass
   232  
   233  		// Set up the attributes.
   234  		for k, v := range t.Node.Attributes {
   235  			t.NodeValues[fmt.Sprintf("%s%s", nodeAttributePrefix, k)] = v
   236  		}
   237  
   238  		// Set up the meta.
   239  		for k, v := range t.Node.Meta {
   240  			t.NodeValues[fmt.Sprintf("%s%s", nodeMetaPrefix, k)] = v
   241  		}
   242  	}
   243  
   244  	// Build the Vault Token
   245  	if t.InjectVaultToken && t.VaultToken != "" {
   246  		t.TaskEnv[VaultToken] = t.VaultToken
   247  	}
   248  
   249  	// Interpret the environment variables
   250  	interpreted := make(map[string]string, len(t.Env))
   251  	for k, v := range t.Env {
   252  		interpreted[k] = hargs.ReplaceEnv(v, t.NodeValues, t.TaskEnv)
   253  	}
   254  
   255  	for k, v := range interpreted {
   256  		t.TaskEnv[k] = v
   257  	}
   258  
   259  	// Clean keys (see #2405)
   260  	cleanedEnv := make(map[string]string, len(t.TaskEnv))
   261  	for k, v := range t.TaskEnv {
   262  		cleanedK := helper.CleanEnvVar(k, '_')
   263  		cleanedEnv[cleanedK] = v
   264  	}
   265  	t.TaskEnv = cleanedEnv
   266  
   267  	return t
   268  }
   269  
   270  // EnvList returns a list of strings with NAME=value pairs.
   271  func (t *TaskEnvironment) EnvList() []string {
   272  	env := []string{}
   273  	for k, v := range t.TaskEnv {
   274  		env = append(env, fmt.Sprintf("%s=%s", k, v))
   275  	}
   276  
   277  	return env
   278  }
   279  
   280  // EnvMap returns a copy of the tasks environment variables.
   281  func (t *TaskEnvironment) EnvMap() map[string]string {
   282  	m := make(map[string]string, len(t.TaskEnv))
   283  	for k, v := range t.TaskEnv {
   284  		m[k] = v
   285  	}
   286  
   287  	return m
   288  }
   289  
   290  // EnvMapAll returns the environment variables that will be set as well as node
   291  // meta/attrs in the map. This is appropriate for interpolation.
   292  func (t *TaskEnvironment) EnvMapAll() map[string]string {
   293  	m := make(map[string]string, len(t.TaskEnv))
   294  	for k, v := range t.TaskEnv {
   295  		m[k] = v
   296  	}
   297  	for k, v := range t.NodeValues {
   298  		m[k] = v
   299  	}
   300  
   301  	return m
   302  }
   303  
   304  // Builder methods to build the TaskEnvironment
   305  func (t *TaskEnvironment) SetAllocDir(dir string) *TaskEnvironment {
   306  	t.AllocDir = dir
   307  	return t
   308  }
   309  
   310  func (t *TaskEnvironment) ClearAllocDir() *TaskEnvironment {
   311  	t.AllocDir = ""
   312  	return t
   313  }
   314  
   315  func (t *TaskEnvironment) SetTaskLocalDir(dir string) *TaskEnvironment {
   316  	t.TaskDir = dir
   317  	return t
   318  }
   319  
   320  func (t *TaskEnvironment) ClearTaskLocalDir() *TaskEnvironment {
   321  	t.TaskDir = ""
   322  	return t
   323  }
   324  
   325  func (t *TaskEnvironment) SetSecretsDir(dir string) *TaskEnvironment {
   326  	t.SecretsDir = dir
   327  	return t
   328  }
   329  
   330  func (t *TaskEnvironment) ClearSecretsDir() *TaskEnvironment {
   331  	t.SecretsDir = ""
   332  	return t
   333  }
   334  
   335  func (t *TaskEnvironment) SetMemLimit(limit int) *TaskEnvironment {
   336  	t.MemLimit = limit
   337  	return t
   338  }
   339  
   340  func (t *TaskEnvironment) ClearMemLimit() *TaskEnvironment {
   341  	t.MemLimit = 0
   342  	return t
   343  }
   344  
   345  func (t *TaskEnvironment) SetCpuLimit(limit int) *TaskEnvironment {
   346  	t.CpuLimit = limit
   347  	return t
   348  }
   349  
   350  func (t *TaskEnvironment) ClearCpuLimit() *TaskEnvironment {
   351  	t.CpuLimit = 0
   352  	return t
   353  }
   354  
   355  func (t *TaskEnvironment) SetNetworks(networks []*structs.NetworkResource) *TaskEnvironment {
   356  	t.Networks = networks
   357  	return t
   358  }
   359  
   360  func (t *TaskEnvironment) clearNetworks() *TaskEnvironment {
   361  	t.Networks = nil
   362  	return t
   363  }
   364  
   365  func (t *TaskEnvironment) SetPortMap(portMap map[string]int) *TaskEnvironment {
   366  	t.PortMap = portMap
   367  	return t
   368  }
   369  
   370  func (t *TaskEnvironment) clearPortMap() *TaskEnvironment {
   371  	t.PortMap = nil
   372  	return t
   373  }
   374  
   375  // Takes a map of meta values to be passed to the task. The keys are capatilized
   376  // when the environent variable is set.
   377  func (t *TaskEnvironment) SetTaskMeta(m map[string]string) *TaskEnvironment {
   378  	t.TaskMeta = m
   379  	return t
   380  }
   381  
   382  func (t *TaskEnvironment) ClearTaskMeta() *TaskEnvironment {
   383  	t.TaskMeta = nil
   384  	return t
   385  }
   386  
   387  func (t *TaskEnvironment) SetEnvvars(m map[string]string) *TaskEnvironment {
   388  	t.Env = m
   389  	return t
   390  }
   391  
   392  // Appends the given environment variables.
   393  func (t *TaskEnvironment) AppendEnvvars(m map[string]string) *TaskEnvironment {
   394  	if t.Env == nil {
   395  		t.Env = make(map[string]string, len(m))
   396  	}
   397  
   398  	for k, v := range m {
   399  		t.Env[k] = v
   400  	}
   401  	return t
   402  }
   403  
   404  // AppendHostEnvvars adds the host environment variables to the tasks. The
   405  // filter parameter can be use to filter host environment from entering the
   406  // tasks.
   407  func (t *TaskEnvironment) AppendHostEnvvars(filter []string) *TaskEnvironment {
   408  	hostEnv := os.Environ()
   409  	if t.Env == nil {
   410  		t.Env = make(map[string]string, len(hostEnv))
   411  	}
   412  
   413  	// Index the filtered environment variables.
   414  	index := make(map[string]struct{}, len(filter))
   415  	for _, f := range filter {
   416  		index[f] = struct{}{}
   417  	}
   418  
   419  	for _, e := range hostEnv {
   420  		parts := strings.SplitN(e, "=", 2)
   421  		key, value := parts[0], parts[1]
   422  
   423  		// Skip filtered environment variables
   424  		if _, filtered := index[key]; filtered {
   425  			continue
   426  		}
   427  
   428  		// Don't override the tasks environment variables.
   429  		if _, existing := t.Env[key]; !existing {
   430  			t.Env[key] = value
   431  		}
   432  	}
   433  
   434  	return t
   435  }
   436  
   437  func (t *TaskEnvironment) ClearEnvvars() *TaskEnvironment {
   438  	t.Env = nil
   439  	return t
   440  }
   441  
   442  // Helper method for setting all fields from an allocation.
   443  func (t *TaskEnvironment) SetAlloc(alloc *structs.Allocation) *TaskEnvironment {
   444  	t.AllocId = alloc.ID
   445  	t.AllocName = alloc.Name
   446  	t.AllocIndex = alloc.Index()
   447  	t.Alloc = alloc
   448  	return t
   449  }
   450  
   451  // Helper method for clearing all fields from an allocation.
   452  func (t *TaskEnvironment) ClearAlloc(alloc *structs.Allocation) *TaskEnvironment {
   453  	return t.ClearAllocId().ClearAllocName().ClearAllocIndex()
   454  }
   455  
   456  func (t *TaskEnvironment) SetAllocIndex(index int) *TaskEnvironment {
   457  	t.AllocIndex = index
   458  	return t
   459  }
   460  
   461  func (t *TaskEnvironment) ClearAllocIndex() *TaskEnvironment {
   462  	t.AllocIndex = -1
   463  	return t
   464  }
   465  
   466  func (t *TaskEnvironment) SetAllocId(id string) *TaskEnvironment {
   467  	t.AllocId = id
   468  	return t
   469  }
   470  
   471  func (t *TaskEnvironment) ClearAllocId() *TaskEnvironment {
   472  	t.AllocId = ""
   473  	return t
   474  }
   475  
   476  func (t *TaskEnvironment) SetAllocName(name string) *TaskEnvironment {
   477  	t.AllocName = name
   478  	return t
   479  }
   480  
   481  func (t *TaskEnvironment) ClearAllocName() *TaskEnvironment {
   482  	t.AllocName = ""
   483  	return t
   484  }
   485  
   486  func (t *TaskEnvironment) SetTaskName(name string) *TaskEnvironment {
   487  	t.TaskName = name
   488  	return t
   489  }
   490  
   491  func (t *TaskEnvironment) SetJobName(name string) *TaskEnvironment {
   492  	t.JobName = name
   493  	return t
   494  }
   495  
   496  func (t *TaskEnvironment) ClearTaskName() *TaskEnvironment {
   497  	t.TaskName = ""
   498  	return t
   499  }
   500  
   501  func (t *TaskEnvironment) ClearJobName() *TaskEnvironment {
   502  	t.JobName = ""
   503  	return t
   504  }
   505  
   506  func (t *TaskEnvironment) SetVaultToken(token string, inject bool) *TaskEnvironment {
   507  	t.VaultToken = token
   508  	t.InjectVaultToken = inject
   509  	return t
   510  }
   511  
   512  func (t *TaskEnvironment) ClearVaultToken() *TaskEnvironment {
   513  	t.VaultToken = ""
   514  	t.InjectVaultToken = false
   515  	return t
   516  }