transform插件

transform插件用於修改POST請求的body和http響應的body，可以理解爲一個過濾器，或者nginx的body_filter。

有兩個地方可以添加transform回調函數，TS_HTTP_REQUEST_TRANSFORM_HOOK和TS_HTTP_RESPONSE_TRANSFORM_HOOK，分別是過濾POST請求的body和http響應的body。在適當的hook點判斷是否需要使transform邏輯生效，需要的話在相應的transform hook點添加回調函數。比如這樣一個場景，在緩存查詢結束的時候判斷查詢結果是hit還是miss，miss的話需要在http響應body末尾添加一些內容，這樣響應的內容連同被添加的內容會當成整個的響應內容被緩存住，後續的請求如果在緩存查詢結束判斷查詢結果是hit就會將transform後的結果響應出去。

有兩個vio，input_vio(write vio)和output_vio(read vio)。每個vio都有配套的TSIOBuffer和TSIOBufferReader一起工作。

數據的流向是input_vio->分片插件->output_vio。transform插件對應一個continuation和一個回調函數，假設這個回調函數叫plutin_handler，plutin_handler根據不同的event有不同的邏輯。分片插件每transform一些東西，都會觸發上游的TS_EVENT_VCONN_WRITE_READY事件通知上游可以繼續發送東西過來，並且觸發下游的TS_EVENT_IMMEDIATE事件通知下游有東西可以讀了。如果分片插件判斷已經沒有數據需要繼續處理了，會觸發上游的TS_EVENT_VCONN_WRITE_COMPLETE事件，並且觸發下游的TS_EVENT_IMMEDIATE事件。或者可以這樣理解。整個流程分爲三部分，上游，分片插件，下游。下游處理完數據會通過TSContCall觸發上游的TS_EVENT_VCONN_WRITE_READY或TS_EVENT_VCONN_WRITE_COMPLETE事件，上游準備好了數據會通過TSVIOReenable觸發下游的TS_EVENT_IMMEDIATE事件。

ats插件的編譯：tsxs -o hello-world.so -c hello-world.c，注意tsxs要使用與其它可執行文件配套一起生成的，不然的話會出錯。

重點函數：

TSVConnWrite：

tsapi TSVIO TSVConnWrite(TSVConn connp, TSCont contp, TSIOBufferReader readerp, int64_t nbytes);

生成一個vio，有什麼操作的話會觸發contp，要寫的字節數就是nbytes

TSVConnWriteVIOGet：

tsapi TSVIO TSVConnWriteVIOGet(TSVConn connp);

獲取一個VConnection的input_vio(write_vio)，VConnection是input_vio的實施者

TSVIOContGet:

tsapi TSCont TSVIOContGet(TSVIO viop);

獲取一個vio的continuation，continuation是vio的使用者，是transform的上游

TSTransformOutputVConnGet:

tsapi TSVConn TSTransformOutputVConnGet(TSVConn connp);

獲取一個下游的VConnection

TSVIONDoneGet：

tsapi TSIOBuffer TSVIOBufferGet(TSVIO viop);

獲取一個vio已經操作完成的字節數

TSVIONBytesSet：

tsapi void TSVIONBytesSet(TSVIO viop, int64_t nbytes);

設置這個vio將要處理的字節數，第二個參數要比TSVIONDoneGet得到的結果大

TSVIONBytesGet:

tsapi int64_t TSVIONBytesGet(TSVIO viop);

獲取這個vio將要處理的字節數

TSIOBufferCopy：

tsapi int64_t TSIOBufferCopy(TSIOBuffer bufp, TSIOBufferReader readerp, int64_t length, int64_t offset);

將內容從一個TSIOBufferReader拷貝到一個TSIOBuffer，拷貝的字節數是length，偏移量是offset

TSIOBufferReaderConsume：

tsapi void TSIOBufferReaderConsume(TSIOBufferReader readerp, int64_t nbytes);

拷貝操作結束後要執行這個函數，要消費input_vio指定字節數

TSVIONDoneSet：

tsapi void TSVIONDoneSet(TSVIO viop, int64_t ndone);

拷貝操作結束後要執行這個函數，要告知input_vio已經消費了多少字節

TSVConnClosedGet：

tsapi int TSVConnClosedGet(TSVConn connp);

判斷connp是否已經被關閉，connp是transform插件本身的contiuation

如下是一個完整的transform插件，實現的功能是在讀源站的響應頭時添加transform插件。如果緩存查詢結果是hit則不需要回源，也就不會添加transform插件了。插件的邏輯是在響應末尾添加"hello world\n"字符串。

#include <limits.h>
#include <stdio.h>
#include <string.h>
#include <ts/ts.h>

