#include "socket.h"

#include "byte.h"

#include "dns.h"

#include "buffer.h"

#include "scan.h"

#include "ip6.h"

#include "str.h"

#include "fmt.h"

#include "ip4.h"

#include "textcode.h"

#include <sys/types.h>

#include <unistd.h>

#include <signal.h>

#include <sys/time.h>

#include <sys/resource.h>

#include <stdlib.h>

#include <errno.h>

#include "havealloca.h"

#include "ndelay.h"

static void carp(const char* routine) {

buffer_flush(buffer_1);

buffer_puts(buffer_2,"httpbench: ");

buffer_puts(buffer_2,routine);

buffer_puts(buffer_2,": ");

buffer_puterror(buffer_2);

buffer_putnlflush(buffer_2);

}

static void panic(const char* routine) {

carp(routine);

exit(111);

}

uint16 bindport=0;

static int make_connection(char* ip,uint16 port,uint32 scope_id) {

int v6=byte_diff(ip,12,V4mappedprefix);

int s;

if (v6) {

s=socket_tcp6();

if (s==-1)

panic("socket_tcp6()");

ndelay_off(s);

if (bindport) {

for (;;) {

int r=socket_bind6_reuse(s,V6any,bindport,0);

if (++bindport<1024) bindport=1024;

if (r==0) break;

if (errno!=EADDRINUSE)

panic("socket_bind6");

}

}

if (socket_connect6(s,ip,port,scope_id)==-1) {

carp("socket_connect6");

close(s);

return -1;

}

} else {

s=socket_tcp4();

if (s==-1)

panic("socket_tcp4()");

ndelay_off(s);

if (bindport) {

for (;;) {

int r=socket_bind4_reuse(s,V6any,bindport);

if (++bindport<1024) bindport=1024;

if (r==0) break;

if (errno!=EADDRINUSE)

panic("socket_bind4");

}

}

if (socket_connect4(s,ip+12,port)==-1) {

carp("socket_connect4");

close(s);

return -1;

}

}

return s;

}

static int readanswer(int s,int measurethroughput) {

char buf[8192];

int i,j,body=-1,r;

unsigned long rest;

unsigned long done=0;

struct timeval a,b;

i=0;

while ((r=read(s,buf+i,sizeof(buf)-i)) > 0) {

i+=r;

for (j=0; j+3<i; ++j) {

if (buf[j]=='\r' && buf[j+1]=='\n' && buf[j+2]=='\r' && buf[j+3]=='\n') {

body=j+4;

break;

}

}

if (body!=-1) {

if (byte_diff(buf,7,"HTTP/1.")) {

buffer_putsflush(buffer_2,"invalid HTTP response!\n");

return -1;

}

break;

}

}

if (r==-1) return -1;

rest=-1;

for (j=0; j<r; j+=str_chr(buf+j,'\n')) {

if (byte_equal(buf+j,17,"\nContent-Length: ")) {

char* c=buf+j+17;

if (c[scan_ulong(c,&rest)]!='\r') {

buffer_putsflush(buffer_2,"invalid Content-Length header!\n");

return -1;

}

break;

}

++j;

}

if (measurethroughput) {

gettimeofday(&a,0);

done=r-body;

}

rest-=(r-body);

while (rest) {

r=read(s,buf,rest>sizeof(buf)?sizeof(buf):rest);

if (r<1) {

if (r==-1)

carp("read from HTTP socket");

else {

buffer_puts(buffer_2,"early HTTP EOF; expected ");

buffer_putulong(buffer_2,rest);

buffer_putsflush(buffer_2,"more bytes!\n");

return -1;

}

} else {

rest-=r;

if (measurethroughput) {

unsigned long x=done/1000000;

unsigned long y;

done+=r;

y=done/1000000;

if (x!=y) {

unsigned long d;

unsigned long long z;

gettimeofday(&b,0);

d=(b.tv_sec-a.tv_sec)*1000000;

d=d+b.tv_usec-a.tv_usec;

buffer_puts(buffer_1,"tput ");

z=1000000000ull/d;

buffer_putulong(buffer_1,z);

buffer_putnlflush(buffer_1);

byte_copy(&a,sizeof(a),&b);

}

}

}

}

return 0;

}

