// stats plugin
//
// Counts packets of types icmp, tcp and udp.
// Counts returned on request.
//

// Standard includes - do not change
#include "stdinc.h"

#include "stats.h"

//
// This defines name of the plugin class to the runtime environment
//
MOD_MISC("stats")

//
// Static plugin class structure. Defined for the CLASS, not each instance.
// Change only the name
//
static struct rp_class stats_class;

//
// Create and initialize an instance of stats.
//
// This is where we initialize any per instance variables.
// Most of this code is standard stuff and should be left
// unchanged, except for the plugin name.
//
struct rp_instance * stats_create_instance(
	struct rp_class *theClass,	// pointer to class structure
	u_int32_t instanceid		// identifier for new instance
) {
	struct stats_instance *newInst; 

	// Allocate memory for new instance of stats.
	MSR_PLUGIN_MALLOC(newInst, struct stats_instance *, 
	                  sizeof(struct stats_instance),
	                  M_MSR, 
	                  M_WAITOK);
	if (newInst == NULL) return NULL;

	// Fill in instance pointers to local methods
	newInst->rootinstance.rpclass         = &stats_class;
	newInst->rootinstance.handle_packet   = stats_handle_packet;
	newInst->rootinstance.free_instance   = stats_free_instance;
	newInst->rootinstance.bind_instance   = stats_bind_instance;
	newInst->rootinstance.unbind_instance = stats_unbind_instance;
	newInst->rootinstance.handle_msg      = stats_handle_msg;

	newInst->rootinstance.instanceid = instanceid;

	//
	// Initialization of per instance variables goes here.
	//
	newInst->icmpCnt = 0; newInst->tcpCnt = 0; newInst->udpCnt = 0;
	newInst->icmpTot = 0; newInst->tcpTot = 0; newInst->udpTot = 0;

	return (struct rp_instance *) newInst;
}

//
// Handle an incoming packet.
//
void stats_handle_packet(
	struct rp_instance *this,	// pointer to instance structure
	void *bufferList		// pointer to list of packet buffers
) {
	struct stats_instance *inst = (struct stats_instance *) this;
	msr_bufhdr_t *buffer = msr_firstBuffer(bufferList);
	struct ip *iph = msr_pkt_iph(buffer);
	int len = msr_iplen(iph);
	int proto = msr_ipproto(iph);

	if (proto == 1) { 
		inst->icmpCnt++; inst->icmpTot++; 
	} else if (proto == 6) {
		inst->tcpCnt++; inst->tcpTot++; 
	} else if (proto == 17) {
		inst->udpCnt++; inst->udpTot++;
	}

	// remove packet from input list
    	msr_removeBuffer(bufferList, buffer);     
    	msr_freeBuffer(buffer);     
}

//
// Handle a control message.
//
// The message is simply a list of integers. The first integer defines
// a message type. The remaining values can be used to specify parameters.
// A response can be returned by over-writing the values in the message buffer.
// The length argument must be updated correctly, to indicate the size of the
// response. Length of message and response limited to xx bytes.
//
// stats defines two message types
//
//	1. Return icmpCnt, tcpCnt and udpCnt, then clear.
//	2. Return icmpTot, tcpTot and udpTot.
//

int stats_handle_msg(
	struct rp_instance *this,	// pointer to instance structure
	void *buf,			// buffer containing integer values
					// buf[0]=instance id
					// buf[1]=command type
					// buf[2]=first parameter value
					// buf[3]=second parameter value ...
	u_int8_t flags,			// ???
	u_int8_t seq,			// sequence number of message
	u_int8_t *len			// number of values in buf
) {
	struct stats_instance *inst = (struct stats_instance *) this;
	u_int32_t *vals = (u_int32_t *) buf;
	u_int32_t id  = (u_int32_t) ntohl(vals[0]);
	u_int32_t typ = (u_int32_t) ntohl(vals[1]);

	if (typ == 1) { // return Cnt values and clear
		vals[0] = (u_int32_t) htonl(inst->icmpCnt);
		vals[1] = (u_int32_t) htonl(inst->tcpCnt);
		vals[2] = (u_int32_t) htonl(inst->udpCnt);
		*len = 3*sizeof(u_int32_t);
		inst->icmpCnt = 0; inst->tcpCnt = 0; inst->udpCnt = 0;
	} else if (typ == 2) { // set count and length
		vals[0] = (u_int32_t) htonl(inst->icmpTot);
		vals[1] = (u_int32_t) htonl(inst->tcpTot);
		vals[2] = (u_int32_t) htonl(inst->udpTot);
		*len = 3*sizeof(u_int32_t);
	}

	return 0;
}

//
// Free instance structure.
//
void stats_free_instance(struct rp_instance *this) {
	// perform any cleanup of per-instance data structures that
	// may be required.
}

//
// Bind/unbind plugin instance to/from a filter
//
void stats_bind_instance(struct rp_instance *this) {
	// perform any initialization that should be done whenever the
	// plugin is bound to a filter
}

void stats_unbind_instance(struct rp_instance *this) {
	// perform any cleanup that should be done whenever the
	// plugin is unbound a filter

}

//
// The remaining code is standard stuff. The only thing that you
// typically need to change is the various places where the
// plugin name (stats) appears in function and type names.
//

//
// Initialize the class. 
//
void stats_init_class() {
	stats_class.classid = stats_ID;
	stats_class.itype = RP_INTERFACE_TYPE_HLIST;
	stats_class.create_instance = stats_create_instance;
	return;
}

//
// Return class structure
//
struct rp_class * stats_get_class() {
	return &stats_class;
}

//
// External kernel module entry point.
//
// It is called each time the class is loaded or unloaded.
// The stat information is not needed here, so we will 
// leave it lkm_nofunc().
//

int stats(struct lkm_table *lkmtp, int cmd, int ver,
		struct kernel_plugin_fct_struct *fctPtr) {
	// Do NOT put anything before the kernel_plugin_fcts is set,
	// especially if it uses any of the MSR_ macros!!!
	if (kernel_plugin_fcts == NULL) {
		kernel_plugin_fcts = fctPtr;
	  	kernel_plugin_variables = fctPtr->pluginVariables;
	}
	
	MSR_PLUGIN_DISPATCH(lkmtp, cmd, ver, stats_load,
			stats_unload, PLUGIN_LKM_NOFUNC_FCT);
}

//
// Load class. This function is called each time the class is loaded
//
int stats_load(struct lkm_table *lkmtp, int cmd) {
	int err;

	if (PLUGIN_LKM_EXISTS_FCT(lkmtp)) { // avoid loading twice
		return (EEXIST);
	}
	stats_init_class();

	err = PLUGIN_PCU_REGISTER_CLASS_FCT(stats_get_class());
	if (err != RP_OK) {
		return -1;
	}

	return 0;
}

//
// Unload class. This function is called each time the class is unloaded.
// Remove all existing instances and then remove class.
//
int stats_unload(struct lkm_table *lkmtp, int cmd) {
	struct rp_class *rpclass;
	u_int32_t cid;

	rpclass = stats_get_class();
	cid = rpclass->classid;

	if (PLUGIN_PCU_FREE_ALL_INSTANCES_FCT(rpclass) != RP_OK) {
		return -1;
	}

	return 0;
}