expand.go

// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package config

import (
	"bytes"

	"golang.org/x/text/transform"
)

// expandTransformer implements transform.Transformer
type expandTransformer struct {
	transform.NopResetter
	expand func(string) (string, error)
}

// First char of shell variable may be [a-zA-Z_]
func isShellNameFirstChar(c byte) bool {
	return c == '_' ||
		(c >= 'a' && c <= 'z') ||
		(c >= 'A' && c <= 'Z')
}

// Char's after the first of shell variable may be [a-zA-Z0-9_]
func isShellNameChar(c byte) bool {
	return c == '_' ||
		(c >= 'a' && c <= 'z') ||
		(c >= 'A' && c <= 'Z') ||
		(c >= '0' && c <= '9')
}

// bytesIndexCFunc returns the index of the byte for which the
// complement of the supplied function is true
func bytesIndexCFunc(buf []byte, f func(b byte) bool) int {
	for i, b := range buf {
		if !f(b) {
			return i
		}
	}
	return -1
}

// Transform expands shell-like sequences like $foo and ${foo} using
// the configured expand function.  The sequence '$$' is replaced with
// a literal '$'.
func (e *expandTransformer) Transform(dst, src []byte, atEOF bool) (int, int, error) {
	var srcPos int
	var dstPos int

	for srcPos < len(src) {

		if dstPos == len(dst) {
			return dstPos, srcPos, transform.ErrShortDst
		}

		end := bytes.IndexByte(src[srcPos:], '$')

		if end == -1 {
			// src does not contain '$', copy into dst
			cnt := copy(dst[dstPos:], src[srcPos:])
			srcPos += cnt
			dstPos += cnt
			continue
		} else if end > 0 {
			// copy chars preceding '$' from src to dst
			cnt := copy(dst[dstPos:], src[srcPos:srcPos+end])
			srcPos += cnt
			dstPos += cnt

			if dstPos == len(dst) {
				return dstPos, srcPos, transform.ErrShortDst
			}
		}

		// src[srcPos] now points to '$', dstPos < len(dst)

		// If we're at the end of src, but we found a starting
		// token, return ErrShortSrc, unless we're also at EOF,
		// in which case just copy it dst.
		if srcPos+1 == len(src) {
			if atEOF {
				dst[dstPos] = src[srcPos]
				srcPos++
				dstPos++
				continue
			}
			return dstPos, srcPos, transform.ErrShortSrc
		}

		// At this point we know that src[srcPos+1] is populated.

		// If this token sequence represents the special '$$'
		// sequence, emit a '$' into dst.
		if src[srcPos+1] == '$' {
			dst[dstPos] = src[srcPos]
			srcPos += 2
			dstPos++
			continue
		}

		var token []byte
		var tokenEnd int

		// Start of bracketed token ${foo}
		if src[srcPos+1] == '{' {
			end := bytes.IndexByte(src[srcPos+2:], '}')
			if end == -1 {
				if atEOF {
					// No closing bracket and we're at
					// EOF, so it's not a valid bracket
					// expression.
					if len(dst[dstPos:]) <
						len(src[srcPos:]) {
						return dstPos, srcPos,
							transform.ErrShortDst
					}

					cnt := copy(dst[dstPos:], src[srcPos:])
					srcPos += cnt
					dstPos += cnt
					continue
				}

				// Otherwise, we need more bytes in src
				return dstPos, srcPos, transform.ErrShortSrc
			}

			// Set tokenEnd so it points to the byte
			// immediately after the closing '}'
			tokenEnd = end + srcPos + 3

			token = src[srcPos+2 : tokenEnd-1]
		} else { // Else start of non-bracketed token $foo
			if !isShellNameFirstChar(src[srcPos+1]) {
				// If it doesn't conform to the naming
				// rules for shell variables, do not
				// try to expand, just copy to dst.
				dst[dstPos] = src[srcPos]
				srcPos++
				dstPos++
				continue
			}

			end := bytesIndexCFunc(src[srcPos+2:], isShellNameChar)

			if end == -1 {
				// Reached the end of src without finding
				// end of shell variable
				if !atEOF {
					// We need more bytes in src
					return dstPos, srcPos,
						transform.ErrShortSrc
				}
				tokenEnd = len(src)
			} else {
				// Set tokenEnd so it points to the byte
				// immediately after the token
				tokenEnd = end + srcPos + 2
			}

			token = src[srcPos+1 : tokenEnd]
		}

		replacement, err := e.expand(string(token))
		if err != nil {
			return dstPos, srcPos, err
		}

		if len(dst[dstPos:]) < len(replacement) {
			return dstPos, srcPos, transform.ErrShortDst
		}

		cnt := copy(dst[dstPos:], replacement)
		srcPos = tokenEnd
		dstPos += cnt
	}

	return dstPos, srcPos, nil
}