int main(int argc,char* argv[]) {

unsigned long count=1000;

unsigned long interval=10;

unsigned long sample=5;

int keepalive=0;

char ip[16];

uint16 port=80;

uint32 scope_id=0;

stralloc ips={0};

int s;

char* request;

int rlen;

signal(SIGPIPE,SIG_IGN);

if (!geteuid()) {

struct rlimit rl;

long l;

#ifdef RLIMIT_NPROC

rl.rlim_cur=RLIM_INFINITY; rl.rlim_max=RLIM_INFINITY;

setrlimit(RLIMIT_NPROC,&rl);

#endif

for (l=0; l<20000; l+=500) {

rl.rlim_cur=l; rl.rlim_max=l;

if (setrlimit(RLIMIT_NOFILE,&rl)==-1) break;

}

}

for (;;) {

int i;

int c=getopt(argc,argv,"c:i:s:kb");

if (c==-1) break;

switch (c) {

case 'k':

keepalive=1;

break;

case 'i':

i=scan_ulong(optarg,&interval);

if (i==0 || optarg[i]) {

buffer_puts(buffer_2,"httpbench: warning: could not parse interval: ");

buffer_puts(buffer_2,optarg+i+1);

buffer_putsflush(buffer_2,"\n");

}

break;

case 'c':

i=scan_ulong(optarg,&count);

if (i==0 || optarg[i]) {

buffer_puts(buffer_2,"httpbench: warning: could not parse count: ");

buffer_puts(buffer_2,optarg+i+1);

buffer_putsflush(buffer_2,"\n");

}

break;

case 's':

i=scan_ulong(optarg,&sample);

if (i==0 || optarg[i]) {

buffer_puts(buffer_2,"httpbench: warning: could not parse sample size: ");

buffer_puts(buffer_2,optarg+i+1);

buffer_putsflush(buffer_2,"\n");

}

break;

case 'b':

bindport=10000;

break;

case '?':

usage:

buffer_putsflush(buffer_2,

"usage: httpbench [-hb] [-c count] [-i interval] [-s sample] url\n"

"\n"

"\t-h\tprint this help\n"

"\t-c n\topen n connections total (default: 1000)\n"

"\t-i n\tevery n connections, measure the latency (default: 10)\n"

"\t-s n\tlatency == average of time to fetch an URL n times (default: 5)\n"

"\t-k\tenable HTTP keep-alive\n"

"\t-b\tbind the sockets ourselves, so the OS doesn't choose the ports\n"

"Setting the number of connections to 1 measures the throughput\n"

"instead of the latency (give URL to a large file).\n");

return 0;

}

}

if (!argv[optind]) goto usage;

if (byte_diff(argv[optind],7,"http://")) goto usage;

{

char* host=argv[optind]+7;

int colon=str_chr(host,':');

int slash=str_chr(host,'/');

char* c;

if (host[0]=='[') { /* ipv6 IP notation */

int tmp;

++host;

--colon; --slash;

tmp=str_chr(host,']');

if (host[tmp]==']') host[tmp]=0;

if (host[tmp+1]==':') colon=tmp+1;

if (colon<tmp+1) colon=tmp+1+str_len(host+tmp+1);

}

if (colon<slash) {

host[colon]=0;

c=host+colon+1;

if (c[scan_ushort(c,&port)]!='/') goto usage;

*c=0;

}

host[colon]=0;

c=host+slash;

*c=0;

{

char* tmp=alloca(str_len(host)+1);

tmp[fmt_str(tmp,host)]=0;

host=tmp;

}

*c='/';

{

int tmp=str_chr(host,'%');

if (host[tmp]) {

host[tmp]=0;

scope_id=socket_getifidx(host+tmp+1);

if (scope_id==0) {

buffer_puts(buffer_2,"httpbench: warning: network interface ");

buffer_puts(buffer_2,host+tmp+1);

buffer_putsflush(buffer_2," not found.\n");

}

}

}

{

stralloc a={0};

stralloc_copys(&a,host);

if (dns_ip6(&ips,&a)==-1) {

buffer_puts(buffer_2,"httpbench: could not resolve IP: ");

buffer_puts(buffer_2,host);

buffer_putnlflush(buffer_2);

return 1;

}

}

request=malloc(300+str_len(host)+str_len(c)*3);

if (!request) panic("malloc");

{

int i;

i=fmt_str(request,"GET ");

i+=fmt_urlencoded(request+i,c,str_len(c));

i+=fmt_str(request+i," HTTP/1.0\r\nHost: ");

i+=fmt_str(request+i,host);

i+=fmt_str(request+i,":");

i+=fmt_ulong(request+i,port);

i+=fmt_str(request+i,"\r\nUser-Agent: httpbench/1.0\r\nConnection: ");

i+=fmt_str(request+i,keepalive?"keep-alive":"close");

i+=fmt_str(request+i,"\r\n\r\n");

rlen=i; request[rlen]=0;

}

}

{

int i;

s=-1;

for (i=0; i+16<=ips.len; i+=16) {

char buf[IP6_FMT];

int v6=byte_diff(ips.s+i,12,V4mappedprefix);

buffer_puts(buffer_1,"connecting to ");

buffer_put(buffer_1,buf,

v6?

fmt_ip6(buf,ips.s+i):

fmt_ip4(buf,ips.s+i+12));

buffer_puts(buffer_1," port ");

buffer_putulong(buffer_1,port);

buffer_putnlflush(buffer_1);

s=make_connection(ips.s+i,port,scope_id);

if (s!=-1) {

byte_copy(ip,16,ips.s+i);

break;

}

}

if (s==-1)

return 1;

}

if (write(s,request,rlen)!=rlen) panic("write");

if (readanswer(s,count==1)==-1) exit(1);

close(s);

if (count==1)

return 0;

{

long i;

long j;

long err=0;

int *socks;

socks=malloc(sizeof(int)*count);

if (!socks) panic("malloc");

for (i=j=0; i<count; ++i) {

struct timeval a,b;

long d;

if (j==0) {

int k,s=0;

long x=0,y=0;

for (k=0; k<sample; ++k) {

if (!keepalive || !k) {

gettimeofday(&a,0);

s=make_connection(ip,port,scope_id);

if (s==-1)

panic("make_connection");

gettimeofday(&b,0);

d=(b.tv_sec-a.tv_sec)*1000000;

d=d+b.tv_usec-a.tv_usec;

x+=d;

}

gettimeofday(&a,0);

write(s,request,rlen);

if (readanswer(s,0)==-1) {

++err;

keepalive=0;

}

gettimeofday(&b,0);

d=(b.tv_sec-a.tv_sec)*1000000;

d=d+b.tv_usec-a.tv_usec;

y+=d;

if (!keepalive) close(s);

}

if (keepalive) close(s);

buffer_puts(buffer_1,"sample ");

buffer_putulong(buffer_1,x);

buffer_puts(buffer_1," ");

buffer_putulong(buffer_1,y/sample);

buffer_putnlflush(buffer_1);

}

++j; if (j==interval) j=0;

gettimeofday(&a,0);

socks[i]=make_connection(ip,port,scope_id);

if (socks[i]==-1)

panic("make_connection");

gettimeofday(&b,0);

d=(b.tv_sec-a.tv_sec)*1000000;

d=d+b.tv_usec-a.tv_usec;

buffer_puts(buffer_1,"clat ");

buffer_putulong(buffer_1,d);

buffer_putnlflush(buffer_1);

}

}

buffer_flush(buffer_1);

return 0;

}