shithub: rtmp

ref: 74213ea11514c175161cc1a61d48a429845d4624
dir: /rtmp.c/

View raw version
#include <u.h>
#include <libc.h>
#include <thread.h>
#include <bio.h>
#include <libsec.h>

enum {
	Port = 1935,

	Anum,
	Aarr = 8,
	Aend,
	Alstr = 12,

	Faudio = 8,
	Fvideo = 9,
	Fscript = 18,

	EncH264 = 7,
	EncAAC = 10,

	FlKey = 1<<0,
	FlHdr = 1<<1,

	Sigsz = 1536,
};

typedef struct Frame Frame;
typedef struct IVF IVF;

struct Frame {
	u64int ts;
	u8int *buf;
	int bufsz;
	int sz;
};

struct IVF {
	u16int w, h;
	u32int tbdenum, tbnum;
};

int mainstacksize = 65536;

static u8int *
amfi16(u8int *p, u8int *e, s16int i)
{
	if(p == nil)
		return nil;
	if(e-p < 2){
		werrstr("buffer short");
		return nil;
	}
	*p++ = i >> 8;
	*p++ = i;

	return p;
}

static u8int *
amfi24(u8int *p, u8int *e, s32int i)
{
	if(p == nil)
		return nil;
	if(e-p < 3){
		werrstr("buffer short");
		return nil;
	}
	*p++ = i >> 16;
	*p++ = i >> 8;
	*p++ = i;

	return p;
}

static u8int *
amfi32(u8int *p, u8int *e, s32int i)
{
	if(p == nil)
		return nil;
	if(e-p < 4){
		werrstr("buffer short");
		return nil;
	}
	*p++ = i >> 24;
	*p++ = i >> 16;
	*p++ = i >> 8;
	*p++ = i;

	return p;
}

static u8int *
amfnum(u8int *p, u8int *e, double v)
{
	union {
		double v;
		u64int u;
	}x;

	if(p == nil)
		return nil;
	if(p+8 > e){
		werrstr("buffer short");
		return nil;
	}
	x.v = v;
	*p++ = x.u >> 56;
	*p++ = x.u >> 48;
	*p++ = x.u >> 40;
	*p++ = x.u >> 32;
	*p++ = x.u >> 24;
	*p++ = x.u >> 16;
	*p++ = x.u >> 8;
	*p++ = x.u;

	return p;
}

static u8int *
amfkvnum(u8int *p, u8int *e, char *name, double v)
{
	int n;

	if(p == nil)
		return nil;
	if((n = strlen(name)) > 0xffff){
		werrstr("string too long");
		return nil;
	}
	if(p+2+n+8 > e){
		werrstr("buffer short");
		return nil;
	}
	p = amfi16(p, e, n);
	p = (u8int*)memmove(p, name, n) + n;

	return amfnum(p, e, v);
}

static u8int *
amfstr(u8int *p, u8int *e, char *s)
{
	int n;

	if(p == nil)
		return nil;
	n = strlen(s);
	if(p+1+4+n > e){
		werrstr("string too long");
		return nil;
	}
	*p++ = Alstr;

	return (u8int*)memmove(amfi32(p, e, n), s, n) + n;
}

static u8int *
amfarr(u8int *p, u8int *e)
{
	if(p == nil)
		return nil;
	if(p == e){
		werrstr("buffer short");
		return nil;
	}
	*p++ = Aarr;

	return p;
}

static u8int *
amfend(u8int *p, u8int *e)
{
	return amfi24(p, e, Aend);
}

static u8int *
flvscript(u8int *p, u8int *e, int w, int h, int audio)
{
	u8int *psz, *d, *p0;
	int stream;
	u32int ts;

	if(p+16 > e){
		werrstr("buffer short");
		return nil;
	}

	/* FIXME ever need to change these? */
	stream = 0;
	ts = 0;

	p0 = p;
	*p++ = Fvideo;
	psz = p;
	p = amfi24(p, e, 0); /* sz set later */
	p = amfi24(p, e, ts);
	*p++ = ts>>24;
	p = amfi24(p, e, stream);

	d = p;
	p = amfstr(p, e, "onMetaData");
	p = amfarr(p, e);
	p = amfi32(p, e, audio ? 5 : 4);
	p = amfkvnum(p, e, "duration", 0.0);
	p = amfkvnum(p, e, "width", w);
	p = amfkvnum(p, e, "height", h);
	p = amfkvnum(p, e, "videocodecid", EncH264);
	if(audio)
		p = amfkvnum(p, e, "audiocodecid", EncAAC);
	p = amfend(p, e);
	amfi24(psz, e, p-d);

	return amfi32(p, e, p-p0);
}