#define ASSERT_SUCCESS(_x) TSAssert ((_x) == TS_SUCCESS)

typedef struct {
	TSVIO output_vio;
	TSIOBuffer output_buffer;
	TSIOBufferReader output_reader;
	int append_needed;
} MyData;

static TSIOBuffer append_buffer;
static TSIOBufferReader append_buffer_reader;
static int append_buffer_length;

static MyData *
my_data_alloc() {
	MyData *data;

	data = (MyData *) TSmalloc(sizeof(MyData));
	TSReleaseAssert(data);
	data->output_vio = NULL;
	data->output_buffer = NULL;
	data->output_reader = NULL;
	data->append_needed = 1;

	return data;
}

static void my_data_destroy(MyData * data) {
	if (data) {
		if (data->output_buffer)
			TSIOBufferDestroy(data->output_buffer);
		TSfree(data);
	}
}

static void handle_transform(TSCont contp) {
	TSVConn output_conn;
	TSVIO write_vio;
	MyData *data;
	int64_t towrite;
	int64_t avail;

	/*獲取transform插件的下游VConnection，這個VConnection會接收transform插件的處理完的數據
	  這個VConnection會觸發contp的TS_EVENT_VCONN_WRITE_READY，TS_EVENT_VCONN_WRITE_COMPLETE
	  或TS_EVENT_ERROR事件*/
	output_conn = TSTransformOutputVConnGet(contp);
	/*獲取input_vio，通過input_vio可以獲取到input_buffer*/
	write_vio = TSVConnWriteVIOGet(contp);
	/*獲取這個contp的數據，沒有的話初始化並且TSContDataSet*/
	data = TSContDataGet(contp);
	if (!data) {
		towrite = TSVIONBytesGet(write_vio);
		if (towrite != INT64_MAX) {
			towrite += append_buffer_length;
		}
		data = my_data_alloc();
		data->output_buffer = TSIOBufferCreate();
		data->output_reader = TSIOBufferReaderAlloc(data->output_buffer);
		/*獲取output_vio，並且指定要向其傳送的字節數，字節數就是http響應body的長度加上"hello world\n"的長度*/
		data->output_vio = TSVConnWrite(output_conn, contp, data->output_reader, towrite);
		TSContDataSet(contp, data);
	}
	/*如果input_buffer已經爲空，說明上游不會再有數據傳到transform插件*/
	if (!TSVIOBufferGet(write_vio)) {
		if (data->append_needed) {
			data->append_needed = 0;
			TSIOBufferCopy(TSVIOBufferGet(data->output_vio), append_buffer_reader, append_buffer_length, 0);
		}
		TSVIONBytesSet(data->output_vio, TSVIONDoneGet(write_vio) + append_buffer_length);
		TSVIOReenable(data->output_vio);

		return;
	}
	/*獲取還有多少數據需要處理，也即還有多少東西會從上游發送來到transform插件*/
	towrite = TSVIONTodoGet(write_vio);
	if (towrite > 0) {
		avail = TSIOBufferReaderAvail(TSVIOReaderGet(write_vio));
		if (towrite > avail) {
			towrite = avail;
		}
		if (towrite > 0) {
			TSIOBufferCopy(TSVIOBufferGet(data->output_vio), TSVIOReaderGet(write_vio), towrite, 0);
			TSIOBufferReaderConsume(TSVIOReaderGet(write_vio), towrite);
			TSVIONDoneSet(write_vio, TSVIONDoneGet(write_vio) + towrite);
		}
	}
	/*獲取還有多少數據需要處理，也即還有多少東西會從上游發送來到transform插件*/
	if (TSVIONTodoGet(write_vio) > 0) {
		/*如果上游還有數據需要傳送*/
		if (towrite > 0) {
			/*通知下游可以讀數據了*/
			TSVIOReenable(data->output_vio);
			/*如果上游尚有數據會發送過來，需要通知上游可以繼續發送數據了*/
			TSContCall(TSVIOContGet(write_vio), TS_EVENT_VCONN_WRITE_READY, write_vio);
		}
	}
	else {
		/*如果上游已經沒有更多的數據，需要執行append的邏輯了，也即將"hello world\n"添加到響應數據末尾*/
		if (data->append_needed) {
			data->append_needed = 0;
			TSIOBufferCopy(TSVIOBufferGet(data->output_vio), append_buffer_reader, append_buffer_length, 0);
		}
		/*設置output_vio需要處理的字節數，這樣output_vio處理完這些內容就可以發送TS_EVENT_VCONN_WRITE_COMPLETE
		  給transform插件了*/
		TSVIONBytesSet(data->output_vio, TSVIONDoneGet(write_vio) + append_buffer_length);
		/*通知下游可以讀數據了*/
		TSVIOReenable(data->output_vio);
		/*如果上游尚有數據會發送過來，需要通知上游可以繼續發送數據了*/
		TSContCall(TSVIOContGet(write_vio), TS_EVENT_VCONN_WRITE_COMPLETE, write_vio);
	}
}

