#!/usr/sbin/dtrace -Cs
/*
 * udpsnoop.d - snoop UDP network I/O by process.
 *              Written using DTrace (Solaris 10 3/05)
 *
 * This analyses UCP network I/O and prints the responsible PID and UID,
 * plus standard details such as IP address and port. This tracks UDP
 * read/writes by payload; this is the same as by packet (although ~42 fewer
 * bytes), unless large datagrams are sent and then fragmented by IP - in 
 * which case it becomes obvious that we are printing by payload (sizes > MTU).
 * 
 * This program can help identify which processes is causing UDP traffic.
 *
 * 04-Feb-2007, ver 0.56        (check for newer versions)
 *
 * USAGE:       udpsnoop.d
 *
 * FIELDS:
 *              UID             user ID
 *              PID             process ID
 *              CMD             command
 *              LADDR           local IP address
 *              RADDR           remote IP address
 *              LPORT           local port number
 *              RPORT           remote port number
 *              DR              direction
 *              SIZE            payload size, bytes
 *
 * SEE ALSO: snoop -rS, tcpsnoop
 *
 * ATTENTION: This script is currently fragile as it uses the "fbt" DTrace
 * provider, which is an unstable interface. This script is likely to fail
 * on future Solaris versions, including minor updates. This fbt based
 * script was released to provide temporary observability until a stable
 * network provider has been released and this script can be rewritten.
 * See the OpenSolaris dtrace-discuss mailing list for more information.
 *
 * COPYRIGHT: Copyright (c) 2005 Brendan Gregg.
 *
 *  This program is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU General Public License
 *  as published by the Free Software Foundation; either version 2
 *  of the License, or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software Foundation,
 *  Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 *
 *  (http://www.gnu.org/copyleft/gpl.html)
 *
 * Author: Brendan Gregg  [Sydney, Australia]
 *
 * Known Bugs:
 * - destination address is sometimes equal to source address when it 
 *   shouldn't. This has been observed with RPC traffic.
 * - outbound traceroute UDP packets are not detected (nor are they by 
 *   mib:::udpOutDatagrams).
 * - I'd rather this was packet based than payload based.
 *
 * ToDo: IPv6
 *
 * 14-Jul-2005  Brendan Gregg   Created this.
 * 04-Feb-2007	   "	  "	Updates for newer Solaris.
 */

#pragma D option quiet

#include <inet/common.h>
#include <sys/byteorder.h>

/*
 * Print Header
 */
dtrace:::BEGIN
{
	/* print main headers */
	printf("%5s %6s %-15s %5s %2s %-15s %5s %5s %s\n",
	    "UID", "PID", "LADDR", "LPORT", "DR", "RADDR", "RPORT",
	    "SIZE", "CMD");
}

/*
 * UDP Process Write
 */

/*
 * UDP Track this thread as interesting (mib is our checkpoint)
 */
mib:::udp*OutDatagrams
{
	self->udpOutDatagrams = 1;
}

/*
 * UDP Save IP header info if tracked
 */
fbt::ip_output:entry
/self->udpOutDatagrams/
{
	self->wipha = (ipha_t *)args[1]->b_rptr;
        this->ipha_hlen = (self->wipha->ipha_version_and_hdr_length & 15) * 4;
	self->wsize = msgdsize(args[1]) - this->ipha_hlen - 8;
}

/*
 * UDP Fetch data and print output line
 */
