001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.commons.io.output;
018
019import java.io.File;
020import java.io.FileNotFoundException;
021import java.io.FileOutputStream;
022import java.io.IOException;
023import java.io.OutputStream;
024import java.io.OutputStreamWriter;
025import java.io.StringWriter;
026import java.io.Writer;
027import java.util.Locale;
028import java.util.regex.Matcher;
029import java.util.regex.Pattern;
030
031import org.apache.commons.io.IOUtils;
032import org.apache.commons.io.input.XmlStreamReader;
033
034/**
035 * Character stream that handles all the necessary Voodoo to figure out the
036 * charset encoding of the XML document written to the stream.
037 *
038 * @see XmlStreamReader
039 * @since 2.0
040 */
041public class XmlStreamWriter extends Writer {
042    private static final int BUFFER_SIZE = IOUtils.DEFAULT_BUFFER_SIZE;
043
044    private final OutputStream out;
045
046    private final String defaultEncoding;
047
048    private StringWriter xmlPrologWriter = new StringWriter(BUFFER_SIZE);
049
050    private Writer writer;
051
052    private String encoding;
053
054    /**
055     * Constructs a new XML stream writer for the specified output stream
056     * with a default encoding of UTF-8.
057     *
058     * @param out The output stream
059     */
060    public XmlStreamWriter(final OutputStream out) {
061        this(out, null);
062    }
063
064    /**
065     * Constructs a new XML stream writer for the specified output stream
066     * with the specified default encoding.
067     *
068     * @param out The output stream
069     * @param defaultEncoding The default encoding if not encoding could be detected
070     */
071    public XmlStreamWriter(final OutputStream out, final String defaultEncoding) {
072        this.out = out;
073        this.defaultEncoding = defaultEncoding != null ? defaultEncoding : "UTF-8";
074    }
075
076    /**
077     * Constructs a new XML stream writer for the specified file
078     * with a default encoding of UTF-8.
079     *
080     * @param file The file to write to
081     * @throws FileNotFoundException if there is an error creating or
082     * opening the file
083     */
084    public XmlStreamWriter(final File file) throws FileNotFoundException {
085        this(file, null);
086    }
087
088    /**
089     * Constructs a new XML stream writer for the specified file
090     * with the specified default encoding.
091     *
092     * @param file The file to write to
093     * @param defaultEncoding The default encoding if not encoding could be detected
094     * @throws FileNotFoundException if there is an error creating or
095     * opening the file
096     */
097    @SuppressWarnings("resource")
098    public XmlStreamWriter(final File file, final String defaultEncoding) throws FileNotFoundException {
099        this(new FileOutputStream(file), defaultEncoding);
100    }
101
102    /**
103     * Returns the detected encoding.
104     *
105     * @return the detected encoding
106     */
107    public String getEncoding() {
108        return encoding;
109    }
110
111    /**
112     * Returns the default encoding.
113     *
114     * @return the default encoding
115     */
116    public String getDefaultEncoding() {
117        return defaultEncoding;
118    }
119
120    /**
121     * Closes the underlying writer.
122     *
123     * @throws IOException if an error occurs closing the underlying writer
124     */
125    @Override
126    public void close() throws IOException {
127        if (writer == null) {
128            encoding = defaultEncoding;
129            writer = new OutputStreamWriter(out, encoding);
130            writer.write(xmlPrologWriter.toString());
131        }
132        writer.close();
133    }
134
135    /**
136     * Flushes the underlying writer.
137     *
138     * @throws IOException if an error occurs flushing the underlying writer
139     */
140    @Override
141    public void flush() throws IOException {
142        if (writer != null) {
143            writer.flush();
144        }
145    }
146
147    /**
148     * Detects the encoding.
149     *
150     * @param cbuf the buffer to write the characters from
151     * @param off The start offset
152     * @param len The number of characters to write
153     * @throws IOException if an error occurs detecting the encoding
154     */
155    private void detectEncoding(final char[] cbuf, final int off, final int len)
156            throws IOException {
157        int size = len;
158        final StringBuffer xmlProlog = xmlPrologWriter.getBuffer();
159        if (xmlProlog.length() + len > BUFFER_SIZE) {
160            size = BUFFER_SIZE - xmlProlog.length();
161        }
162        xmlPrologWriter.write(cbuf, off, size);
163
164        // try to determine encoding
165        if (xmlProlog.length() >= 5) {
166            if (xmlProlog.substring(0, 5).equals("<?xml")) {
167                // try to extract encoding from XML prolog
168                final int xmlPrologEnd = xmlProlog.indexOf("?>");
169                if (xmlPrologEnd > 0) {
170                    // ok, full XML prolog written: let's extract encoding
171                    final Matcher m = ENCODING_PATTERN.matcher(xmlProlog.substring(0,
172                            xmlPrologEnd));
173                    if (m.find()) {
174                        encoding = m.group(1).toUpperCase(Locale.ROOT);
175                        encoding = encoding.substring(1, encoding.length() - 1);
176                    } else {
177                        // no encoding found in XML prolog: using default
178                        // encoding
179                        encoding = defaultEncoding;
180                    }
181                } else if (xmlProlog.length() >= BUFFER_SIZE) {
182                    // no encoding found in first characters: using default
183                    // encoding
184                    encoding = defaultEncoding;
185                }
186            } else {
187                // no XML prolog: using default encoding
188                encoding = defaultEncoding;
189            }
190            if (encoding != null) {
191                // encoding has been chosen: let's do it
192                xmlPrologWriter = null;
193                writer = new OutputStreamWriter(out, encoding);
194                writer.write(xmlProlog.toString());
195                if (len > size) {
196                    writer.write(cbuf, off + size, len - size);
197                }
198            }
199        }
200    }
201
202    /**
203     * Writes the characters to the underlying writer, detecting encoding.
204     *
205     * @param cbuf the buffer to write the characters from
206     * @param off The start offset
207     * @param len The number of characters to write
208     * @throws IOException if an error occurs detecting the encoding
209     */
210    @Override
211    public void write(final char[] cbuf, final int off, final int len) throws IOException {
212        if (xmlPrologWriter != null) {
213            detectEncoding(cbuf, off, len);
214        } else {
215            writer.write(cbuf, off, len);
216        }
217    }
218
219    static final Pattern ENCODING_PATTERN = XmlStreamReader.ENCODING_PATTERN;
220}