go-tg-screenshot-bot

1
// Copyright 2009 The Go Authors. All rights reserved.
2
// Use of this source code is governed by a BSD-style
3
// license that can be found in the LICENSE file.
4

5
// Fork, exec, wait, etc.
6

7
package windows
8

9
import (
10
	errorspkg "errors"
11
	"unsafe"
12
)
13

14
// EscapeArg rewrites command line argument s as prescribed
15
// in http://msdn.microsoft.com/en-us/library/ms880421.
16
// This function returns "" (2 double quotes) if s is empty.
17
// Alternatively, these transformations are done:
18
//   - every back slash (\) is doubled, but only if immediately
19
//     followed by double quote (");
20
//   - every double quote (") is escaped by back slash (\);
21
//   - finally, s is wrapped with double quotes (arg -> "arg"),
22
//     but only if there is space or tab inside s.
23
func EscapeArg(s string) string {
24
	if len(s) == 0 {
25
		return `""`
26
	}
27
	n := len(s)
28
	hasSpace := false
29
	for i := 0; i < len(s); i++ {
30
		switch s[i] {
31
		case '"', '\\':
32
			n++
33
		case ' ', '\t':
34
			hasSpace = true
35
		}
36
	}
37
	if hasSpace {
38
		n += 2 // Reserve space for quotes.
39
	}
40
	if n == len(s) {
41
		return s
42
	}
43

44
	qs := make([]byte, n)
45
	j := 0
46
	if hasSpace {
47
		qs[j] = '"'
48
		j++
49
	}
50
	slashes := 0
51
	for i := 0; i < len(s); i++ {
52
		switch s[i] {
53
		default:
54
			slashes = 0
55
			qs[j] = s[i]
56
		case '\\':
57
			slashes++
58
			qs[j] = s[i]
59
		case '"':
60
			for ; slashes > 0; slashes-- {
61
				qs[j] = '\\'
62
				j++
63
			}
64
			qs[j] = '\\'
65
			j++
66
			qs[j] = s[i]
67
		}
68
		j++
69
	}
70
	if hasSpace {
71
		for ; slashes > 0; slashes-- {
72
			qs[j] = '\\'
73
			j++
74
		}
75
		qs[j] = '"'
76
		j++
77
	}
78
	return string(qs[:j])
79
}
80

81
// ComposeCommandLine escapes and joins the given arguments suitable for use as a Windows command line,
82
// in CreateProcess's CommandLine argument, CreateService/ChangeServiceConfig's BinaryPathName argument,
83
// or any program that uses CommandLineToArgv.
84
func ComposeCommandLine(args []string) string {
85
	if len(args) == 0 {
86
		return ""
87
	}
88

89
	// Per https://learn.microsoft.com/en-us/windows/win32/api/shellapi/nf-shellapi-commandlinetoargvw:
90
	// “This function accepts command lines that contain a program name; the
91
	// program name can be enclosed in quotation marks or not.”
92
	//
93
	// Unfortunately, it provides no means of escaping interior quotation marks
94
	// within that program name, and we have no way to report them here.
95
	prog := args[0]
96
	mustQuote := len(prog) == 0
97
	for i := 0; i < len(prog); i++ {
98
		c := prog[i]
99
		if c <= ' ' || (c == '"' && i == 0) {
100
			// Force quotes for not only the ASCII space and tab as described in the
101
			// MSDN article, but also ASCII control characters.
102
			// The documentation for CommandLineToArgvW doesn't say what happens when
103
			// the first argument is not a valid program name, but it empirically
104
			// seems to drop unquoted control characters.
105
			mustQuote = true
106
			break
107
		}
108
	}
109
	var commandLine []byte
110
	if mustQuote {
111
		commandLine = make([]byte, 0, len(prog)+2)
112
		commandLine = append(commandLine, '"')
113
		for i := 0; i < len(prog); i++ {
114
			c := prog[i]
115
			if c == '"' {
116
				// This quote would interfere with our surrounding quotes.
117
				// We have no way to report an error, so just strip out
118
				// the offending character instead.
119
				continue
120
			}
121
			commandLine = append(commandLine, c)
122
		}
123
		commandLine = append(commandLine, '"')
124
	} else {
125
		if len(args) == 1 {
126
			// args[0] is a valid command line representing itself.
127
			// No need to allocate a new slice or string for it.
128
			return prog
129
		}
130
		commandLine = []byte(prog)
131
	}
132

133
	for _, arg := range args[1:] {
134
		commandLine = append(commandLine, ' ')
135
		// TODO(bcmills): since we're already appending to a slice, it would be nice
136
		// to avoid the intermediate allocations of EscapeArg.
137
		// Perhaps we can factor out an appendEscapedArg function.
138
		commandLine = append(commandLine, EscapeArg(arg)...)
139
	}
140
	return string(commandLine)
141
}
142