static u8int *
flvdata(u8int *p, u8int *e, u32int pts, u32int dts, void *data, int sz, int type, int fl)
{
	u8int *p0, *psz, *d;
	int stream;

	/* FIXME ever need to change these? */
	stream = 0;

	assert(type == Faudio || type == Fvideo);
	p0 = p;
	*p++ = type;
	psz = p;
	p = amfi24(p, e, 0); /* size to be set later */
	p = amfi24(p, e, dts);
	*p++ = dts >> 24;
	p = amfi24(p, e, stream);

	d = p;
	if(type == Faudio){
		*p++ = (EncAAC<<4) | 0x0f;
		*p++ = (fl & FlHdr) ? 0 : 1;
	}
	if(type == Fvideo){
		*p++ = ((fl & FlKey) ? 0x10 : 0x20) | EncH264;
		*p++ = (fl & FlHdr) ? 0 : 1;
		pts = ((fl & FlHdr) || pts < dts) ? 0 : (pts - dts);
		p = amfi24(p, e, pts);
		if((fl & FlHdr) == 0)
			p = amfi32(p, e, sz);
	}
	p = (u8int*)memmove(p, data, sz) + sz;
	amfi24(psz, e, p-d);

	return amfi32(p, e, p-p0);
}

static int
Bu16le(Biobuf *b, u16int *o)
{
	int x;

	x = Bgetc(b);
	x |= Bgetc(b)<<8;
	*o = x;
	if(x < 0)
		werrstr("failed to read 2 bytes");

	return x < 0 ? -1 : 0;
}

static int
Bu32le(Biobuf *b, u32int *o)
{
	int x, i;

	*o = 0;
	for(i = 0; i < 4; *o |= x<<(i*8), i++){
		if((x = Bgetc(b)) < 0){
			werrstr("failed to read 4 bytes");
			return -1;
		}
	}

	return 0;
}

static int
Bu64le(Biobuf *b, u64int *o)
{
	int x, i;

	*o = 0;
	for(i = 0; i < 8; *o |= x<<(i*8), i++){
		if((x = Bgetc(b)) < 0){
			werrstr("failed to read 8 bytes");
			return -1;
		}
	}

	return 0;
}

static int
ivfopen(Biobuf *v, IVF *o)
{
	u16int hlen;
	u8int b[6];

	if(Bread(v, b, 6) != 6 || Bu16le(v, &hlen) < 0 ||
	   hlen < 0x20 || memcmp(b, "DKIF", 4) != 0 ||
	   Bread(v, b, 4) != 4){
		werrstr("invalid header");
		goto err;
	}
	if(memcmp(b, "AVC1", 4) != 0){
		werrstr("not H.264");
		goto err;
	}
	if(Bu16le(v, &o->w) < 0 ||
	   Bu16le(v, &o->h) < 0 ||
	   Bu32le(v, &o->tbdenum) < 0 ||
	   Bu32le(v, &o->tbnum) < 0){
		werrstr("invalid data");
		goto err;
	}
	if(Bseek(v, hlen, 0) != hlen){
		werrstr("broken stream");
		goto err;
	}

	return 0;
err:
	werrstr("ivf: %r");
	return -1;
}

static int
ivfread(Biobuf *v, Frame *f)
{
	u8int *buf;
	u64int ts;
	u32int sz;
	int n;

	if(Bu32le(v, &sz) < 0 || Bu64le(v, &ts) < 0 || (int)sz < 0){
		/* eof */
		f->sz = 0;
		return 0;
	}
	buf = f->buf;
	if(sz > f->bufsz){
		if((buf = realloc(f->buf, sz)) == nil){
			werrstr("frame is too big: %d bytes", sz);
			goto err;
		}
		f->buf = buf;
	}
	if((n = Bread(v, buf, sz)) != sz){
		werrstr("short read (%d < %d)", n, sz);
		goto err;
	}
	f->buf = buf;
	f->sz = sz;
	f->ts = ts;

	return 0;
err:
	werrstr("ivf: %r");
	return -1;
}

