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 }