编写JMeter扩展（1）编写采样器代码

Apache JMeter自带了许多的采样器供我们使用，而且能够满足大部分的测试需求。但是在实际使用过程中，难免需要针对项目自身的特点和需求对Apache JMeter进行扩展。虽然直接继承AbstractJavaSamplerClient即可编写Java采样器，但是其相对应的GUI界面不是很友好，或者我们想编写一个类似HTTP Request那样的sampler，该如何做呢？比如我们需要编写一个TLS的采样器，该采样器允许用户在发送TLS请求的时候制定TLS的版本，以及客户端证书等信息。请参考如下代码及注释：

package kg.apc.jmeter.samplers;

import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.Serializable;
import java.net.Socket;
import java.security.KeyStore;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLParameters;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;

import org.apache.jmeter.samplers.AbstractSampler;
import org.apache.jmeter.samplers.Entry;
import org.apache.jmeter.samplers.Interruptible;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.jorphan.logging.LoggingManager;
import org.apache.log.Logger;

// FIXME: actually keep-alive does not work!
/**
 *
 * @author undera
 */
public class TLSRawSampler extends AbstractSampler implements Serializable, Cloneable, Interruptible{
    private static final Logger log = LoggingManager.getLoggerForClass();
    
    //following static final String are property keys in JMeter
    private static final String HOST_NAME = "TLS_hostName";
    private static final String PORT = "TLS_port";
    private static final String TLS_VERSON = "TLS_tlsversion";
    private static final String CIPHER_LIST = "TLS_cipherlist";
    private static final String DATA = "TLS_data";
    private static final String CLIENT_CERT = "TLS_client_cert";
    private static final String CLIENT_CERT_PASSWORD = "TLS_client_cert_password";
 
    public TLSRawSampler() {
    }
    
    /**
     * Process sample here
     */
    @Override
    public SampleResult sample(Entry entry) {
       SampleResult sr = new SampleResult();
       sr.sampleStart();
       sr.setSuccessful(true);
       sr.setSampleLabel(getName());
       sr.setResponseCode("200");
       
       //do the sample here
       try{
    	   Socket socket = getTLSSocket();
    	   socket.getOutputStream().write(getData().getBytes());
    	   StringBuilder headerSB = new StringBuilder();
    	   StringBuilder contentSB = new StringBuilder();
    	   BufferedReader reader = new BufferedReader(new InputStreamReader(socket.getInputStream()));
    	   String line = null;
    	   boolean header = true;
    	   while((line = reader.readLine()) != null){
    		   if(line.length() == 0){
    			   header = false;
    			   continue;
    		   }
    		   if(header){
	    		   headerSB.append(line);
	    		   headerSB.append("\n");
    		   }
    		   else{
    			   contentSB.append(line);
    			   contentSB.append("\n");
    		   }
    	   }
    	   sr.setResponseHeaders(headerSB.toString());
    	   sr.setResponseData(contentSB.toString(), "");
    	   socket.close();
       }catch(Exception e){
    	   sr.setSuccessful(false);
    	   sr.setResponseCode("500");
    	   sr.setResponseMessage(e.getMessage());
       }
       
       sr.sampleEnd();
       return sr;
    }

	@Override
	public boolean interrupt() {
		// TODO Auto-generated method stub
		return false;
	}

	//following get and set methods are used for get and set properties in JMeter
    
    public String getHostName(){
    	return getPropertyAsString(TLSRawSampler.HOST_NAME);
    }
    
    public void setHostName(String hostName){
    	setProperty(TLSRawSampler.HOST_NAME, hostName);
    }
    
    public String getPort(){
    	return getPropertyAsString(TLSRawSampler.PORT);
    }
    
    public void setPort(String port){
    	setProperty(TLSRawSampler.PORT, port);
    }
    
    public String getTLSVersion(){
    	return getPropertyAsString(TLSRawSampler.TLS_VERSON);
    }
    
    public void setTLSVersion(String version){
    	setProperty(TLSRawSampler.TLS_VERSON, version);
    }
    
    public String getClientCert(){
    	return getPropertyAsString(TLSRawSampler.CLIENT_CERT);
    }
    
    public void setClientCert(String clientCert){
    	setProperty(TLSRawSampler.CLIENT_CERT, clientCert);
    }
    
    public String getCipherList(){
    	return getPropertyAsString(TLSRawSampler.CIPHER_LIST);
    }
    
    public void setCipherList(String list){
    	setProperty(TLSRawSampler.CIPHER_LIST, list);
    }
    
    public String getData(){
    	return getPropertyAsString(TLSRawSampler.DATA);
    }
    
    public void setData(String data){
    	setProperty(TLSRawSampler.DATA, data);
    }
    
    public String getClientCertPassword(){
    	return getPropertyAsString(TLSRawSampler.CLIENT_CERT_PASSWORD);
    }
    
    public void setClientCertPassword(String password){
    	setProperty(TLSRawSampler.CLIENT_CERT_PASSWORD, password);
    }
    
    private Socket getTLSSocket() throws Exception{
		TrustManager[] allTrusted = {new X509TrustManager(){

			@Override
			public void checkClientTrusted(X509Certificate[] chain,
					String authType) throws CertificateException {
				// TODO Auto-generated method stub

			}

			@Override
			public void checkServerTrusted(X509Certificate[] chain,
					String authType) throws CertificateException {
				// TODO Auto-generated method stub

			}

			@Override
			public X509Certificate[] getAcceptedIssuers() {
				// TODO Auto-generated method stub
				return null;
			}

		}};

		SSLContext ctx = SSLContext.getInstance(getTLSVersion(), "SunJSSE");//use 1.2

		//get client cert
		String clientCert = getClientCert();
		if(clientCert != null && clientCert.length() > 0){
			String clientCertPassword = getClientCertPassword();
			KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance("SunX509");
			KeyStore keyStore = KeyStore.getInstance("PKCS12");

			InputStream keyInput = new FileInputStream(clientCert);
			keyStore.load(keyInput, clientCertPassword.toCharArray());
			keyInput.close();

			keyManagerFactory.init(keyStore, clientCertPassword.toCharArray());

			ctx.init(keyManagerFactory.getKeyManagers(), allTrusted, null);
		}
		else{
			ctx.init(null, allTrusted, null);
		}
		SSLParameters parameters = ctx.getDefaultSSLParameters();
		String cipherList = getCipherList();
		if(cipherList != null && cipherList.length() > 0){
			parameters.setCipherSuites(cipherList.trim().split(","));
		}
		for(String cipher : parameters.getCipherSuites()){
			log.info(cipher);
		}

		SSLSocket socket = (SSLSocket)ctx.getSocketFactory().createSocket(getHostName(), Integer.valueOf(getPort()));
		socket.setTcpNoDelay(true);
		socket.setSSLParameters(parameters);

		return socket;
	}
}