fbt::ip_output:return
/self->udpOutDatagrams/
{

	/* fetch ports */
        this->ipha_hlen = (self->wipha->ipha_version_and_hdr_length & 15) * 4;
        this->udpha = (udpha_t *)(((uchar_t *)self->wipha) + this->ipha_hlen);
#if defined(_BIG_ENDIAN)
	this->sport = this->udpha->uha_src_port;
	this->dport = this->udpha->uha_dst_port;
#else
	this->sport = BSWAP_16(this->udpha->uha_src_port);
	this->dport = BSWAP_16(this->udpha->uha_dst_port);
#endif

	/* fetch IPv4 addresses */
	this->s0 = (self->wipha->ipha_src & 0x000000ff);
	this->s1 = (self->wipha->ipha_src & 0x0000ff00) >> 8;
	this->s2 = (self->wipha->ipha_src & 0x00ff0000) >> 16;
	this->s3 = (self->wipha->ipha_src & 0xff000000) >> 24;
	this->d0 = (self->wipha->ipha_dst & 0x000000ff);
	this->d1 = (self->wipha->ipha_dst & 0x0000ff00) >> 8;
	this->d2 = (self->wipha->ipha_dst & 0x00ff0000) >> 16;
	this->d3 = (self->wipha->ipha_dst & 0xff000000) >> 24;

	/* convert type for use with lltostr() */
	this->s0 = this->s0 < 0 ? 256 + this->s0 : this->s0;
	this->s1 = this->s1 < 0 ? 256 + this->s1 : this->s1;
	this->s2 = this->s2 < 0 ? 256 + this->s2 : this->s2;
	this->s3 = this->s3 < 0 ? 256 + this->s3 : this->s3;
	this->d0 = this->d0 < 0 ? 256 + this->d0 : this->d0;
	this->d1 = this->d1 < 0 ? 256 + this->d1 : this->d1;
	this->d2 = this->d2 < 0 ? 256 + this->d2 : this->d2;
	this->d3 = this->d3 < 0 ? 256 + this->d3 : this->d3;

	/* stringify addresses */
	self->saddr = strjoin(lltostr(this->s0), ".");
	self->saddr = strjoin(self->saddr, strjoin(lltostr(this->s1), "."));
	self->saddr = strjoin(self->saddr, strjoin(lltostr(this->s2), "."));
	self->saddr = strjoin(self->saddr, lltostr(this->s3));
	self->daddr = strjoin(lltostr(this->d0), ".");
	self->daddr = strjoin(self->daddr, strjoin(lltostr(this->d1), "."));
	self->daddr = strjoin(self->daddr, strjoin(lltostr(this->d2), "."));
	self->daddr = strjoin(self->daddr, lltostr(this->d3));

	/* print output line */
	printf("%5d %6d %-15s %5d -> %-15s %5d %5d %s\n", uid, pid,
	    self->saddr, this->sport, self->daddr, this->dport, 
	    self->wsize, execname);

	/* clear variables */
	self->wipha = 0;
	self->wsize = 0;
	self->saddr = 0;
	self->daddr = 0;
	self->udpOutDatagrams = 0;
}

/*
 * UDP Process Read
 */

/*
 * The following stores details of the UDP connection. In the
 * future this will be replaced with a struct, when support for
 * zeroing entries in an array of structs is added to DTrace.
 */
int queue[int];
int sport[int];
int dport[int];
string saddr[int];
string daddr[int];

/*
 * The following is some stateful code that is intended to identify
 * getq_noneab()s that began with a udp_rput which also triggered
 * a mib:::udpInDatagrams. This combination tracks an inbound UDP
 * packet from kernel to userland. The following code blocks are
 * listed in chronological order.
 */

/* 
 * UDP Track all rputs
 */ 
fbt::udp_rput:entry
{
	self->qp = args[0];
	self->ripha = (ipha_t *)args[1]->b_rptr;
}

/*
 * UDP Mark this rput as interesting (mib is our checkpoint)
 */