static int append_transform(TSCont contp, TSEvent event, void *edata) {
	/*如果transform插件已經被關閉，需要釋放資源*/
	if (TSVConnClosedGet(contp)) {
		my_data_destroy(TSContDataGet(contp));
		TSContDestroy(contp);
		return 0;
	} else {
		switch (event) {
		case TS_EVENT_ERROR: {
			TSVIO write_vio;
			write_vio = TSVConnWriteVIOGet(contp);
			/*通知上游出錯了*/
			TSContCall(TSVIOContGet(write_vio), TS_EVENT_ERROR, write_vio);
		}
			break;
		case TS_EVENT_VCONN_WRITE_COMPLETE:
			/*如果下游觸發了TS_EVENT_VCONN_WRITE_COMPLETE事件，需要關閉下游的VConnection*/
			TSVConnShutdown(TSTransformOutputVConnGet(contp), 0, 1);
			break;
		case TS_EVENT_VCONN_WRITE_READY:
		default:
			handle_transform(contp);
			break;
		}
	}

	return 0;
}

static int transformable(TSHttpTxn txnp) {
	TSMBuffer bufp;
	TSMLoc hdr_loc;
	TSHttpStatus resp_status;

	TSHttpTxnServerRespGet(txnp, &bufp, &hdr_loc);
	//判斷源站給的狀態碼，只有源站返回200時纔會執行transform邏輯
	if (TS_HTTP_STATUS_OK == (resp_status = TSHttpHdrStatusGet(bufp, hdr_loc))) {
		ASSERT_SUCCESS(TSHandleMLocRelease(bufp, TS_NULL_MLOC, hdr_loc));
		return 1;
	} else {
		ASSERT_SUCCESS(TSHandleMLocRelease(bufp, TS_NULL_MLOC, hdr_loc));
		return 0;
	}
}

static int transform_plugin(TSCont contp, TSEvent event, void *edata) {
	TSHttpTxn txnp = (TSHttpTxn) edata;

	switch (event) {
	case TS_EVENT_HTTP_READ_RESPONSE_HDR:
		//只有緩存查詢結果爲miss的時候纔會回源，纔會執行這個case的邏輯
		if(transformable(txnp))
		{
			TSVConn connp;
			connp = TSTransformCreate(append_transform, txnp);
			TSHttpTxnHookAdd(txnp, TS_HTTP_RESPONSE_TRANSFORM_HOOK, connp);
		}

		TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
		return 0;
	default:
		break;
	}

	return 0;
}

static int data_prepared() {
	append_buffer = TSIOBufferCreate();
	append_buffer_reader = TSIOBufferReaderAlloc(append_buffer);
	TSIOBufferWrite(append_buffer, "hello world\n", sizeof("hello world\n") - 1);
	append_buffer_length = TSIOBufferReaderAvail(append_buffer_reader);

	return 1;
}

int check_ts_version() {

	const char *ts_version = TSTrafficServerVersionGet();
	int result = 0;

	if (ts_version) {
		int major_ts_version = 0;
		int minor_ts_version = 0;
		int patch_ts_version = 0;

		if (sscanf(ts_version, "%d.%d.%d", &major_ts_version, &minor_ts_version, &patch_ts_version) != 3) {
			return 0;
		}
		if (major_ts_version >= 2) {
			result = 1;
		}
	}

	return result;
}

void TSPluginInit(int argc, const char *argv[]) {
	TSPluginRegistrationInfo info;
	info.plugin_name = "append-transform";
	info.vendor_name = "MyCompany";
	info.support_email = "[email protected]";

	if (TSPluginRegister(TS_SDK_VERSION_3_0, &info) != TS_SUCCESS) {
		TSError("Plugin registration failed.\n");
		goto Lerror;
	}
	if (!check_ts_version()) {
		TSError("Plugin requires Traffic Server 3.0 or later\n");
		goto Lerror;
	}
	if (!data_prepared()) {
		TSError("[append-transform] Failed to prepared data\n");
		goto Lerror;
	}
	TSHttpHookAdd(TS_HTTP_READ_RESPONSE_HDR_HOOK,TSContCreate(transform_plugin, NULL));
	return;

Lerror:
	TSError("[append-transform] Unable to initialize plugin\n");
}