static int
handshake(int f, char *path)
{
	u8int cl[1+Sigsz], sv[1+Sigsz];

	cl[0] = 3; /* no encryption */
	memset(cl+1, 0, 8);
	prng(cl+1+8, Sigsz-8);
	if(write(f, cl, sizeof(cl)) != sizeof(cl))
		goto err;
	if(readn(f, sv, sizeof(sv)) != sizeof(sv))
		goto err;
	if(cl[0] != sv[0]){
		werrstr("expected %d (no encryption), got %d", cl[0], sv[0]);
		goto err;
	}
	if(write(f, sv+1, Sigsz) != Sigsz)
		goto err;
	if(readn(f, sv+1, Sigsz) != Sigsz)
		goto err;
	if(memcmp(cl, sv, sizeof(cl)) != 0){
		werrstr("signature mismatch");
		goto err;
	}

	return f;

err:
	werrstr("handshake: %r");
	close(f);
	return -1;
}

static int
rtmpdial(char *url)
{
	char *s, *e, *path;
	int f, port, ctl;

	if(memcmp(url, "rtmp://", 7) != 0){
		werrstr("invalid url");
		goto err;
	}
	s = url + 7;
	if((e = strpbrk(s, ":/")) == nil){
		werrstr("no path");
		goto err;
	}
	port = 1935;
	if(*e == ':'){
		if((port = strtol(e+1, &path, 10)) < 1 || path == e+1 || *path != '/'){
			werrstr("invalid port");
			goto err;
		}
	}else{
		path = e;
	}

	s = smprint("tcp!%.*s!%d", (int)(e-s), s, port);
	f = dial(s, nil, nil, &ctl);
	free(s);
	if(f < 0)
		goto err;

	return handshake(f, path);
err:
	werrstr("rtmpdial: %r");
	return -1;
}

static void
usage(void)
{
	fprint(2, "usage: %s [-a AUDIO] -v VIDEO [URL]\n", argv0);
	threadexitsall("usage");
}

void
threadmain(int argc, char **argv)
{
	Biobuf *a, *v, o;
	u8int *b, *p, *e;
	int bufsz, fd;
	Frame f;
	IVF ivf;

	a = nil;
	v = nil;
	ARGBEGIN{
	case 'a':
		if((a = Bopen(EARGF(usage()), OREAD)) == nil)
			sysfatal("%r");
		break;
	case 'v':
		if((v = Bopen(EARGF(usage()), OREAD)) == nil)
			sysfatal("%r");
		break;
	default:
		usage();
	}ARGEND

	if(argc != 1)
		usage();
	if(v == nil)
		sysfatal("no video specified");
	if(ivfopen(v, &ivf) != 0)
		sysfatal("%r");
	srand(time(nil));
	if((fd = rtmpdial(argv[0])) < 0 || Binit(&o, fd, OWRITE) < 0)
		sysfatal("%r");

	bufsz = 65536;
	if((b = malloc(bufsz)) == nil)
		sysfatal("memory");
	e = b + bufsz;

	if((p = flvscript(b, e, ivf.w, ivf.h, 0)) == nil || Bwrite(&o, b, p-b) < 0)
		sysfatal("%r");

	memset(&f, 0, sizeof(f));
	for(;;){
		if(ivfread(v, &f) != 0)
			sysfatal("%r");
		if(f.sz == 0)
			break;
		if(bufsz < f.sz+64){
			free(b);
			bufsz *= 2;
			if((b = malloc(bufsz)) == nil)
				sysfatal("memory");
			e = b + bufsz;
		}
		if((p = flvdata(b, e, f.ts, f.ts, f.buf, f.sz, Fvideo, FlHdr)) == nil)
			sysfatal("video: flvdata: %r");
		if(Bwrite(&o, b, p-b) < 0)
			sysfatal("%r");
		Bflush(&o);
	}

	threadexitsall(nil);
}