mib:::udp*InDatagrams
/self->qp/
{
	self->udpInDatagrams = 1;

	/* fetch ports */
	this->ipha_hlen = (self->ripha->ipha_version_and_hdr_length & 15) * 4;
	this->udpha = (udpha_t *)(((uchar_t *)self->ripha) + this->ipha_hlen);
#if defined(_BIG_ENDIAN)
	self->sport = this->udpha->uha_src_port;
	self->dport = this->udpha->uha_dst_port;
#else
	self->sport = BSWAP_16(this->udpha->uha_src_port);
	self->dport = BSWAP_16(this->udpha->uha_dst_port);
#endif

	/* fetch IPv4 addresses */
	this->s0 = (self->ripha->ipha_src & 0x000000ff);
	this->s1 = (self->ripha->ipha_src & 0x0000ff00) >> 8;
	this->s2 = (self->ripha->ipha_src & 0x00ff0000) >> 16;
	this->s3 = (self->ripha->ipha_src & 0xff000000) >> 24;
	this->d0 = (self->ripha->ipha_dst & 0x000000ff);
	this->d1 = (self->ripha->ipha_dst & 0x0000ff00) >> 8;
	this->d2 = (self->ripha->ipha_dst & 0x00ff0000) >> 16;
	this->d3 = (self->ripha->ipha_dst & 0xff000000) >> 24;

        /* convert type for use with lltostr() */
        this->s0 = this->s0 < 0 ? 256 + this->s0 : this->s0;
        this->s1 = this->s1 < 0 ? 256 + this->s1 : this->s1;
        this->s2 = this->s2 < 0 ? 256 + this->s2 : this->s2;
        this->s3 = this->s3 < 0 ? 256 + this->s3 : this->s3;
        this->d0 = this->d0 < 0 ? 256 + this->d0 : this->d0;
        this->d1 = this->d1 < 0 ? 256 + this->d1 : this->d1;
        this->d2 = this->d2 < 0 ? 256 + this->d2 : this->d2;
        this->d3 = this->d3 < 0 ? 256 + this->d3 : this->d3;

        /* stringify addresses */
        self->saddr = strjoin(lltostr(this->s0), ".");
        self->saddr = strjoin(self->saddr, strjoin(lltostr(this->s1), "."));
        self->saddr = strjoin(self->saddr, strjoin(lltostr(this->s2), "."));
        self->saddr = strjoin(self->saddr, lltostr(this->s3));
        self->daddr = strjoin(lltostr(this->d0), ".");
        self->daddr = strjoin(self->daddr, strjoin(lltostr(this->d1), "."));
        self->daddr = strjoin(self->daddr, strjoin(lltostr(this->d2), "."));
        self->daddr = strjoin(self->daddr, lltostr(this->d3));
}

/*
 * UDP Save bridge data by queue (kernel -> userland)
 */
fbt:genunix:putq:entry
/self->udpInDatagrams/
{
	/* save IP data */
	queue[arg0] = 1;
	sport[arg0] = self->sport;
	dport[arg0] = self->dport;
	saddr[arg0] = self->saddr;
	daddr[arg0] = self->daddr;
}

/*
 * UDP Done with rput
 */
fbt::udp_rput:return
/self->qp/
{
	/* clear variables */
	self->qp = 0;
	self->ripha = 0;
	self->sport = 0;
	self->dport = 0;
	self->saddr = 0;
	self->daddr = 0;
	self->udpInDatagrams = 0;
}

/*
 * UDP Track this getq_noenab, if queue was marked
 */
fbt:genunix:getq_noenab:entry
/queue[arg0]/
{
	/* track userland queue read */
	self->gq = args[0];
}

/*
 * UDP Retrieve data from queue marked by putq
 */
fbt:genunix:getq_noenab:return
/self->gq/
{
	/* fetch read size */
	this->size = msgdsize((struct msgb *)arg1);
	this->size = (int)this->size > 0 ? this->size : 0;

	/* print output line */
	printf("%5d %6d %-15s %5d <- %-15s %5d %5d %s\n", uid, pid,
	    daddr[(int)self->gq], dport[(int)self->gq], saddr[(int)self->gq], 
	    sport[(int)self->gq], this->size, execname);
}

/*
 * UDP Clear bridge data, if queue empty
 */
fbt:genunix:getq_noenab:return
/self->gq && (self->gq->q_count == 0)/
{
	/* clear IP data variables */
	queue[(int)self->gq] = 0;
	saddr[(int)self->gq] = 0;
	daddr[(int)self->gq] = 0;
	sport[(int)self->gq] = 0;
	dport[(int)self->gq] = 0;
}

/*
 * UDP Done with getq_noenab
 */
fbt:genunix:getq_noenab:return
/self->gq/
{	
	self->gq = 0;
}