1/*
2 * Copyright (c) 2013, 2016, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation.
8 *
9 * This code is distributed in the hope that it will be useful, but WITHOUT
10 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
11 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
12 * version 2 for more details (a copy is included in the LICENSE file that
13 * accompanied this code).
14 *
15 * You should have received a copy of the GNU General Public License version
16 * 2 along with this work; if not, write to the Free Software Foundation,
17 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
18 *
19 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
20 * or visit www.oracle.com if you need additional information or have any
21 * questions.
22 */
23
24package toolbox;
25
26import java.io.BufferedReader;
27import java.io.ByteArrayOutputStream;
28import java.io.File;
29import java.io.IOException;
30import java.io.InputStream;
31import java.io.InputStreamReader;
32import java.io.PrintStream;
33import java.io.PrintWriter;
34import java.io.StringWriter;
35import java.util.EnumMap;
36import java.util.HashMap;
37import java.util.Map;
38import static toolbox.ToolBox.lineSeparator;
39
40/**
41 * A utility base class to simplify the implementation of tasks.
42 * Provides support for running the task in a process and for
43 * capturing output written by the task to stdout, stderr and
44 * other writers where applicable.
45 * @param <T> the implementing subclass
46 */
47abstract class AbstractTask<T extends AbstractTask<T>> implements Task {
48    protected final ToolBox toolBox;
49    protected final Mode mode;
50    private final Map<OutputKind, String> redirects = new EnumMap<>(OutputKind.class);
51    private final Map<String, String> envVars = new HashMap<>();
52    private Expect expect = Expect.SUCCESS;
53    int expectedExitCode = 0;
54
55    /**
56     * Create a task that will execute in the specified mode.
57     * @param mode the mode
58     */
59    protected AbstractTask(ToolBox tb, Mode mode) {
60        toolBox = tb;
61        this.mode = mode;
62    }
63
64    /**
65     * Sets the expected outcome of the task and calls {@code run()}.
66     * @param expect the expected outcome
67     * @return the result of calling {@code run()}
68     */
69    public Result run(Expect expect) {
70        expect(expect, Integer.MIN_VALUE);
71        return run();
72    }
73
74    /**
75     * Sets the expected outcome of the task and calls {@code run()}.
76     * @param expect the expected outcome
77     * @param exitCode the expected exit code if the expected outcome
78     *      is {@code FAIL}
79     * @return the result of calling {@code run()}
80     */
81    public Result run(Expect expect, int exitCode) {
82        expect(expect, exitCode);
83        return run();
84    }
85
86    /**
87     * Sets the expected outcome and expected exit code of the task.
88     * The exit code will not be checked if the outcome is
89     * {@code Expect.SUCCESS} or if the exit code is set to
90     * {@code Integer.MIN_VALUE}.
91     * @param expect the expected outcome
92     * @param exitCode the expected exit code
93     */
94    protected void expect(Expect expect, int exitCode) {
95        this.expect = expect;
96        this.expectedExitCode = exitCode;
97    }
98
99    /**
100     * Checks the exit code contained in a {@code Result} against the
101     * expected outcome and exit value
102     * @param result the result object
103     * @return the result object
104     * @throws TaskError if the exit code stored in the result object
105     *      does not match the expected outcome and exit code.
106     */
107    protected Result checkExit(Result result) throws TaskError {
108        switch (expect) {
109            case SUCCESS:
110                if (result.exitCode != 0) {
111                    result.writeAll();
112                    throw new TaskError("Task " + name() + " failed: rc=" + result.exitCode);
113                }
114                break;
115
116            case FAIL:
117                if (result.exitCode == 0) {
118                    result.writeAll();
119                    throw new TaskError("Task " + name() + " succeeded unexpectedly");
120                }
121
122                if (expectedExitCode != Integer.MIN_VALUE
123                        && result.exitCode != expectedExitCode) {
124                    result.writeAll();
125                    throw new TaskError("Task " + name() + "failed with unexpected exit code "
126                        + result.exitCode + ", expected " + expectedExitCode);
127                }
128                break;
129        }
130        return result;
131    }
132
133    /**
134     * Sets an environment variable to be used by this task.
135     * @param name the name of the environment variable
136     * @param value the value for the environment variable
137     * @return this task object
138     * @throws IllegalStateException if the task mode is not {@code EXEC}
139     */
140    public T envVar(String name, String value) {
141        if (mode != Mode.EXEC)
142            throw new IllegalStateException();
143        envVars.put(name, value);
144        return (T) this;
145    }
146
147    /**
148     * Redirects output from an output stream to a file.
149     * @param outputKind the name of the stream to be redirected.
150     * @param path the file
151     * @return this task object
152     * @throws IllegalStateException if the task mode is not {@code EXEC}
153     */
154    public T redirect(OutputKind outputKind, String path) {
155        if (mode != Mode.EXEC)
156            throw new IllegalStateException();
157        redirects.put(outputKind, path);
158        return (T) this;
159    }
160
161    /**
162     * Returns a {@code ProcessBuilder} initialized with any
163     * redirects and environment variables that have been set.
164     * @return a {@code ProcessBuilder}
165     */
166    protected ProcessBuilder getProcessBuilder() {
167        if (mode != Mode.EXEC)
168            throw new IllegalStateException();
169        ProcessBuilder pb = new ProcessBuilder();
170        if (redirects.get(OutputKind.STDOUT) != null)
171            pb.redirectOutput(new File(redirects.get(OutputKind.STDOUT)));
172        if (redirects.get(OutputKind.STDERR) != null)
173            pb.redirectError(new File(redirects.get(OutputKind.STDERR)));
174        pb.environment().putAll(envVars);
175        return pb;
176    }
177
178    /**
179     * Collects the output from a process and saves it in a {@code Result}.
180     * @param tb the {@code ToolBox} containing the task {@code t}
181     * @param t the task initiating the process
182     * @param p the process
183     * @return a Result object containing the output from the process and its
184     *      exit value.
185     * @throws InterruptedException if the thread is interrupted
186     */
187    protected Result runProcess(ToolBox tb, Task t, Process p) throws InterruptedException {
188        if (mode != Mode.EXEC)
189            throw new IllegalStateException();
190        ProcessOutput sysOut = new ProcessOutput(p.getInputStream()).start();
191        ProcessOutput sysErr = new ProcessOutput(p.getErrorStream()).start();
192        sysOut.waitUntilDone();
193        sysErr.waitUntilDone();
194        int rc = p.waitFor();
195        Map<OutputKind, String> outputMap = new EnumMap<>(OutputKind.class);
196        outputMap.put(OutputKind.STDOUT, sysOut.getOutput());
197        outputMap.put(OutputKind.STDERR, sysErr.getOutput());
198        return checkExit(new Result(toolBox, t, rc, outputMap));
199    }
200
201    /**
202     * Thread-friendly class to read the output from a process until the stream
203     * is exhausted.
204     */
205    static class ProcessOutput implements Runnable {
206        ProcessOutput(InputStream from) {
207            in = new BufferedReader(new InputStreamReader(from));
208            out = new StringBuilder();
209        }
210
211        ProcessOutput start() {
212            new Thread(this).start();
213            return this;
214        }
215
216        @Override
217        public void run() {
218            try {
219                String line;
220                while ((line = in.readLine()) != null) {
221                    out.append(line).append(lineSeparator);
222                }
223            } catch (IOException e) {
224            }
225            synchronized (this) {
226                done = true;
227                notifyAll();
228            }
229        }
230
231        synchronized void waitUntilDone() throws InterruptedException {
232            boolean interrupted = false;
233
234            // poll interrupted flag, while waiting for copy to complete
235            while (!(interrupted = Thread.interrupted()) && !done)
236                wait(1000);
237
238            if (interrupted)
239                throw new InterruptedException();
240        }
241
242        String getOutput() {
243            return out.toString();
244        }
245
246        private final BufferedReader in;
247        private final StringBuilder out;
248        private boolean done;
249    }
250
251    /**
252     * Utility class to simplify the handling of temporarily setting a
253     * new stream for System.out or System.err.
254     */
255    static class StreamOutput {
256        // Functional interface to set a stream.
257        // Expected use: System::setOut, System::setErr
258        interface Initializer {
259            void set(PrintStream s);
260        }
261
262        private final ByteArrayOutputStream baos = new ByteArrayOutputStream();
263        private final PrintStream ps = new PrintStream(baos);
264        private final PrintStream prev;
265        private final Initializer init;
266
267        StreamOutput(PrintStream s, Initializer init) {
268            prev = s;
269            init.set(ps);
270            this.init = init;
271        }
272
273        /**
274         * Closes the stream and returns the contents that were written to it.
275         * @return the contents that were written to it.
276         */
277        String close() {
278            init.set(prev);
279            ps.close();
280            return baos.toString();
281        }
282    }
283
284    /**
285     * Utility class to simplify the handling of creating an in-memory PrintWriter.
286     */
287    static class WriterOutput {
288        private final StringWriter sw = new StringWriter();
289        final PrintWriter pw = new PrintWriter(sw);
290
291        /**
292         * Closes the stream and returns the contents that were written to it.
293         * @return the contents that were written to it.
294         */
295        String close() {
296            pw.close();
297            return sw.toString();
298        }
299    }
300}
301