postwriter.java

测试工具

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * 
 */

package org.apache.jmeter.protocol.http.sampler;

import java.io.BufferedInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;
import java.net.URLConnection;

import org.apache.jmeter.protocol.http.util.HTTPArgument;
import org.apache.jmeter.protocol.http.util.HTTPConstants;
import org.apache.jmeter.testelement.property.PropertyIterator;

/**
 * Class for setting the necessary headers for a POST request, and sending the
 * body of the POST.
 */
public class PostWriter {
    
	private static final String DASH_DASH = "--";  // $NON-NLS-1$
	private static final byte[] DASH_DASH_BYTES = DASH_DASH.getBytes();

    /** The bounday string between multiparts */
    protected final static String BOUNDARY = "---------------------------7d159c1302d0y0"; // $NON-NLS-1$

	private final static byte[] CRLF = { 0x0d, 0x0A };

	public static final String ENCODING = "ISO-8859-1"; // $NON-NLS-1$

    /** The form data that is going to be sent as url encoded */
    protected byte[] formDataUrlEncoded;    
    /** The form data that is going to be sent in post body */
    protected byte[] formDataPostBody;
    /** The start of the file multipart to be sent */
    private byte[] formDataFileStartMultipart;
    /** The boundary string for multipart */
    private final String boundary;
    
    /**
     * Constructor for PostWriter.
     * Uses the PostWriter.BOUNDARY as the boundary string
     * 
     */
    public PostWriter() {
        this(BOUNDARY);
    }

    /**
     * Constructor for PostWriter
     * 
     * @param boundary the boundary string to use as marker between multipart parts
     */
    public PostWriter(String boundary) {
        this.boundary = boundary;
    }

	/**
	 * Send POST data from Entry to the open connection.
     * 
     * @return the post body sent. Actual file content is not returned, it
     * is just shown as a placeholder text "actual file content"
	 */
	public String sendPostData(URLConnection connection, HTTPSampler sampler) throws IOException {
        // Buffer to hold the post body, except file content
        StringBuffer postedBody = new StringBuffer(1000);
        
        // Check if we should do a multipart/form-data or an
        // application/x-www-form-urlencoded post request
        if(sampler.getUseMultipartForPost()) {
            OutputStream out = connection.getOutputStream();
            
            // Write the form data post body, which we have constructed
            // in the setHeaders. This contains the multipart start divider
            // and any form data, i.e. arguments
            out.write(formDataPostBody);
            // We get the posted bytes as UTF-8, since java is using UTF-8
            postedBody.append(new String(formDataPostBody, "UTF-8")); // $NON-NLS-1$
            
            // Add any files
            if(sampler.hasUploadableFiles()) {
                // First write the start multipart file
                out.write(formDataFileStartMultipart);
                // We get the posted bytes as UTF-8, since java is using UTF-8
                postedBody.append(new String(formDataFileStartMultipart, "UTF-8")); // $NON-NLS-1$
                
                // Write the actual file content
                writeFileToStream(sampler.getFilename(), out);
                // We just add placeholder text for file content
                postedBody.append("<actual file content, not shown here>"); // $NON-NLS-1$

                // Write the end of multipart file
                byte[] fileMultipartEndDivider = getFileMultipartEndDivider(); 
                out.write(fileMultipartEndDivider);
                // We get the posted bytes as UTF-8, since java is using UTF-8
                postedBody.append(new String(fileMultipartEndDivider, "UTF-8")); // $NON-NLS-1$
            }

            // Write end of multipart
            byte[] multipartEndDivider = getMultipartEndDivider(); 
            out.write(multipartEndDivider);
            // We get the posted bytes as UTF-8, since java is using UTF-8
            postedBody.append(new String(multipartEndDivider, "UTF-8")); // $NON-NLS-1$

            out.flush();
            out.close();
        }
        else {
            // If there are no arguments, we can send a file as the body of the request
            if(sampler.getArguments() != null && !sampler.hasArguments() && sampler.getSendFileAsPostBody()) {
                OutputStream out = connection.getOutputStream();
                writeFileToStream(sampler.getFilename(), out);
                out.flush();
                out.close();

                // We just add placeholder text for file content
                postedBody.append("<actual file content, not shown here>"); // $NON-NLS-1$
            }
            else if (formDataUrlEncoded != null){ // may be null for PUT           
                // In an application/x-www-form-urlencoded request, we only support
                // parameters, no file upload is allowed
                OutputStream out = connection.getOutputStream();
                out.write(formDataUrlEncoded);
                out.flush();
                out.close();

                // We get the posted bytes as UTF-8, since java is using UTF-8
                postedBody.append(new String(formDataUrlEncoded, "UTF-8")); // $NON-NLS-1$
            }            
        }
        return postedBody.toString();
	}
    
    public void setHeaders(URLConnection connection, HTTPSampler sampler) throws IOException {
    	// Get the encoding to use for the request
        String contentEncoding = sampler.getContentEncoding();
        if(contentEncoding == null || contentEncoding.length() == 0) {
            contentEncoding = ENCODING;
        }
        long contentLength = 0L;
    	
        // Check if we should do a multipart/form-data or an
        // application/x-www-form-urlencoded post request
        if(sampler.getUseMultipartForPost()) {
            // Set the content type
            connection.setRequestProperty(
            		HTTPConstants.HEADER_CONTENT_TYPE,
            		HTTPConstants.MULTIPART_FORM_DATA + "; boundary=" + getBoundary()); // $NON-NLS-1$
            
            // Write the form section
            ByteArrayOutputStream bos = new ByteArrayOutputStream();

            // First the multipart start divider
            bos.write(getMultipartDivider());
            // Add any parameters
            PropertyIterator args = sampler.getArguments().iterator();
            while (args.hasNext()) {
                HTTPArgument arg = (HTTPArgument) args.next().getObjectValue();
                String parameterName = arg.getName();
                if (parameterName.length()==0){
                    continue; // Skip parameters with a blank name (allows use of optional variables in parameter lists)
                }
                // End the previous multipart
                bos.write(CRLF);
                // Write multipart for parameter
                writeFormMultipart(bos, parameterName, arg.getValue(), contentEncoding);
            }
            // If there are any files, we need to end the previous multipart
            if(sampler.hasUploadableFiles()) {
                // End the previous multipart
                bos.write(CRLF);
            }
            bos.flush();
            // Keep the content, will be sent later
            formDataPostBody = bos.toByteArray();
            bos.close();
            contentLength = formDataPostBody.length;

            // Now we just construct any multipart for the files
            // We only construct the file multipart start, we do not write
            // the actual file content
            if(sampler.hasUploadableFiles()) {
                bos = new ByteArrayOutputStream();
                // Write multipart for file
                writeStartFileMultipart(bos, sampler.getFilename(), sampler.getFileField(), sampler.getMimetype());
                bos.flush();
                formDataFileStartMultipart = bos.toByteArray();
                bos.close();
                contentLength += formDataFileStartMultipart.length;
                // Add also the length of the file content
                File uploadFile = new File(sampler.getFilename());
                contentLength += uploadFile.length();
                // And the end of the file multipart