143
// DecomposeCommandLine breaks apart its argument command line into unescaped parts using CommandLineToArgv,
144
// as gathered from GetCommandLine, QUERY_SERVICE_CONFIG's BinaryPathName argument, or elsewhere that
145
// command lines are passed around.
146
// DecomposeCommandLine returns an error if commandLine contains NUL.
147
func DecomposeCommandLine(commandLine string) ([]string, error) {
148
	if len(commandLine) == 0 {
149
		return []string{}, nil
150
	}
151
	utf16CommandLine, err := UTF16FromString(commandLine)
152
	if err != nil {
153
		return nil, errorspkg.New("string with NUL passed to DecomposeCommandLine")
154
	}
155
	var argc int32
156
	argv, err := commandLineToArgv(&utf16CommandLine[0], &argc)
157
	if err != nil {
158
		return nil, err
159
	}
160
	defer LocalFree(Handle(unsafe.Pointer(argv)))
161

162
	var args []string
163
	for _, p := range unsafe.Slice(argv, argc) {
164
		args = append(args, UTF16PtrToString(p))
165
	}
166
	return args, nil
167
}
168

169
// CommandLineToArgv parses a Unicode command line string and sets
170
// argc to the number of parsed arguments.
171
//
172
// The returned memory should be freed using a single call to LocalFree.
173
//
174
// Note that although the return type of CommandLineToArgv indicates 8192
175
// entries of up to 8192 characters each, the actual count of parsed arguments
176
// may exceed 8192, and the documentation for CommandLineToArgvW does not mention
177
// any bound on the lengths of the individual argument strings.
178
// (See https://go.dev/issue/63236.)
179
func CommandLineToArgv(cmd *uint16, argc *int32) (argv *[8192]*[8192]uint16, err error) {
180
	argp, err := commandLineToArgv(cmd, argc)
181
	argv = (*[8192]*[8192]uint16)(unsafe.Pointer(argp))
182
	return argv, err
183
}
184

185
func CloseOnExec(fd Handle) {
186
	SetHandleInformation(Handle(fd), HANDLE_FLAG_INHERIT, 0)
187
}
188

189
// FullPath retrieves the full path of the specified file.
190
func FullPath(name string) (path string, err error) {
191
	p, err := UTF16PtrFromString(name)
192
	if err != nil {
193
		return "", err
194
	}
195
	n := uint32(100)
196
	for {
197
		buf := make([]uint16, n)
198
		n, err = GetFullPathName(p, uint32(len(buf)), &buf[0], nil)
199
		if err != nil {
200
			return "", err
201
		}
202
		if n <= uint32(len(buf)) {
203
			return UTF16ToString(buf[:n]), nil
204
		}
205
	}
206
}
207

208
// NewProcThreadAttributeList allocates a new ProcThreadAttributeListContainer, with the requested maximum number of attributes.
209
func NewProcThreadAttributeList(maxAttrCount uint32) (*ProcThreadAttributeListContainer, error) {
210
	var size uintptr
211
	err := initializeProcThreadAttributeList(nil, maxAttrCount, 0, &size)
212
	if err != ERROR_INSUFFICIENT_BUFFER {
213
		if err == nil {
214
			return nil, errorspkg.New("unable to query buffer size from InitializeProcThreadAttributeList")
215
		}
216
		return nil, err
217
	}
218
	alloc, err := LocalAlloc(LMEM_FIXED, uint32(size))
219
	if err != nil {
220
		return nil, err
221
	}
222
	// size is guaranteed to be ≥1 by InitializeProcThreadAttributeList.
223
	al := &ProcThreadAttributeListContainer{data: (*ProcThreadAttributeList)(unsafe.Pointer(alloc))}
224
	err = initializeProcThreadAttributeList(al.data, maxAttrCount, 0, &size)
225
	if err != nil {
226
		return nil, err
227
	}
228
	return al, err
229
}
230

231
// Update modifies the ProcThreadAttributeList using UpdateProcThreadAttribute.
232
func (al *ProcThreadAttributeListContainer) Update(attribute uintptr, value unsafe.Pointer, size uintptr) error {
233
	al.pointers = append(al.pointers, value)
234
	return updateProcThreadAttribute(al.data, 0, attribute, value, size, nil, nil)
235
}
236

237
// Delete frees ProcThreadAttributeList's resources.
238
func (al *ProcThreadAttributeListContainer) Delete() {
239
	deleteProcThreadAttributeList(al.data)
240
	LocalFree(Handle(unsafe.Pointer(al.data)))
241
	al.data = nil
242
	al.pointers = nil
243
}
244

245
// List returns the actual ProcThreadAttributeList to be passed to StartupInfoEx.
246
func (al *ProcThreadAttributeListContainer) List() *ProcThreadAttributeList {
247
	return al.data
248
}
249